Organizations: University of California, Los Angeles · Shanghai Jiao Tong University · University of Pennsylvania · RIKEN Center for Advanced Intelligence Project · The Institute of Statistical Mathematics
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.
Figures & tables
Figure 1: Overview of the MTP architecture and its advantage in reasoning tasks. (a) The MTP architecture. Building upon Gloeckle et al. (2024) , our MTP employs a shared backbone with multiple independent output heads to predict several future tokens simultaneously. (b) Illustration of Star Graph and Binary Tree. In standard forward planning, early steps can be difficult due to a large search space. MTP facilitates “reverse reasoning” by looking ahead to the goal ( G ), transforming complex forward searches into easy backward steps.
Figure 2: Path-finding accuracy of MTP versus NTP under varying data and parameter scales. (a, b) 2-path 5-node star graph: test accuracy as a function of training data size (16M model) and parameter size (0.5M data). (c, d) 5-node binary tree: test accuracy as a function of data size (16M model) and parameter size (1M data). Error bars represent 90% confidence intervals calculated from three independent replicates ( n=3 ). Notably, while NTP stagnates at 50% accuracy on star graph, it begins to learn and scale on binary trees. Across both topologies, MTP consistently demonstrates superior planning abilities.
TASK
NTP
2-MTP
3-MTP
4-MTP
5-MTP
6-MTP
7-MTP
Countdown
60.27
60.36
63.20
62.09
62.75
63.17
64.93
3-SAT
10.40
28.17
69.50
63.10
82.83
82.00
87.47
Table 1: Test accuracy of models trained with NTP versus MTP configurations on Countdown and 3-SAT tasks. Models trained with MTP achieve consistently higher test accuracy than the NTP baseline across both tasks.
Figure 3: Mechanism comparison between the minimal structures of NTP and 2-MTP on the star graph task. (a–d) Query-key weights ( W(1) , W(2) ) and attention patterns ( attn(1) , attn(2) ) for NTP. (e–h) Corresponding weights and attention patterns for 2-MTP. Attention heatmaps are computed on a sequence with input graph [3,7,6,0,7,2,3,6,0,3] and target [3,6,0] . When predicting the difficult intermediate node, MTP (g) exhibits a clear predecessor-pointing pattern consistent with the reverse reasoning circuit ( Theorem 1 ): Layer 1 attends sharply to the end node, whereas NTP (c) shows no such structure. (i) The output projection V0 for NTP (identical in MTP) activates exclusively on the content block of the Layer 2 output. (j) The MTP-specific projection V1 activates on the content block of the Layer 1 output. These block-sparse structures motivate the reduced model in Section 5.1 .
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 4: Structural representations of reasoning tasks. (a) Countdown: An expression tree showing the hierarchical computation to reach the target 19 using the numbers {11,14,40,97} . The internal nodes represent arithmetic operators, and the root node represents the final result. (b) SAT: A bipartite factor graph representing a Boolean Satisfiability problem \quantity(¬x1)∧\quantity(x1∨¬x2)∧\quantity(x1∨x2∨x3) . Nodes {x1,x2,x3} denote variables, and nodes {c1,c2,c3} denote logical clauses. Solid lines represent positive literals, while dashed lines indicate negated literals.
Figure 5: Attention heatmaps for the model trained with the standard NTP objective. The NTP model exhibits severe overfitting, achieving a training accuracy of 97% while the test accuracy only reaches 20%. Furthermore, the visualization demonstrates that the model primarily attends to the start node, specifically in Layers 7 and 8.
Figure 6: Attention heatmaps for the model trained with the 3-MTP objective. In contrast to the NTP model, the 3-MTP model achieves perfect generalization with 100% accuracy on both the training and test sets. Notably, the visualization highlights that the 3-MTP model primarily attends to the end node, predominantly in Layers 3 and 4.
Multi-token prediction trains models to forecast several future tokens at once, supporting more efficient generation and richer predictive representations.
Next-token prediction (NTP) is the standard pretraining objective for large language models, yet it provides an explicit training signal only for the immediate next token, which can lead models to exploit local patterns instead of capturing longer-range structure and ideas. Multi-token prediction (MTP) addresses this by training models to predict several future tokens. However, existing MTP methods often introduce a large number of new parameters with limited improvements in downstream performance. Latent MTP approaches address this efficiency issue by encoding future tokens into a vector representation. However, these approaches usually rely on external helper models for future token encoding. We propose LightMTP, a lightweight, i.e., parameter-efficient, latent MTP approach that bootstraps the future token representations from the model's own hidden states. Our two LightMTP variants extend supervision to more future tokens without requiring the additional computational overhead of conventional MTP nor the external supervision latent MTP normally relies on. LightMTP adds at most 1% extra parameters, retains better performance on general language modeling benchmarks, and achieves similar gains in planning, coding, and reasoning.
Tamara Czinczoll, Julie Kallini, Gerard de Melo +1
Hasso Plattner Institute / University of Potsdam · Stanford University · Tel-Aviv University
Looped transformers have emerged as a parameter-efficient alternative to scaling depth for strong reasoning. By reusing one stack of layers across T iterations, they attain the effective depth and reasoning capabilities of larger models at a fixed parameter count. Yet existing approaches suffer from latent overthinking and undifferentiated computation, largely because intermediate representations receive no guidance across loops. Multi-token prediction (MTP) supplies exactly the dense, forward-looking supervision the loop is missing. We propose LoopMTP, which links the two through a structural correspondence in latent space: a model that loops T times can anticipate T future tokens. LoopMTP realizes this by softly aligning the hidden state of loop t with the embedding of the token t steps ahead, while a lightweight gate preserves useful information across iterations. LoopMTP improves average accuracy by up to 8.1% (relative) over the non-looped baseline, with training remaining stable for up to 15 loops.
Behzad Shomali, Markus Frey, David Berghaus +2
Lamarr Institute · University of Bonn · Fraunhofer IAIS
Existing policy-gradient methods for auto-regressive language models typically select subsequent tokens one at a time as actions in the policy. While effective for many generation tasks, such an approach may not fully capture the structure of complex reasoning tasks, where a single semantic decision is often realized across multiple tokens--for example, when defining variables or composing equations. This introduces a potential mismatch between token-level optimization and the inherently block-level nature of reasoning in these settings. To bridge this gap, we propose Multi-token Policy Gradient Optimization (MPO), a framework that treats sequences of K consecutive tokens as unified semantic actions. This block-level perspective enables our method to capture the compositional structure of reasoning trajectories and supports optimization over coherent, higher-level objectives. Experiments on mathematical reasoning and coding benchmarks show that MPO outperforms standard token-level policy gradient baselines, highlight the limitations of token-level policy gradients for complex reasoning, motivating future research to look beyond token-level granularity for reasoning-intensive language tasks.
Mufan Xu, Kehai Chen, Xuefeng Bai +4
School of Computer Science and Technology, Harbin Institute of Technology, China · Baidu Inc., Beijing, China