Multi-Token Prediction and Alternative Tokenization Improve Learning Latent Multi-Hop Graph Search at Scale

Arvid Frydenlund; FuTe Wong; Katherine Lambert; Frank Rudzicz

cs.LGOct 22, 2026

Multi-Token Prediction and Alternative Tokenization Improve Learning Latent Multi-Hop Graph Search at Scale

Authors: Arvid Frydenlund, FuTe Wong, Katherine Lambert, Frank Rudzicz

Organizations: University of Toronto, Vector Institute, Manulife Applied AI Research · University of Toronto, Vector Institute · University of Toronto · Dalhousie University, Vector Institute

Abstract

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.

Figures & tables

Figure 1: A path from 14 to 8, s.t. 14, 61, 33, and 47 are choice nodes. The CHC absorbs supervision from non-choice nodes (zigzag red) under NTP. Whereas MTP induces a dense quadratic latent multihop supervision supergraph (dashed rainbow). The rainbow displays the hop-length mixture. This supergraph is agnostic to the underlying topology (See Fig. 10 in E for the full example).

Panel of Figure 2 — Figure 2: Baseline ( $|\mathcal{N}|\in\{30,\,40,\,50\}$ ) and ablation ( $|\mathcal{N}|=50$ ) results using prefix LMs on ERGs for the shortest-path task. Each experiment represents different seeds, illustrating the large variance.

Panel of Figure 3 — Figure 3: STAN and $N$ results using prefix LMs on ERGs. The experiments at $|\mathcal{N}|=50$ (blue) are also shown in Fig. 2 , demonstrating a dramatic improvement in learning and variance compared to the baselines. This then allows us to scale to much larger graph sizes: $50$ through $75,\,100$ , and $125$ . Fig. 2 and Fig. 4 display a max of 75k steps vs. 25k here, showing a large difference between models.

Panel of Figure 4 — Figure 4: STAN and $N$ results using causal LMs on ERGs, for comparison to prefix LMs in Fig. 3 . We see that MTP halves the cost of learning compared to NTP, thus demonstrating its importance to scaling, even if the tokenization methods account for the largest empirical gains.

Panel of Figure 5 — Figure 5: STAN and $N$ results on PSGs. Each colour increments $d$ . Notice that in the STAN and $N$ setting, MTP makes the task trivial to learn across all scales (please zoom in for the MTP results).

Panel of Figure 6 — Figure 6: The PSG experiment where $k=9$ and $d=8$ from Fig. 5 . Here, three seeds learn the task, and two seeds do not. Jumps in SACC indicate learning the task for a specific $k^{\prime}$ and $d^{\prime}$ graph size.

Appendix figures & tables27 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 7: An illustration of the work’s scope

Figure 8: A path-star graph with query,

Q=(29,\,23)

. Node 29 is the sole choice node, with 55 and 23 predictable by shortcuts.

Figure 9: A directed graph with

Q=(33,\,2)

. 33, 6, and 8 are choice nodes. 21 is not, despite having out-degree 2; as 33 is in the path, it is removed from consideration. The backward edge from 21 to 33 is why the choice set must be defined over unseen neighbours rather than by out-degree alone.

Panel of Figure 10 — Figure 10: The Benefits of MTP. MTP supervision (c) is dense, unlike NTP (b), and topology-invariant as shown by (d)-(e).

Panel of Figure 11 — Figure 11: Fragmented routing. A sparse graph with $|\mathcal{V}|=26$ , $|\mathcal{E}|=33$ and $k=9$ . Panels (c)–(f) show the graph as a stream of edges under a causal LM. As information can only be routed across an edge once that edge has been observed, routing will be fragmented (with dotted edges and faded nodes not yet observed). The strip above each of (c)–(f) reproduces the stream of $E$ from (b) at a given time in the stream.

Panel of Figure 12 — Figure 12: Attention pattern (causal vs. prefix) and tokenization (SEAN vs. STAN). The bottom row of every panel is the annotation channel : SEAN fills it with positional annotations $\phi$ , whereas STAN overwrites it with structural markers ( $\mid$ for edge boundaries, $/$ for node-list entries). The shaded region marks what the highlighted token ( $\blacktriangledown$ ) can attend to.

Figure 13: The Plackett-Luce loss of the MTP head (Alg. 1 ), with concrete example logits. A gather puts the logits into the rank order

