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.
Figures & tables
Figure 1: Standard autoregressive inference vs latent lookahead. Left: in standard next token prediction, the model samples from the hidden state of the latest generated token after applying the final unembedding head, and appends the generated token to the context. Right: in our approach, the model enters the latent lookahead thinking, where the hidden states are fed directly into the context instead of sampled visible tokens. This procedure is repeated τ times, and only then the visible token is sampled from the first latent position. In the figure above, the tokens x2 and x4 are selected, τ=3 , and zi,j indicates the j -th latent token relative to the i -th visible token. See also Fig. 3
Figure 2: Lookahead behaviour when solving a Sudoku. In the first slot, both 1 and 3 are viable options. However, when thinking ahead to the second empty slot, where 3 is the only plausible entry, it is easy to realize that 1 is the right choice for the first slot.
Figure 3: Left: the visible tokens (in blue) are supervised with standard next token prediction objective. The latent tokens are supervised explicitly with an equivalent number of steps ahead, i.e. the zi,j latent token is supervised with the xi+j visible token. Right: to ensure that the model is able to refine the next token prediction based on its own assessment of the future steps, we allow for full (non-causal) masking between the latent tokens. The visible tokens follow the standard causal masking. Finally, the latent thoughts causally attend to previous visible, but not to the previous latent. This allows for parallel generation of the latent thoughts during training.
Model
Mini 4 × 4 Sudoku
Full 9 × 9 Sudoku
ProsQA
Maze
Ours
93.5
35.5
91.8
21.5
Pause
86.0
12.5
82.5
19.5
Standard NTP
78.0
12.5
80.5
18.5
Table 1: Accuracy (%) for each dataset. The number of tokens used are the same as the answer length: 19, 70, 5 for mini-Sudoku, Sudoku and ProsQA, respectively.
Figure 4: Effect of increasing the number of latent tokens in three tasks. As the compute allocated increases, the <pause> token method either saturates or fails to improve over the NTP baseline. In contrast, latent lookahead continues to benefit from larger τ .
Figure 6: (a): Ablation comparing latent lookahead with Looped variant and MTP (mini and full Sudoku). (b): Comparison between sequential and random latent position strategies on mini Sudoku .
Figure 7: Eval curves ablating over the following design choices: (1) use a causal mask within the latent thoughts (orange), and (2) use the model as a multi-token predictor (dashed).
Depth
Accuracy (%)
2
78.0
4
79.0
8
79.5
16
82.4
32
80.0
Ours
93.5
Table 2: Accuracy vs. depth on small Sudoku.
Figure 8: Visualization of latent lookahead distributions in the probability simplex over {1,2,3,4} . Each point corresponds to one latent token, projected as a convex combination of the four vertices. (a) Refinement of the first token over successive lookahead steps (color = step index). The distribution progressively sharpens towards the correct vertex, illustrating iterative refinement under the non-causal mask. (b) Representation of all latent tokens at the final step of the thought process (color = latent token position). Tokens cluster near vertices when predictions are confident, while intermediate points reflect uncertainty between candidate symbols.
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
Parameter
Value
Model
Architecture
2-layer GPT-2 style Transformer
Hidden size
768
Attention heads
8
Layers
2
Vocabulary size
varies across datasets
Appendix
Table 3: Main hyperparameters and settings used for the scratch experiments (Sudoku, ProsQA, Maze) with latent lookahead.
Parameter
Value
Base model
allenai/OLMo-2-0425-1B
Dataset
Tulu 3 SFT (0225)
Hidden size
2048
Intermediate size
8192
Attention heads
16
Layers
16
Appendix
Table 4: Main hyperparameters and settings used for the OLMo-2-0425-1B finetuning run with latent tokens.
Figure 9: Visualization of Sudoku predictions. Each point corresponds to a decoded token, with green indicating a correct prediction and red a wrong one. The vertices (1–4) denote the possible output classes.
Table 13
Parameter
Value
Base model
Qwen/Qwen2.5-0.5B
Dataset
Sudoku and GSM8k
Hidden size
896
Intermediate size
4864
Attention heads
14
Layers
24
Appendix
Table 7: Main hyperparameters and settings used for the Qwen2.5-0.5B finetuning run
Dataset
# shots
Shots (in order)
Prompt prefix
GSM8K
1
Shot 1 — Question: There are 15 trees in the grove. Grove workers will plant trees in the grove today. After they are done, there will be 21 trees. How many trees did the grove workers plant today? Answer: There are 15 trees originally. Then there were 21 trees after some more were planted. So there must have been 21 - 15 = 6. So the answer is 6.
Answer the following math question and explain step by step.
Aqua
2
Shot 1 — Question: Two friends plan to walk along a 43-km trail, starting at opposite ends of the trail at the same time. If Friend P’s rate is 15% faster than Friend Q’s, how many kilometers will Friend P have walked when they pass each other? Options: A) 21; B) 21.5; C) 22; D) 22.5; E) 23. Answer: If Q complete x kilometers, then P completes 1.15x kilometers. x + 1.15x = 43; 2.15x = 43; x = 43/2.15 = 20. Then P will have walked 1.15*20 = 23 km. The answer is E. Shot 2 — Question: In the coordinate plane, points (x, 1) and (5, y) are on line k. If line k passes through the origin and has slope 1/5, then what are the values of x and y respectively? Options: A) 4 and 1; B) 1 and 5; C) 5 and 1; D) 3 and 5; E) 5 and 3. Answer: Line k passes through the origin and has slope 1/5, so its equation is y = (1/5)x. Thus (x,1) = (5,1) and (5,y) = (5,1) ⇒x=5 , y=1 . The answer is C.
Answer the following math questions by choosing one of the options A,B,C,D,E. Explain step by step.
BBH
0
—
Answer the following question by indicating the correct option.
Appendix
Table 8: Few-shot configurations and prompt prefixes used per dataset. Shots are taken in order from the provided pools.