Better & Faster Large Language Models via Multi-token Prediction

Key Takeaways:

• Multi-token prediction is a simple yet powerful modification to LLM training, improving sample efficiency and performance on various tasks.
• This approach is particularly effective at scale, with larger models showing significant gains on coding benchmarks like MBPP and HumanEval.
• Multi-token prediction enables faster inference through self-speculative decoding, potentially reaching 3x speedup compared to next-token prediction.
• The technique promotes learning global patterns and improves algorithmic reasoning capabilities in LLMs.
• While effective for generative tasks, the paper finds mixed results on benchmarks based on multiple-choice questions.

FAIR (Facebook AI Research team) researchers authored this paper, and the results look promising. I would love to see this turn into an actual product. I think this paper has all the ingredients to be the kind of groundbreaking paper that Mixture of Experts proved out to be for the current generation of models.

Introduction:

The paper starts by highlighting the limitations of current LLM training methods based on next-token prediction. Despite their impressive capabilities, these models require massive amounts of data compared to humans to achieve similar levels of fluency. The authors argue that next-token prediction focuses too much on local patterns and overlooks “hard” decisions, leading to inefficiencies in learning. They propose multi-token prediction as a solution to overcome these limitations.

Background:

Traditional language models are trained using a next-token prediction loss where the model predicts the next token in a sequence based on the preceding context. This paper proposes a more general approach where the model predicts n future tokens at once using n independent output heads connected to a shared model trunk. This forces the model to consider longer-term dependencies and global patterns in the text.

How is it different?

Unlike previous works that explored multi-token prediction, this paper:

• Proposes a memory-efficient architecture that avoids significant overhead in training time and memory usage.
• Demonstrates the benefits of the approach at scale with large models (up to 13B parameters).
• Shows that multi-token prediction allows for faster inference through self-speculative decoding.

Method:

The paper presents a simple and efficient architecture for multi-token prediction. The model consists of:

• Shared Transformer trunk: Processes the input context and generates a latent representation.
• Independent output heads: Each head predicts one of the n future tokens based on the latent representation.
• Shared unembedding matrix: Converts the predictions from each head back into tokens.

The authors also introduce a memory-efficient implementation to address the challenge of high GPU memory usage in multi-token prediction models. They achieve this by sequentially computing the forward and backward passes for each head, thereby reducing peak memory requirements.

Experiments on Real Data:

The paper conducts a series of experiments to evaluate the effectiveness of multi-token prediction on both code and natural language tasks.

Benefits scale with model size:

• Larger models show more significant improvements with multi-token prediction compared to smaller models.
• 13B parameter models solve 12% more problems on HumanEval and 17% more on MBPP than comparable next-token models.

Faster inference:

• Using self-speculative decoding, 4-token prediction models achieve a 3x speedup on code and 2.7x speedup on text compared to next-token models.
• 8-byte prediction models achieve a 6.4x speedup, demonstrating the potential for efficient byte-level training.

Learning global patterns with multi-byte prediction:

• Training a byte-level model with 8-byte prediction leads to significant improvements on MBPP and HumanEval compared to next-byte prediction.
• This highlights the ability of multi-token prediction to capture longer-term patterns and dependencies.

Searching for optimal n:

Experiments with different values of n (number of predicted tokens) show that n=4 generally yields the best performance on code benchmarks.

Training with multiple epochs:

The benefits of multi-token prediction remain even when training for multiple epochs, albeit with diminishing improvements.

Finetuning multi-token predictors:

Models pretrained with multi-token prediction also perform better than next-token models when fine-tuned on the CodeContests dataset, demonstrating their richer representations.

Multi-token prediction on natural language:

• While results on multiple-choice NLP benchmarks show no significant improvement, multi-token prediction models achieve better performance on summarization tasks, as measured by ROUGE scores.
• This suggests that multi-token prediction is more beneficial for generative tasks than for discriminative ones.

