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
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).
Figure 2: Baseline ( ∣N∣∈{30,40,50} ) and ablation ( ∣N∣=50 ) results using prefix LMs on ERGs for the shortest-path task. Each experiment represents different seeds, illustrating the large variance.
Figure 3: STAN and N results using prefix LMs on ERGs. The experiments at ∣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.
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.
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).
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′ and d′ 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.
Figure 10: The Benefits of MTP. MTP supervision (c) is dense, unlike NTP (b), and topology-invariant as shown by (d)-(e).
Figure 11: Fragmented routing. A sparse graph with ∣V∣=26 , ∣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.
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 ϕ , whereas STAN overwrites it with structural markers ( ∣ for edge boundaries, / for node-list entries). The shaded region marks what the highlighted token ( ▼ ) 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 ri , 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 Zij 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 WoL 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 WoL−1 – the only new parameters (or WoL 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.
Figure 15: Full-page version of Fig. 2 .
Figure 16: Full-page version of Fig. 3 .
Figure 17: Full-page version of Fig. 4 .
Figure 18: Full-page version of Fig. 5 .
Figure 19: Full-page version of Fig. 6 . Note that, unlike in bachmann2024the who did not use MMHT or a mixture of d′ , the plateaus are above 1/d because SACC marginalizes over the (k′,d′) 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′ fails to learn the task.
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′ , 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 xN (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′≤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′ 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.
Figure 31: The BFS ordering inductive bias changes the topology of the task and reduces the latent reasoning graph to a linear chain.
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 π≤ . The sampled ground-truth path π 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).
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.
Zhongjian Zhang, Yue Yu, Mengmei Zhang +3
Beijing University of Posts and Telecommunications · China Telecom Bestpay · Beihang University
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
Solving reasoning tasks supported on graphs typically requires providing a Language Reasoning Model (LRM) with a serialized description of the graph. This approach results in the input length growing with the number of edges, which limits the size of the problems that can be solved. Furthermore, even when graphs fit in context, performance sharply decays as the graphs grow. We propose to address this by augmenting LRMs with scalable and permutation equivariant graph positional encodings that modulate attention logits. An additional post-fusion-style module incorporates structural information into hidden states. We compute these positional encodings using a GNN followed by a Graph Transformer. This eliminates the need for serialized adjacencies, reducing the context burden to O(N). We jointly fine-tune the positional encoder and the LRM LoRA adapters using optimal solutions generated by graph-theoretic solvers, resulting in a computationally efficient method. Trained on small graphs, the model exhibits improved size-generalization across graph-theoretic tasks on graphs with up to 1,000 nodes, and produces fewer tokens than a conventional LRM. Additionally, we observe that our model remains robust to node relabelings. Evaluations on graph-theoretic benchmarks also show that performance on small graphs is preserved.
Javier Porras-Valenzuela, Arush Arora, Alejandro Ribeiro