让Transformer在生成前‘预演’未来,提升复杂任务决策能力
Thinking into the Future: Latent Lookahead Training for Transformers
- 在特定位置通过潜空间递归回溯,模拟多步前瞻预测
- 在迷宫求解、数独等任务上显著超越自回归与非自回归基线
- 适合需要规划与推理的生成任务,如逻辑谜题与复杂对话
基于下一个词预测训练的自回归语言模型逐个采样离散词元生成文本。尽管可扩展性好,但该目标迫使模型每一步都立即确定输出,无法探索或反思多种可能的后续内容。此外,每个词元的计算分配均匀;每个词元仅依赖一次前向传播,可能限制模型在难点词元上表达能力。为解决上述问题,我们提出潜空间前瞻(latent lookahead)训练策略,使模型能在生成前“思考”:在序列的选定位置,模型在提交下一个词元前,于潜空间执行多步前瞻。具体而言,不直接采样未来词元,而是将网络隐藏状态递归反馈至上下文,进行τ步潜空间回溯,投入更多计算预测该词元。这生成τ个潜空间预测,以真实下一τ个词元为监督信号,激励模型“前瞻”并优化预测。实验表明,潜空间前瞻在迷宫求解、数独及ProsQA等需前瞻性规划的任务中,显著优于自回归与非自回归基线。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。