r_{i}

, which for MTP is the temporal order: position

j

holds the token

j

steps ahead. Bar widths are the softmax of the logits shown, on one absolute scale across all three rows, so mass leaving the partition function is visible in place. At rank-step

j

the already-ranked tokens are struck out and

Z_{i}^{j}

spans only the remainder: each ranked token is removed from the partition function, so the next step renormalizes and competes only over what is left, i.e. sampling without replacement in the context of it modelling a generative process.

Figure 14: Architecture. The NTP head uses the last layer

L

, feeding the usual unembedding

W_{o}^{L}

and a softmax-CE loss over the next token: the semantic distribution. It is active during both training and inference. The MTP head instead taps the second-to-last layer

L\!-\!1

, feeding a second projection

W_{o}^{L-1}

– the only new parameters (or

W_{o}^{L}

can be reused) – into a Plackett-Luce loss over the ordering of the next

n

tokens: the sequential distribution. The MTP head exists at training time only, so inference is unchanged. Squared boxes are parameters and operations, while rounded boxes are data.

Experiment	Figure	Num Sub Exps.	Num $k$	Num Trials	Total Runs
ERG, Prefix LM, Baselines	Fig. 2	10	10	5	500
ERG, Prefix LM	Fig. 3	8	10	5	400
ERG, Causal LM, Baseline	Fig. 20	6	10	5	300
ERG, Causal LM	Fig. 4	8	10	5	400
EG, Prefix LM	Fig. 21	8	10	5	400
PSG, Prefix LM	Fig. 5	18	10	5	900

Appendix

Table 1: Calculation of the total number of runs needed for the main experiments broken down by figure. Overlapping experiments were not double counted.

Panel of Figure 15 — Figure 15: Full-page version of Fig. 2 .

Panel of Figure 16 — Figure 16: Full-page version of Fig. 3 .

Panel of Figure 17 — Figure 17: Full-page version of Fig. 4 .

Panel of Figure 18 — Figure 18: Full-page version of Fig. 5 .

Panel of Figure 19 — Figure 19: Full-page version of Fig. 6 . Note that, unlike in bachmann2024the who did not use MMHT or a mixture of $d^{\prime}$ , the plateaus are above $1/d$ because SACC marginalizes over the $(k^{\prime},\,d^{\prime})$ mixture. That is, failure to learn the task no longer results in an accuracy of $1/d$ .

Figure 20: The baseline and ablation for the causal LM.

Figure 21: Results using Euclidean Graphs with Prefix LMs

Figure 22: Restricting the mixture of MMHT to higher

k^{\prime}

fails to learn the task.

Panel of Figure 23 — Figure 23: The Khops task , $k=4$ , one instance throughout. Colours and letters denote symbol values , the numbers denote positions i.e. node IDs. Dashed arrows are the latent content-based match while solid is the deterministic $+1$ shift.

Figure 24: Khops experiments. Prior work sampled

k^{\prime}

, providing inner sample MMHT. We show that this is necessary for learning. Using multi-label ranking provides implicit intra sample MMHT.

Figure 25: A simplified version of Fig. 24 with the where we have removed the ‘- inner sample MMHT’ experiment and any experiment which did not converge (‘x’s at top). This better allows us to see the effect of the ranking loss vs. the baseline.

Figure 26: Khops-Gen. Sequential generation of the full compositional chain from the single query

x_{N}

(position 31). The source-side is on the left in red, with the target-side being generated. Unlike Khops, only one chain is considered in a single instance of the task. Because of this, and as the source-side spans the entire

N

nodes, there is no implicit MMLT training as in Khops. Thus, MMLT must instead be induced explicitly by sampling

N^{\prime}\leq N

( L.1.1 ). This also indicates that there is much less supervision in Khops-Gen as there is in Khops. However, the supervision is different. Whereas Khops has implicit MMLT but no implicit MMHT, Khops-Gen is the opposite. Unlike Khops, the chain runs until the context is exhausted and

k

is not supplied to the model. Note that we sample

k^{\prime}

uniformly, but that does not mean the context window between hops is the same size. On the target side, we show both NTP and MTP. Because this is a genuine generation task, MTP applies directly rather than through the ranking-loss stand-in required for Khops ( L.1.2 ). Importantly, MTP induces the same quadratic-to-

k

supergraph (rainbow) as it does in path search (and will for any multihop reasoning task). The MTP target block is a triangle of

