Thinking into the Future: Latent Lookahead Training for Transformers
This addresses the problem of limited foresight and compute allocation in language models for tasks requiring planning, though it is incremental as it builds on existing transformer architectures.
The paper tackled the limitations of autoregressive language models by introducing latent lookahead training, which allows models to perform multi-step lookahead in latent space before generating tokens, resulting in substantial performance improvements on planning tasks like maze solving, Sudoku, and ProsQA.
Autoregressive language models trained with next-token prediction generate text by sampling one discrete token at a time. Although very scalable, this objective forces the model to commit at every step, preventing it from exploring or reflecting upon multiple plausible continuations. Furthermore, the compute allocation across tokens is uniform; every token is formed based on a single forward-pass, potentially limiting the model's expressiveness in cases where difficult tokens require inherently more compute. Towards addressing these limitations, we introduce latent lookahead, a training strategy that enables models to "think" before generating: at selected positions in the sequence, before committing to the next token, the model performs a multi-step lookahead in latent space. More precisely, instead of sampling future tokens, we leverage the network's latent space by recursively feeding its hidden states back into the context for $Ï$ steps, investing more compute on predicting that token. This produces $Ï$ latent predictions that are supervised against the next $Ï$ ground-truth tokens, encouraging the model to "lookahead" and refine its prediction. We show that latent lookahead substantially outperforms both autoregressive and non-autoregressive baselines on planning tasks such as maze solving, Sudoku, and ProsQA, where foresight is essential.