Better & Faster Large Language Models via Multi-token Prediction

Fabian Gloeckle; Badr Youbi Idrissi; Baptiste Rozière; David Lopez-Paz; Gabriel Synnaeve

Source-linked AI summary

Better & Faster Large Language Models via Multi-token Prediction

Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Rozière, David Lopez-Paz, Gabriel Synnaeve

arXiv:2404.19737v1Apr 30, 2024cs.CL

TL;DR

Next-token prediction may inefficiently acquire language, knowledge, and reasoning capabilities because teacher forcing focuses on short-term predictions. The paper trains models to predict multiple future tokens with independent heads over a shared trunk, reporting stronger downstream performance and up to 3× faster inference, while leaving choices such as the prediction horizon for future work.

Problem
Next-token prediction remains an inefficient way to acquire language, world knowledge, and reasoning capabilities, while teacher forcing emphasizes short-term predictions over longer-term dependencies.
Method
The method predicts n future tokens at each training position using independent output heads on top of a shared model trunk, without train-time or memory overhead.
Results
Experiments show multi-token prediction is increasingly useful for larger models, with strong improvements on code tasks; exact inference becomes up to 3 times faster with speculative decoding.
Takeaways & Limitations
Multi-token prediction offers a simple auxiliary training objective for stronger and faster transformer models, particularly on generative and reasoning tasks.
Takeaways & Limitations
The paper leaves automatic selection of n, vocabulary-size tuning, and improved embedding-space auxiliary losses for future work.

Abstract

from arXiv · show

Large language models such as GPT and Llama are trained with a next-token prediction loss. In this work, we suggest that training language models to predict multiple future tokens at once results in higher sample efficiency. More specifically, at each position in the training corpus, we ask the model to predict the following n tokens using n independent output heads, operating on top of a shared model trunk. Considering multi-token prediction as an auxiliary training task, we measure improved downstream capabilities with no overhead in training time for both code and natural language models. The method is increasingly useful for larger model sizes, and keeps its appeal when training for multiple epochs. Gains are especially pronounced on generative benchmarks like coding, where our models consistently outperform strong baselines by several percentage points. Our 13B parameter models solves 12 % more problems on HumanEval and 17 % more on MBPP than comparable next-token models. Experiments on small algorithmic tasks demonstrate that multi-token prediction is favorable for the development of induction heads and algorithmic reasoning capabilities. As an additional benefit, models trained with 4-token prediction are up to 3 times faster at inference, even with large batch sizes.

1. Introduction

The paper argues that next-token prediction is inefficient because teacher forcing emphasizes local patterns and overlooks hard decisions. It proposes multi-token prediction as a simple auxiliary objective that improves sample efficiency, code performance, and inference speed.

Next-token prediction can overlook hard decisions and require orders of magnitude more data than human children for comparable fluency.
Multi-token prediction trains models to predict n future tokens from each corpus position simultaneously and in parallel.
The proposed architecture adds multi-token prediction without train-time or memory overhead.
Models up to 13B parameters solve around 15% more code problems on average than comparable next-token predictors.
Multi-token prediction enables self-speculative decoding, making models up to 3 times faster across a wide range of batch sizes.
The method is presented as a cost-free modification for training stronger and faster transformer models.

2. Method

The method generalizes language modeling from predicting one future token to predicting n future tokens with independent heads over a shared representation. Sequential head computation controls memory use, while the extra heads can also accelerate autoregressive inference through speculative decoding.

Standard language modeling minimizes cross-entropy by predicting the next token from the preceding context.
Multi-token prediction instead asks the model to predict n future tokens simultaneously at every training-corpus position.
A shared trunk produces the context representation, which n independent heads use to predict future tokens in parallel.
Sequentially computing each head’s forward and backward pass reduces peak GPU memory from O(nV + d) to O(V + d).
During inference, the next-token head supports vanilla autoregressive generation, while the additional heads enable self-speculative decoding.

3. Experiments on real data

Across real-data experiments, multi-token prediction becomes more effective at larger scales and improves code generation, inference speed, and several finetuning outcomes. Benefits are mixed for 7B choice tasks but positive for summarization, while performance gains persist after repeated epochs.