k(k+1)/2

terms, matching Eq. 3 .

Figure 27: Khops-Gen experiments. MTP improves learnability at scale. Explicit MMLT is also necessary for efficiently learning the task.

Figure 28: A simplified version of Fig. 27

Figure 29: BFS-SP SP-only results. Given the SP, generate the shortest path, i.e. the partial task. This task is (nearly) computationally equivalent to Khops-Gen.

Figure 30: BFS-SP results. This is both the generation of the SP and the generation of the shortest path, i.e. the full SP task.

Panel of Figure 31 — Figure 31: The BFS ordering inductive bias changes the topology of the task and reduces the latent reasoning graph to a linear chain.

Panel of Figure 32 — Figure 32: Multiple valid shortest paths from $s=9$ to $t=52$ require label smoothing. Light marron solid-border nodes are choice nodes , dark dashed-border nodes are not, and grey nodes have been removed from consideration by conditioning on $\pi_{\leq}$ . The sampled ground-truth path $\pi$ is drawn with heavy edges. Below each choice node we give the NTP one-hot target and the correct uniform target over all valid next nodes, using the same encoding. The NTP row greys out the unchosen nodes, asserting they are off-path even though each begins an equally valid shortest path. Smoothing is needed only at the choice nodes, so the two rows coincide in (b).

LLM Reasoning with Graphs

Large language model reasoning with graphs uses structured relationships to support multi-step inference, evidence retrieval, and planning.

See all 209 papers

Revisiting Graph-Tokenizing Large Language Models: A Systematic Evaluation of Graph Token Understanding
May 5, 2026Zhongjian Zhang, Yue Yu, Mengmei Zhang +3LLM Evaluation LLM Reasoning with Graphs
Beijing University of Posts and Telecommunications · China Telecom Bestpay · Beihang University
LightMTP: Lightweight Latent Multi-Token Prediction
Oct 5, 2026Tamara Czinczoll, Julie Kallini, Gerard de Melo +1Language Model Pretraining Multi-Token Prediction
Hasso Plattner Institute / University of Potsdam · Stanford University · Tel-Aviv University
Scaling Graph Reasoning in Language Models with Learnable Graph Positional Encodings
Oct 20, 2026Javier Porras-Valenzuela, Arush Arora, Alejandro RibeiroLLM Reasoning with Graphs Graph Positional Encoding
University of Pennsylvania

May 5, 2026J/K move · Enter open · S save

Revisiting Graph-Tokenizing Large Language Models: A Systematic Evaluation of Graph Token Understanding

Zhongjian Zhang, Yue Yu, Mengmei Zhang, Junping Du+2

Beijing University of Posts and Telecommunications · China Telecom Bestpay · Beihang University

LLM Evaluation LLM Reasoning with Graphs Graph Foundation Models LLMs for Graph Learning

Open paper AskPDF arXiv

The remarkable success of large language models (LLMs) has motivated researchers to adapt them as universal predictors for various graph tasks. As a widely recognized paradigm, Graph-Tokenizing LLMs (GTokenLLMs) compress complex graph data into graph tokens and treat them as prefix tokens for querying LLMs, leading many to believe that LLMs can understand graphs more effectively and efficiently. In this paper, we challenge this belief: Do GTokenLLMs fully understand graph tokens in the natural-language embedding space? Motivated by this question, we formalize a unified framework for GTokenLLMs and propose an evaluation pipeline, GTEval, to assess graph-token understanding via instruction transformations at the format and content levels. We conduct extensive experiments on 6 representative GTokenLLMs with GTEval. The primary findings are as follows: (1) Existing GTokenLLMs do not fully understand graph tokens. They exhibit over-sensitivity or over-insensitivity to instruction changes, and rely heavily on text for reasoning; (2) Although graph tokens preserve task-relevant graph information and receive attention across LLM layers, their utilization varies across models and instruction variants; (3) Additional instruction tuning can improve performance on the original and seen instructions, but it does not fully address the challenge of graph-token understanding, calling for further improvement.

Abstract

LLM Reasoning with Graphs

Explore similar work

Revisiting Graph-Tokenizing Large Language Models: A Systematic Evaluation of Graph Token Understanding

LightMTP: Lightweight Latent Multi-Token Prediction

Scaling Graph Reasoning in Language Models with Learnable Graph Positional Encodings