MULTI-TOKEN Prediction Promises and results

Next Token Prediction (NTP) has long been the dominant training method for LLMs. However, the Multi-Token Prediction (MTP) approach is gaining increasing traction among researchers, driven by two main reasons:

  1. Improved model quality: Research shows that MTP achieves better results on a range of tasks, such as planning and code generation.
  2. Faster inference: If a model generates 10 tokens in a single forward pass, the theoretical speedup can approach 10x (in practice the gain is smaller, but the number illustrates the potential).

However, a key challenge of MTP remains the assumption of token independence during parallel generation. Tokens are most commonly generated in parallel by different transformer heads, which requires heuristics (e.g. speculative decoding) to compensate for inter-token dependencies. For instance, the authors of [5] attempt to work around the independence problem, though there is clearly still room for improvement.

Figure from [8].

Figure from [5].

Teacherless Learning represents an intermediate approach between NTP and MTP — in some sense the opposite of teacher-forcing. For example, when generating the word ABUNDANCE, after the character A, instead of feeding in the actual next characters, the model inserts the placeholder > three times in a row:

0: A >
1: AB
2: AB >
3: ABU >
4: ABUN
…

Unlike MTP, Teacherless Learning has a rigorous mathematical foundation. While it does not speed up inference, it exposes fundamental shortcomings of traditional NTP.

Work [7] shows that NTP with teacher-forcing struggles with tasks requiring long-range dependency prediction (e.g. Path-Finding on Path-Star Graphs). The authors of [6] state directly:

«Teacherless training and diffusion models comparatively excel in producing diverse and original output».

Interestingly, diffusion models — which are rapidly gaining popularity — also generate groups of tokens jointly, and notably show strong performance on code-related tasks.

Against this backdrop, adapting existing models to MTP is emerging as a promising research direction for resource-constrained settings.

References:

[1] On multi-token prediction for efficient LLM inference

[2] Ideas in Inference-time Scaling can Benefit Generative Pre-training Algorithms

[3] Multi-Token Prediction Needs Registers

[4] Improving Large Language Models with Concept-Aware Fine-Tuning

[5] Efficient Joint Prediction of Multiple Future Tokens

[6] Roll the dice & look before you leap: Going beyond the creative limits of next-token prediction

[7] The Pitfalls of Next-Token Prediction

[8] Better & Faster Large Language Models via Multi-token Prediction

Author
dalime
Published

HoTPP & DETPP

A review of two papers from our colleagues at Sber AI Lab: DeTPP & HoTPP.

  • What do we want?
    To predict future user behavior — for example, generating purchase categories and timestamps.
  • The problems.
    [1] This used to be solved via TPP methods, where we predict the user’s next event and its timestamp. But in practice, a single event isn’t enough. We could predict multiple events recurrently, but in that case our objective is still next-token-prediction, whereas what we actually want is next-tokens-prediction, right?
    [2] Autoregressive methods assume that events occur strictly one after another, but people behave fairly randomly even if patterns exist — on any given day you might buy bread, pay a loan, and send money to your mom, all in a random order, yet the distribution of events within the day remains the same.

HoTPP

In HoTPP, the authors articulate all these problems and propose a new benchmark along with a metric called T-mAP for comparing such temporal distributions.

The T-mAP Metric

The metric combines three ideas:

  1. We sweep a threshold — similar to how RocAuc is computed — and observe how quality changes.
  2. We align our predictions with targets rather than comparing them one-to-one.
  3. For time, we introduce a left-and-right window $\delta$ — essentially dividing the time interval into buckets of size $\delta$ .

The metric ends up with two hyperparameters: the prediction horizon and the time step size $\delta$ .

Results:

  • Statistical baselines such as mode and history-based distribution land in the top 3 methods.
  • Models frequently collapse to the mode. Autoregressive methods tend to have low entropy — focusing on a few specific, most frequent tokens.
  • Methods like HYPRO and diffusion-based approaches perform best, though they were specifically designed for long-horizon prediction.
  • Methods trained on the next-token-prediction objective perform significantly worse.
  • The authors also show that standard metrics are not very sensitive to rare events, whereas T-mAP is considerably more sensitive to them.

DeTPP

Naturally, once the problem is well understood, you can build a method to solve it. That is exactly what the authors do in the follow-up paper DeTPP.

The core idea: let’s train the model to align its predictions on the fly.

The final loss factorizes as: P(y)=P(o)P(t)P(l)

  • The probability of an event occurring at a given time P(t),
  • The probability of an event occurring at all P(o),
  • The probability of the event’s category P(l).

Once we obtain the predicted classes and timestamps, we additionally align them in a manner similar to the T-mAP metric.

The authors then describe several architectural decisions — which transformer components are used to predict each part — which I will skip here. The main hyperparameters are the prediction horizon and the allowable time shift (mirroring those of the metric).

Results

Everything performs as expected. Notably, this approach also works well for next-token prediction, matching methods that were specifically trained for that task.

Author
dalime
Published