Ablations on Synthetic Data:

To understand the reasons behind the improvements, the paper conducts experiments on synthetic data, focusing on induction capability and algorithmic reasoning.

Induction capability:

• Multi-token prediction significantly improves the ability of small models to perform induction, i.e., completing patterns based on their most recent continuation.
• This advantage diminishes as models get larger, suggesting that multi-token prediction helps in forming the necessary circuits for induction, which can later be learned with next-token prediction alone.

Algorithmic reasoning:

• Multi-token prediction leads to better generalization on a polynomial arithmetic task, even surpassing the gains achieved by simply increasing model size.
• This further supports the idea that multi-token prediction promotes the development of circuits for complex reasoning tasks.

Why does it work?

The paper provides two main explanations for the effectiveness of multi-token prediction:

Lookahead reinforces choice points:

• Multi-token prediction implicitly assigns higher weights to "choice points" in the text, i.e., tokens that significantly influence the subsequent text generation.
• This allows the model to focus on learning important decision-making steps during training, leading to better generation quality.

Information-theoretic argument:

• Multi-token prediction increases the importance of mutual information between tokens in the loss function.
• This encourages the model to learn relationships between tokens that are relevant for predicting future tokens, improving long-term dependencies and overall text coherence.

Business Implications:

The findings of this paper have several potential business implications for LLMs:

Improved efficiency

‍Multi-token prediction can lead to more efficient LLM training, requiring less data and computational resources to achieve similar or even better performance. This translates to lower costs and faster development cycles.

Enhanced capabilities

‍LLMs trained with multi-token prediction can exhibit better performance on generative tasks such as code generation, summarization, and creative writing. This opens up new possibilities for applications in various industries, including software development, content creation, and education.

Faster response times

‍The ability to use self-speculative decoding for faster inference can significantly improve the responsiveness of LLM-based applications, making them more user-friendly and efficient.

New product opportunities

‍The improved reasoning capabilities facilitated by multi-token prediction can enable the development of LLM-powered tools for tasks that require complex problem-solving and decision-making, such as scientific research, financial analysis, and strategic planning.

Alternative Structures

The paper explores several alternative architectures for multi-token prediction, including:

Replicated unembeddings

‍ Replicating the unembedding matrix for each head, although this approach is not memory-efficient for large models.

Linear heads

‍Using simple linear layers instead of transformer layers for the prediction heads.

Causal and anticausal variants

‍Allowing prediction heads to depend on the outputs of other heads in a causal or anticausal manner.

The experiments show that these alternatives can lead to performance improvements, but not as consistently as the proposed parallel architecture with independent transformer heads.

Conclusion

The paper presents a compelling case for multi-token prediction as a valuable technique for training better and faster LLMs. This approach offers significant benefits in terms of sample efficiency, performance on generative tasks, and inference speed. The findings open up new avenues for research and development in the field of large language models, potentially leading to even more powerful and versatile AI systems in the future.

Critical Analysis

Strengths

• The paper addresses a crucial limitation of current LLM training methods and proposes a simple yet effective solution.
• The experiments are comprehensive and well-designed, providing strong evidence for the benefits of multi-token prediction across various tasks and model sizes.
• The paper provides insightful explanations for the observed improvements, offering valuable theoretical understanding of the technique.

Weaknesses

• The paper mainly focuses on code generation tasks, and further exploration is needed to evaluate the effectiveness of multi-token prediction on a wider range of natural language understanding and generation tasks.
• The optimal number of predicted tokens (n) appears to be task-dependent, and the paper does not provide a clear method for automatically selecting the best value for different scenarios.
• While the proposed architecture is memory-efficient, the training process might still require substantial computational resources, especially for very large models.

Overall, this paper presents a significant contribution to the field of large language models. The proposed multi-token prediction technique offers a promising approach to improve the efficiency and capabilities of LLMs, potentially paving the way for new and exciting applications in the future.