Beyond Multi-Token Prediction: Pretraining LLMs with Future Summaries

TL;DR

This work identifies key limitations of standard next-token prediction and existing multi-token prediction approaches in capturing long-horizon reasoning. It proposes Future Summary Prediction (FSP), which uses a single auxiliary head to predict a compact representation of the long-term future, using either hand-crafted multi-hot summaries or a learned summary from a reverse language model. Across 3B and 8B parameter scales, FSP demonstrates robust improvements over NTP and MTP on math, coding, and reasoning benchmarks, with learned summaries (FSP-RevLM) offering the strongest gains and faster convergence in adaptive tasks. The results suggest that focusing on abstract, predictive aspects of the future can provide a scalable and effective pretraining objective for next-generation LLMs, with broad implications for long-horizon generation and planning tasks.

Abstract

Next-token prediction (NTP) has driven the success of large language models (LLMs), but it struggles with long-horizon reasoning, planning, and creative writing, with these limitations largely attributed to teacher-forced training. Multi-token prediction (MTP) partially mitigates these issues by predicting several future tokens at once, but it mostly captures short-range dependencies and offers limited improvement. We propose future summary prediction (FSP), which trains an auxiliary head to predict a compact representation of the long-term future, preserving information relevant for long-form generations. We explore two variants of FSP: handcrafted summaries, for example, a bag of words summary of the future of the sequence, and learned summaries, which use embeddings produced by a reverse language model trained from right to left. Large-scale pretraining experiments (3B and 8B-parameter models) demonstrate that FSP provides improvements over both NTP and MTP across math, reasoning, and coding benchmarks.

Figures (5)

• Figure 1: A comparison of future-aware pretraining objectives. All methods take a prefix $x_{\leq t}$ as input. NTP: Only predicts the immediate next token. MTP: Uses multiple auxiliary heads, each predicting a specific future token. FSP-BCE: Our proposed hand-crafted summary method that predicts a multi-hot "bag-of-tokens" summary of a long future window using a single auxiliary head. FSP-RevLM: Our proposed learned summary method predicts a compact hidden representation of the future, which is generated by a reverse language model (RevLM).
• Figure 2: An abstraction of the architecture that subsumes NTP, MTP, and FSP.
• Figure 3: Analysis of FSP-BCE on the path-star task, which tests long-horizon planning.Left: Accuracy (mean over $5$ random seeds) of different pretraining objectives on degree $2$ graphs with path lengths $6$ and $8$. Standard NTP generalizes poorly, and MTP's accuracy degrades as the path length increases, while FSP-BCE achieves perfect accuracy. Right: An illustration of why NTP fails while FSP-BCE succeeds, where the input context is shown in grey and the target is in beige.
• Figure 4: Analysis on the modified sibling discovery task which requires adaptive future summaries.Left: Convergence speed (mean (s.e.) over $3$ random seeds) relative to NTP, where lower values imply faster convergence. FSP-RevLM converges faster than NTP while FSP-BCE improves only in the cases with few components. Right: Task setup illustration-- given the prefix, only future tokens in the highlighted component are informative. FSP-BCE, which summarizes all future tokens, suffers from irrelevant information, whereas FSP-RevLM summarizes only the informative aspects.
• Figure 5: Enhanced diversity through learned future summaries. We compare standard MTP, which predicts the immediate future token, with FSP-RevLM, which enriches the auxiliary head target using learned future summaries. FSP-RevLM substantially increases output diversity compared to MTP on GSM8K and MATH benchmarks.