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.