How Transformers Learn to Plan via Multi-Token Prediction

尽管“下一词预测”（NTP）一直是训练语言模型的标准目标，但它在推理任务中往往难以捕捉全局结构。近年来，“多词预测”（MTP）作为一种颇具前景的替代方案崭露头角，然而其内在机制仍缺乏深入理解。本文聚焦于MTP如何促进推理能力，尤其关注规划类任务。实验表明，MTP在合成图路径搜索任务以及更具现实意义的推理基准测试（如“倒计时”游戏和布尔可满足性问题）上，均持续优于NTP。理论上，我们在一种简化的双层Transformer架构上，针对星形图任务展开分析；我们证明，MTP会诱导出一种两阶段的逆向推理过程：模型首先关注终点节点，再通过反向追踪中间节点逐步重构整条路径。这一行为源于MTP所具有的梯度解耦特性——相比NTP，它能提供更清晰、更纯净的训练信号。最终，我们的研究结果揭示：多词预测这类目标函数，天然地将模型优化过程导向鲁棒且可解释的推理回路。

传统语言模型训练目标——逐词预测（NTP）在需要全局结构建模的推理任务（如路径规划、逻辑推演）中表现受限；本文旨在验证：多词预测（MTP）能否更有效地诱导模型学习可解释、鲁棒的结构化推理能力，尤其在规划类任务中是否具有本质优势。该问题新颖且重要——它超越了单纯性能比较，直指训练目标如何塑造模型内部推理机制这一基础性认知问题。

提出MTP通过梯度解耦（gradient decoupling）提供更干净的训练信号，从而在理论上诱导出两阶段逆向推理机制（先关注终点，再反向回溯路径）；该机制在简化Transformer上被严格证明，并非启发式设计，而是MTP目标函数内在优化动态的必然结果，首次从理论层面揭示了多词目标如何‘塑造’可解释推理电路。

实证上，在合成星图路径任务、Countdown数字游戏和布尔可满足性（SAT）三个异构推理基准上一致超越NTP；理论分析基于可验证的两层Transformer+星图设定，给出可证明的逆向注意力机制；论文强调MTP的‘归纳偏置’本质——其目标函数天然鼓励模块化、因果清晰的推理路径；暂未提及其代码是否开源；值得深入的方向包括：MTP在长链推理中的泛化边界、与思维链（CoT）提示的协同机制、以及在真实世界规划任务（如机器人指令生成）中的迁移验证。

《Chain-of-Thought Prompting Elicits Reasoning in Large Language Models》（Wei et al., 2022）；《Large Language Models Are Zero-Shot Reasoners》（Kojima et al., 2023）；《Reasoning with Language Model Is Planning with World Model》（Huang et al., 2023）；《Self-Refine: Iterative Refinement with Self-Feedback》（Madaan et al., 2024）；《Training Verifiers to Solve Math Word Problems》（Cobbe et al., 2021）