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

arXiv: https://arxiv.org/abs/2509.23186 · HTML: https://mlsanity.com/paper/2509.23186/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 36 of its 42 references that are in the map.
- First posted on arXiv 27 Sep 2025.

## Builds on (most cited first, 30 of 36)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [End-to-End Object Detection with Transformers](https://mlsanity.com/paper/2005.12872/) (2020)
- [Tree of Thoughts: Deliberate Problem Solving with Large Language Models](https://mlsanity.com/paper/2305.10601/) (2023)
- [A Survey on Large Language Model based Autonomous Agents](https://mlsanity.com/paper/2308.11432/) (2023)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [HuggingGPT: Solving AI Tasks with ChatGPT and its Friends in Hugging Face](https://mlsanity.com/paper/2303.17580/) (2023)
- [CodeGen: An Open Large Language Model for Code with Multi-Turn Program Synthesis](https://mlsanity.com/paper/2203.13474/) (2022)
- [Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads](https://mlsanity.com/paper/2401.10774/) (2024)
- [Offline Reinforcement Learning as One Big Sequence Modeling Problem](https://mlsanity.com/paper/2106.02039/) (2021)
- [Better & Faster Large Language Models via Multi-token Prediction](https://mlsanity.com/paper/2404.19737/) (2024)
- [On the Planning Abilities of Large Language Models : A Critical Investigation](https://mlsanity.com/paper/2305.15771/) (2023)
- [Towards Revealing the Mystery behind Chain of Thought: A Theoretical Perspective](https://mlsanity.com/paper/2305.15408/) (2023)
- [Stabilizing Transformers for Reinforcement Learning](https://mlsanity.com/paper/1910.06764/) (2019)
- [Can Language Models Solve Graph Problems in Natural Language?](https://mlsanity.com/paper/2305.10037/) (2023)
- [GraphGPT: Graph Instruction Tuning for Large Language Models](https://mlsanity.com/paper/2310.13023/) (2023)
- [GPT4Graph: Can Large Language Models Understand Graph Structured Data ? An Empirical Evaluation and Benchmarking](https://mlsanity.com/paper/2305.15066/) (2023)
- [The Parallelism Tradeoff: Limitations of Log-Precision Transformers](https://mlsanity.com/paper/2207.00729/) (2022)
- [The pitfalls of next-token prediction](https://mlsanity.com/paper/2403.06963/) (2024)
- [Hydra: Sequentially-Dependent Draft Heads for Medusa Decoding](https://mlsanity.com/paper/2402.05109/) (2024)
- [GraphLLM: Boosting Graph Reasoning Ability of Large Language Model](https://mlsanity.com/paper/2310.05845/) (2023)
- [Physics of Language Models: Part 1, Learning Hierarchical Language Structures](https://mlsanity.com/paper/2305.13673/) (2023)
- [Evaluating Cognitive Maps and Planning in Large Language Models with CogEval](https://mlsanity.com/paper/2309.15129/) (2023)
- [Parsel: Algorithmic Reasoning with Language Models by Composing Decompositions](https://mlsanity.com/paper/2212.10561/) (2022)
- [GraphInstruct: Empowering Large Language Models with Graph Understanding and Reasoning Capability](https://mlsanity.com/paper/2403.04483/) (2024)

## Built on by (most cited first, 6 of 6)

- [Self-Distillation for Multi-Token Prediction](https://mlsanity.com/paper/2603.23911/) (2026)
- [Toward Consistent World Models with Multi-Token Prediction and Latent Semantic Enhancement](https://mlsanity.com/paper/2604.06155/) (2026)
- [Understanding Parallel Samplers in Masked Diffusion via Random Walks on Graphs](https://mlsanity.com/paper/2606.22976/) (2026)
- [How Transformers Learn to Plan via Multi-Token Prediction](https://mlsanity.com/paper/2604.11912/) (2026)
- [SPLIT: Cross-Lingual Empathy and Cultural Grounding in English and Ukrainian LLM Responses](https://mlsanity.com/paper/2607.02049/) (2026)
- [EchoRec: Multi-Item Prediction-Empowered Generative Recommendation via Cycle-Consistent Preference Alignment](https://mlsanity.com/paper/2608.14011/) (2026)

## Abstract (arXiv)

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.