多词预测让Transformer学会反向规划,提升推理能力。
How Transformers Learn to Plan via Multi-Token Prediction
- 用多词预测替代单词预测,引导模型反向推理路径
- 在图路径查找、计数谜题等任务上超越传统方法
- 适合研究大模型推理机制与可解释性的人参考
尽管下一个词预测(NTP)是语言模型训练的标准目标,但在捕捉推理任务中的全局结构方面常表现不佳。多词预测(MTP)作为新兴替代方案,其内在机制仍不明确。本文研究MTP如何促进推理,重点关注规划能力。实证结果显示,MTP在合成图路径查找任务及更现实的推理基准(如Countdown和布尔可满足性问题)中均持续优于NTP。理论分析针对简化两层Transformer在星型图任务上的表现,证明MTP诱导出两阶段逆向推理过程:模型先关注终点节点,再通过回溯中间节点重建路径。这一行为源于MTP的梯度解耦特性,提供比NTP更清晰的训练信号。最终结果表明,多词目标天然倾向于引导优化朝向稳健且可解释的推理电路。
原文摘要 · Abstract (English)
While next-token prediction (NTP) has been the standard objective for training language models, it often struggles to capture global structure in reasoning tasks. Multi-token prediction (MTP) has recently emerged as a promising alternative, yet its underlying mechanisms remain poorly understood. In this paper, we study how MTP facilitates reasoning, with a focus on planning. Empirically, we show that MTP consistently outperforms NTP on both synthetic graph path-finding tasks and more realistic reasoning benchmarks, such as Countdown and boolean satisfiability problems. Theoretically, we analyze a simplified two-layer Transformer on a star graph task. We prove that MTP induces a two-stage reverse reasoning process: the model first attends to the end node and then reconstructs the path by tracing intermediate nodes backward. This behavior arises from a gradient decoupling property of MTP, which provides a cleaner training signal compared to NTP. Ultimately, our results highlight how multi-token objectives inherently bias optimization toward robust and interpretable reasoning circuits.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。