3.1. Benefits scale with model size: Multi-token prediction becomes beneficial at scale, outperforming next-token baselines on code benchmarks under the same computational budget.Experiments span six model sizes from 300M to 13B parameters, with evaluations on MBPP and HumanEval.
3.2. Faster inference: 3.0× faster code decoding and 2.7× faster text decoding were achieved with a 7B 4-token prediction model using self-speculative decoding.The 7B model averaged 2.5 accepted tokens out of 3 suggestions on code; an 8-byte model reached 6.4× speedup.
3.3. Learning global patterns with multi-byte prediction: 67% more MBPP pass@1 problems and 20% more HumanEval pass@1 problems were solved by the 7B 8-byte prediction model than by next-byte prediction.The model was trained on 314B bytes, equivalent to around 116B tokens, and approached token-based model performance with 1.7× less data.
3.4. Choosing the prediction window: At 7B scale, 4-future-token training led the ablation on MBPP and HumanEval, while APPS/Intro favored a 6-token window.For MBPP, improvements were +3.8%, +2.1%, and +3.2% on pass@1, pass@10, and pass@100; HumanEval gains were +1.2%, +3.7%, and +4.1%.
3.5. Training for multiple epochs: Multi-token training retained a +2.4% MBPP pass@1 gain and a +3.2% HumanEval pass@100 gain after training on multiple epochs of the same data.The improvements diminished, with similar performance on the remaining reported metrics.
3.6. Finetuning multi-token predictors: Both finetuning variants of a 4-token pretrained model outperformed the next-token baseline on CodeContests pass@k, with next-token finetuning appearing best overall.The comparison used 7B models and the challenging CodeContests benchmark; the temperature oracle selected the best temperature for each k.

4. Ablations on synthetic data

Synthetic-data experiments show that multi-token prediction improves induction and algorithmic reasoning, especially for smaller models and out-of-domain arithmetic generalization.

4.1. Induction capability: Induction measures whether models predict the second token of previously mentioned two-token character names.The first name token depends on preceding semantics, while the second tests pure induction after the name has appeared.
4.1. Induction capability: Multi-token prediction improves induction capability most strongly in models with 30M or fewer nonembedding parameters.Its advantage disappears at 100M nonembedding parameters and above.
4.1. Induction capability: The induction benchmark cannot reach a perfect score because some evaluation-name tokens never occur in training and embedding and unembedding parameters are untied.Once induction is already formed, multi-token prediction can hurt this restricted benchmark because the task becomes locally solvable.
4.2. Algorithmic reasoning: Multi-token prediction improves polynomial-arithmetic accuracy across task difficulties and produces particularly strong out-of-domain generalization despite low absolute scores.Replacing next-token prediction has a larger effect than increasing model size from 30M to 100M parameters.

5. Why does it work? Some speculation

The paper speculates that multi-token prediction improves generation by emphasizing decisions whose consequences affect later tokens and by reducing teacher-forcing mismatch.

Why does it work?: Multi-token prediction may reduce the distributional discrepancy between teacher-forced training and autoregressive generation.Teacher forcing supplies future ground-truth tokens during training, whereas test-time generation is unguided and errors can accumulate.
Lookahead reinforces choice points: Multi-token prediction implicitly assigns larger weights to tokens whose successors are more strongly correlated.Choice points receive greater effective emphasis through the difficulty of predicting their consequences.
Lookahead reinforces choice points: For n-token prediction, choice points receive weight n(n+1)/2, whereas inconsequential points receive the smaller weight n.The authors therefore speculate that n-token losses promote decisions important for the quality of later text.
Information-theoretic argument: With n = 2, the information-theoretic decomposition doubles the importance of mutual information I(X; Y) relative to vanilla next-token prediction.Here X and Y denote the next and second-next future tokens, respectively.

6. Related work

Related work spans denoising, permuted language modeling, multi-token prediction, speculative decoding, and multi-task learning, while distinguishing this paper’s architecture and training setup.

Language modeling losses: Earlier denoising and permuted-language objectives mix attention patterns or predict spans and sequence parts using future information.These objectives train on only a small percentage of input tokens, averaging 15% in the cited examples.
Multi-token prediction in language modelling: Qi et al. propose multi-token prediction for planning and better representations, whereas this work avoids residual-stream replication and enables compute-matched comparisons.The present architecture makes residual representations participate more directly in auxiliary loss terms.
Self-speculative decoding: The paper’s self-speculative decoding architecture uses transformer layers as prediction heads and reorganizes forward/backward computation to use all loss terms.This contrasts with earlier approaches that use linear prediction heads or stochastically select one head for loss computation.
Multi-target prediction: Multi-target prediction frames joint training as auxiliary-task learning that exploits dependencies between target variables.The paper situates multi-token prediction within the broader multi-task-learning literature.

