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

• URL: http://arxiv.org/abs/2509.23186v1
• Date: Sat, 27 Sep 2025 08:40:15 GMT
• Title: Understanding and Enhancing the Planning Capability of Language Models via Multi-Token Prediction
• Authors: Qimin Zhong, Hao Liao, Siwei Wang, Mingyang Zhou, Xiaoqun Wu, Rui Mao, Wei Chen,
• Abstract summary: We theoretically analyze the Multi-Token Prediction (MTP) paradigm using a Transformer architecture.<n>We propose two strategies to enhance the transfer layer and overall learning quality: Next-Token Injection (NTI) and a Transformer-based transfer layer.<n>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.
• Score: 17.631683772383017
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• TopoCurate:Modeling Interaction Topology for Tool-Use Agent Training [53.93696896939915]
  Training tool-use agents typically rely on Supervised Fine-Tuning (SFT) on successful trajectories and Reinforcement Learning (RL) on pass-rate-selected tasks.<n>We propose TopoCurate, an interaction-aware framework that projects multi-trial rollouts from the same task into a unified semantic quotient topology.<n>TopoCurate achieves consistent gains of 4.2% (SFT) and 6.9% (RL) over state-of-the-art baselines.
  arXiv  Detail & Related papers  (2026-03-02T10:38:54Z)
• Foundation Model for Skeleton-Based Human Action Understanding [56.89025287217221]
  This paper presents a Unified Skeleton-based Dense Representation Learning framework.<n>USDRL consists of a Transformer-based Dense Spatio-Temporal (DSTE), Multi-Grained Feature Decorrelation (MG-FD), and Multi-Perspective Consistency Training (MPCT)
  arXiv  Detail & Related papers  (2025-08-18T02:42:16Z)
• Latent Diffusion Planning for Imitation Learning [78.56207566743154]
  Latent Diffusion Planning (LDP) is a modular approach consisting of a planner and inverse dynamics model.<n>By separating planning from action prediction, LDP can benefit from the denser supervision signals of suboptimal and action-free data.<n>On simulated visual robotic manipulation tasks, LDP outperforms state-of-the-art imitation learning approaches.
  arXiv  Detail & Related papers  (2025-04-23T17:53:34Z)
• A representational framework for learning and encoding structurally enriched trajectories in complex agent environments [1.904851064759821]
  The ability of artificial intelligence agents to make optimal decisions and generalise them to different domains and tasks is compromised in complex scenarios.<n>One way to address this issue has focused on learning efficient representations of the world and on how the actions of agents affect them, such as disentangled representations that exploit symmetries.<n>We propose to enrich the agent's ontology and extend the traditionalisation of trajectories to provide a more nuanced view of task execution.
  arXiv  Detail & Related papers  (2025-03-17T14:04:27Z)
• On Adversarial Robustness of Language Models in Transfer Learning [13.363850350446869]
  We show that transfer learning, while improving standard performance metrics, often leads to increased vulnerability to adversarial attacks.<n>Our findings demonstrate that larger models exhibit greater resilience to this phenomenon, suggesting a complex interplay between model size, architecture, and adaptation methods.
  arXiv  Detail & Related papers  (2024-12-29T15:55:35Z)
• Interpreting token compositionality in LLMs: A robustness analysis [10.777646083061395]
  Constituent-Aware Pooling (CAP) is a methodology designed to analyse how large language models process linguistic structures.<n>CAP intervenes in model activations through constituent-based pooling at various model levels.<n>Our findings highlight fundamental limitations in current transformer architectures regarding compositional semantics processing and model interpretability.
  arXiv  Detail & Related papers  (2024-10-16T18:10:50Z)
• ALPINE: Unveiling the Planning Capability of Autoregressive Learning in Language Models [48.559185522099625]
  Planning is a crucial element of both human intelligence and contemporary large language models (LLMs) This paper investigates the emergence of planning capabilities in Transformer-based LLMs via their next-word prediction mechanisms.
  arXiv  Detail & Related papers  (2024-05-15T09:59:37Z)
• End-to-End Meta-Bayesian Optimisation with Transformer Neural Processes [52.818579746354665]
  This paper proposes the first end-to-end differentiable meta-BO framework that generalises neural processes to learn acquisition functions via transformer architectures. We enable this end-to-end framework with reinforcement learning (RL) to tackle the lack of labelled acquisition data.
  arXiv  Detail & Related papers  (2023-05-25T10:58:46Z)
• The Multiscale Surface Vision Transformer [10.833580445244094]
  We introduce the Multiscale Surface Vision Transformer (MS-SiT) as a backbone architecture for surface deep learning. Results demonstrate that the MS-SiT outperforms existing surface deep learning methods for neonatal phenotyping prediction tasks.
  arXiv  Detail & Related papers  (2023-03-21T15:00:17Z)
• SIM-Trans: Structure Information Modeling Transformer for Fine-grained Visual Categorization [59.732036564862796]
  We propose the Structure Information Modeling Transformer (SIM-Trans) to incorporate object structure information into transformer for enhancing discriminative representation learning. The proposed two modules are light-weighted and can be plugged into any transformer network and trained end-to-end easily. Experiments and analyses demonstrate that the proposed SIM-Trans achieves state-of-the-art performance on fine-grained visual categorization benchmarks.
  arXiv  Detail & Related papers  (2022-08-31T03:00:07Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.