Multi-Token Prediction and Alternative Tokenization Improve Learning Latent Multi-Hop Graph Search at Scale
Language models fail to learn shortest-path search on unseen symbolic graphs as graphs scale beyond about 50 nodes, a wall this paper diagnoses and breaks through. The authors show that next-token prediction provides too little and too skewed supervision, while standard tokenization introduces routing and permutation artifacts. By combining multi-token prediction—which densely supervises all sub-paths and implicitly mixes hop lengths—with structural annotation (STAN) and a node list (+NN), they stably train small transformers on graphs up to 125 nodes with far lower variance. The work makes a strong case that objective design and input formatting matter more than model size for latent multi-hop reasoning.
The paper delivers a compelling theoretical-empirical package. It formalizes supervision degradation via choice nodes, proves that MTP achieves near-topology-independent optimal sequential sample efficiency, and shows that tokenization changes (STAN and +NN) produce the largest practical gains in stability and scale. Yet the explanatory story for those tokenization gains is incomplete: the authors candidly note that their ablations contradict the motivating hypotheses, since +NN also helps prefix LMs and STAN also helps causal LMs. The scope is deliberately limited to in-domain synthetic graphs, which is justified as a prerequisite, but it means the broad conjecture about natural-language multi-hop reasoning remains speculative.
The theoretical analysis of sequential sample efficiency (SSE) is rigorous and clear. Extending the Clever Hans Cheat from path-star graphs to arbitrary topologies via choice nodes $c(\pi)$ is elegant, and the proof that MTP provides $k(k+1)/2 - (k - c(\pi))$ supervision terms—quadratic in path length and nearly topology-agnostic—is a genuine contribution. Empirically, the variance reduction from STAN and the scaling to $|\mathcal{N}|=125$ nodes, in roughly $1/3$ the steps of baselines, are substantial and well documented.
The central gap is between the tokenization methods and their motivating hypotheses. The authors hypothesize that +NN helps by bypassing the causal constraint and fragmented routing, yet it improves prefix LMs that are not causally masked; they hypothesize STAN helps by restoring permutation invariance, yet it improves causal LMs that cannot be permutation-invariant. They are honest about this mismatch, but it leaves the mechanism behind the largest empirical gains unexplained.
Second, the work is scoped entirely to in-domain learning on synthetic symbol-to-symbol graphs, so the conjecture that MTP transfers to natural-language multi-hop reasoning is just that—a conjecture. Third, the reported 'lockout' phenomenon on path-star graphs, where models solve easy sub-problems and then cease learning harder ones, indicates that NTP dynamics are brittle in ways not fully resolved even by MMHT.
The evidence that MTP outperforms NTP is solid across three topologies (ERG, EG, PSG) and multiple metrics, with the gap widening as $k$ and $|\mathcal{N}|$ grow. The comparison to prior work is fair: the authors replicate the scaling wall around 50 nodes reported by prior work and benchmark against the adversarial path-star results of Bachmann et al. (2024). However, the paper explicitly declines to benchmark against other MTP implementations such as Gloeckle et al. (2024) or Cai et al. (2024) on this task, framing the contribution as theoretical rather than engineering. The thorough related work discussion is slightly undermined by the convenient dismissal of discrete diffusion as 'a much harder model to study' rather than an empirical baseline.
The experimental setup is detailed in Appendix J, including $L=11$ layers, batch size 1000, curriculum sampling of $k'$ and $|\mathcal{N}|'$, and early stopping at $>95\%$ TACC. The authors report 5 seeds per experiment and evaluate on 10k held-out instances. They have released the online dataset generator and the Plackett-Luce RITF cost function ahead of publication, with the full code promised upon acceptance. Because training uses small transformers from scratch on uniquely generated synthetic graphs, the main barrier to reproduction is compute: roughly 2900 experiments were run. The omission of wall-clock times or FLOP counts is acknowledged and defended as off-topic for a paper focused on supervision efficiency, but it does make cost estimation harder for follow-up work.
Given an unseen graph, a start node, $s$, and an end node, $t$, can a language model generate the shortest path from $s$ to $t$? Prior work has found that models fail to learn this task as the graphs scale. We show that this failure has two causes: the supervision signal and the tokenization. First, we theoretically establish and empirically demonstrate the effectiveness of multi-token prediction (MTP), and show that the benefits of MTP are topologically agnostic. Second, we introduce alternative tokenizations and show that the two causes are orthogonal. Using these methods, we break through the scaling wall observed by prior work, consistently and stably learning to solve the task at scales far beyond previous limits.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.