[论文解读] Thinking into the Future: Latent Lookahead Training for Transformers

Lorenzo Noci, Gregor Bachmann+2 more|arXiv (Cornell University)|Mar 3, 2026

Multimodal Machine Learning Applications被引用 0

一句话总结

本论文提出潜在前瞻（latent lookahead），一种训练策略：将变换器在发出下一个标记前向展开 tau 步的隐藏状态，并以未来 tau 个真值标记作为监督，以提升计划与推理任务的表现。

ABSTRACT

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.

研究动机与目标

• 通过使多步潜在前瞻来纠正自回归下一个标记预测中的短视行为并提供动机。
• 引入一个可微分的训练目标，对潜在预测与未来的真实标记进行监督。
• 通过扩展上下文中的潜在思考来实现对每个标记的非一致计算分配。
• 在需要前瞻性的计划性任务（如数独、ProsQA 与 Maze）上展示收益。

提出的方法

• 定义扩展在每个可见标记 x_i 之后的上下文的潜在标记 z_{i,j}。
• 使用带增强上下文 e^{aug} 的变换器对 tau 次潜在步进行展开以生成 z_{i,j}。
• 训练 z_{i,j} 以预测 x_{i+j}，以未来 tau 个真值标记为监督。
• 在使用常规下一个标记预测来训练可见标记的同时，通过增强注意力让可见标记能够看到潜在思考。
• 使用非完全因果的注意力掩码，允许潜在思考的并行生成以及思考内的双向潜在注意力。
• 将 L_NTP 与 L_latent 合并成一个目标 L = L_NTP + L_latent。

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 mode

实验结果

研究问题

• RQ1 潜在前瞻相较自回归基线和暂停标记基线，在需要规划的任务上是否提升性能？
• RQ2 增加潜在视野 tau 和潜在位置数量 n 会如何影响性能？
• RQ3 潜在前瞻是否比多标记预测或循环-改进基线更有效？
• RQ4 注意力掩码和潜在解码策略如何影响学习与推理？

主要发现

• 潜在前瞻在 Sudoku、ProsQA、Maze 等规划任务上显著超越自回归和暂停基线。
• 完整的 9x9 数独的准确率从 12.5%（NTP 基线）提升到 35.5%（潜在前瞻）。
• 在 Mini-Sudoku 4x4 上，潜在前瞻达到 93.5% 的准确率，超过了更深的基线。
• 增加潜在视野 tau，在各任务上均呈单调提升的准确率，超越基线且上升幅度比暂停标记更稳定。
• 在序列前期分配潜在思考比随机分配带来更大收益。
• 可视化结果显示潜在标记集中在决策顶点附近，表明预测的迭代改进。

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.