Understanding and Enhancing the Planning Capability of Language Models via Multi-Token Prediction

📄 arXiv: 2509.23186v1 📥 PDF

作者: Qimin Zhong, Hao Liao, Siwei Wang, Mingyang Zhou, Xiaoqun Wu, Rui Mao, Wei Chen

分类: cs.AI, cs.LG

发布日期: 2025-09-27

💡 一句话要点

通过多Token预测增强语言模型在复杂规划中的传递关系学习能力

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 语言模型 复杂规划 传递关系学习 多Token预测 Transformer 路径规划 Next-Token Injection

📋 核心要点

  1. 现有大型语言模型在复杂规划中难以学习传递关系，限制了其应用。
  2. 论文通过分析多Token预测范式，提出增强传递层和学习质量的策略。
  3. 实验表明，所提方法显著提升了模型在合成图和Blocksworld上的路径规划能力。

📝 摘要（中文）

大型语言模型(LLMs)在各种任务中取得了令人瞩目的性能，但仍然难以学习传递关系，这是复杂规划的基石。为了解决这个问题，我们研究了多Token预测(MTP)范式及其对传递关系学习的影响。我们使用由共享输出头和传递层组成的Transformer架构对MTP范式进行了理论分析。我们的分析表明，传递层逐渐学习多步邻接信息，从而使骨干模型能够捕获未观察到的传递可达关系，即使这些关系并非直接存在于训练数据中，尽管在邻接估计中存在一些不可避免的噪声。在此基础上，我们提出了两种策略来增强传递层和整体学习质量：Next-Token Injection (NTI)和基于Transformer的传递层。我们在合成图和Blocksworld规划基准上的实验验证了我们的理论发现，并表明这些改进显著增强了模型的路径规划能力。这些发现加深了我们对具有MTP的Transformer如何在复杂规划任务中学习的理解，并提供了克服传递瓶颈的实用策略，为结构感知和通用规划模型铺平了道路。

🔬 方法详解

问题定义：论文旨在解决大型语言模型在复杂规划任务中学习传递关系的困难。现有方法难以有效地学习和泛化到未见过的传递关系，导致规划能力受限。痛点在于模型无法很好地利用训练数据中隐含的传递信息，从而难以进行多步推理。

核心思路：论文的核心思路是利用多Token预测（MTP）范式，通过一个专门设计的传递层来显式地学习和传递邻接信息。MTP允许模型一次预测多个token，从而能够更好地捕捉序列中的长程依赖关系。通过增强传递层，模型可以更好地学习到训练数据中未直接观察到的传递可达关系。

技术框架：整体框架基于Transformer架构，包含一个共享输出头和一个传递层。传递层负责学习和传递邻接信息，而共享输出头则负责最终的token预测。论文提出了两种增强传递层的方法：Next-Token Injection (NTI) 和 Transformer-based transfer layer。NTI通过注入下一个token的信息来增强传递层的学习能力。Transformer-based transfer layer则使用Transformer结构来建模传递层，从而更好地捕捉复杂的邻接关系。

关键创新：论文的关键创新在于对MTP范式的理论分析，揭示了传递层在学习传递关系中的作用。此外，提出的NTI和Transformer-based transfer layer是增强传递层学习能力的新方法。与现有方法相比，该方法能够更有效地学习和泛化到未见过的传递关系，从而提升了模型的规划能力。

关键设计：NTI的关键设计在于如何有效地将下一个token的信息注入到传递层中。Transformer-based transfer layer的关键设计在于如何选择合适的Transformer结构和参数，以更好地建模邻接关系。损失函数通常采用交叉熵损失，用于衡量预测token的准确性。具体的参数设置和网络结构会根据不同的实验设置进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，所提出的NTI和Transformer-based transfer layer能够显著提升模型在合成图和Blocksworld上的路径规划能力。在Blocksworld数据集上，相比于基线模型，该方法能够将规划成功率提升10%以上，验证了其有效性。此外，实验还验证了理论分析的正确性，即传递层在学习传递关系中起着关键作用。

🎯 应用场景

该研究成果可应用于各种需要复杂规划能力的领域，如机器人导航、任务调度、游戏AI和自动驾驶。通过提升语言模型的规划能力，可以使其更好地理解和执行复杂的指令，从而实现更智能化的决策和控制。未来，该研究有望推动通用人工智能的发展，使其能够更好地解决现实世界中的复杂问题。

📄 摘要（原文）

Large Language Models (LLMs) have achieved impressive performance across diverse tasks but continue to struggle with learning transitive relations, a cornerstone for complex planning. To address this issue, we investigate the Multi-Token Prediction (MTP) paradigm and its impact to transitive relation learning. We theoretically analyze the MTP paradigm using a Transformer architecture composed of a shared output head and a transfer layer. Our analysis reveals that the transfer layer gradually learns the multi-step adjacency information, which in turn enables the backbone model to capture unobserved transitive reachability relations beyond those directly present in the training data, albeit with some inevitable noise in adjacency estimation. Building on this foundation, we propose two strategies to enhance the transfer layer and overall learning quality: Next-Token Injection (NTI) and a Transformer-based transfer layer. Our experiments on both synthetic graphs and the Blocksworld planning benchmark validate our theoretical findings and demonstrate that the improvements significantly enhance the model's path-planning capability. These findings deepen our understanding of how Transformers with MTP learn in complex planning tasks, and provide practical strategies to overcome the transitivity bottleneck, paving the way toward structurally aware and general-purpose planning models.