7. Conclusion

The conclusion presents multi-token prediction as an improvement over next-token training for generative and reasoning tasks, with growing value at larger scales and faster exact inference via speculative decoding.

Conclusion: Experiments up to 7B parameters and 1T tokens show that multi-token prediction becomes increasingly useful for larger models, especially on code tasks.The authors posit that it reduces distribution mismatch between teacher-forced training and autoregressive generation.
Conclusion: Exact inference becomes 3 times faster when multi-token prediction is combined with speculative decoding.This conclusion reports the method’s principal inference benefit.
Future work: Future work includes automatically choosing n, tuning vocabulary sizes, and developing auxiliary prediction losses in embedding spaces.These directions indicate unresolved design choices and possible efficiency trade-offs.

Impact statement

The paper aims to improve language-model compute and data efficiency, while cautioning that environmental and societal effects require careful consideration.

The paper’s goal is to make language models more compute and data efficient.
Improved efficiency may in principle reduce the ecological impact of training language models.
The paper cautions that rebound effects and the broader societal advantages and risks of language models should be considered.

Environmental impact

The supplied results examine inference speed, alternative architectures, training cost, finetuning, scaling, and code evaluation, with benefits depending on the design and setting.

Alternative architectures: Alternative anticausal, causal, and linear architectures showed no significant improvement over the parallel architecture.
Alternative architectures: Replicating the unembedding matrix n times is prohibitive for large-scale training because it requires matrices with shapes (d, nV).
Training cost: A slight training-time overhead arose from suboptimal Fully Sharded Data Parallel use and could be removed through correct reimplementation.
Finetuning: Finetuning Llama 2 with 4-token prediction did not significantly improve performance compared with the baseline.
Code evaluation: The code evaluation generates 1000 samples per problem across temperatures T ∈{0.5, 0.6, 0.7, 0.8, 0.9} and estimates pass@k.

G. Additional results on natural language benchmarks

Additional natural-language benchmark results show that multi-token prediction does not improve performance for the evaluated 7B models on standard choice and knowledge tasks.

The evaluated benchmarks include ARC Challenge, COPA, Hellaswag, Natural Questions, PIQA, SIQA, and TriviaQA.
7B models trained on 200B language-data tokens show no performance improvement from 2-token prediction on standard NLP benchmarks.
The 4-future-token model regresses slightly on the evaluated standard NLP benchmarks.

H. Additional results on abstractive text summarization

Additional natural-language experiments cover summarization, mathematical reasoning, and induction capability, with summarization gains and earlier, stronger induction reported in selected settings.

Abstractive summarization: The summarization evaluations report ROUGE-n and ROUGE-L precision, recall, and F1 averages across datasets.
Abstractive summarization: Both 2-token and 4-token prediction models outperform the next-token baseline on the reported abstractive summarization averages.
Abstractive summarization: For models trained on 500B tokens, 4-token prediction appears better on recall metrics while 2-token prediction appears better on precision metrics.
Mathematical reasoning: The GSM8K evaluation measures the frequency that the correct final answer appears among k samples for k = 1, 10, 100.
Induction capability: Training on a 9:1 books-to-children’s-story mix produces induction capability earlier and to a higher degree.

K. Additional results on algorithmic reasoning

Additional experiments and analyses examine how multi-token prediction affects computation sharing, algorithmic reasoning, and sequence prediction beyond standard next-token training.

Computation sharing: Multi-token prediction models outperform next-token models on polynomial arithmetic variants across task difficulties and model sizes.With pause tokens added before answers, the performance gap does not show strong evidence of widening or shrinking.
Computation sharing: Whitespace-based pause tokens give multi-token prediction models a slight but marginal advantage on HumanEval and MBPP.Spaces and newlines provide additional computation in a more natural code-benchmark setting.
Computation sharing: Tripling model size has a smaller performance effect than replacing next-token prediction with multi-token prediction on polynomial arithmetic.The comparison uses 30M- and 100M-parameter models, with two independent runs per configuration.
Information-theoretic analysis: The proposed decomposition interprets 2-token prediction as combining local cross-entropy, weighted mutual information, and shifted next-token cross-entropy terms.The added future-token term incentivizes precomputing features useful for predicting the later token.
Factorization orders: A 4-token prediction loss explicitly encourages latent representations to contain information about successive tokens, facilitating non-causal factorization within activations.This is motivated by dependencies such as grammatical agreement, where later sentence elements can constrain earlier choices.