Your LLM Knows the Future: Uncovering Its Multi-Token Prediction Potential

TL;DR

This paper introduces a novel framework enabling autoregressive language models to predict multiple tokens simultaneously, significantly increasing inference speed without sacrificing output quality.

Contribution

It presents a new multi-token prediction method using masked-input formulation, gated LoRA, and auxiliary losses, enhancing speed and coherence in autoregressive models.

Findings

Achieves nearly 5x faster code and math generation

Improves chat and knowledge task speed by 2.5x

Maintains high output quality with speedup

Abstract

Autoregressive language models are constrained by their inherently sequential nature, generating one token at a time. This paradigm limits inference speed and parallelism, especially during later stages of generation when the direction and semantics of text are relatively certain. In this work, we propose a novel framework that leverages the inherent knowledge of vanilla autoregressive language models about future tokens, combining techniques to realize this potential and enable simultaneous prediction of multiple subsequent tokens. Our approach introduces several key innovations: (1) a masked-input formulation where multiple future tokens are jointly predicted from a common prefix; (2) a gated LoRA formulation that preserves the original LLM's functionality, while equipping it for multi-token prediction; (3) a lightweight, learnable sampler module that generates coherent sequences from the predicted future tokens; (4) a set of auxiliary training losses, including a consistency loss, to enhance the coherence and accuracy of jointly generated tokens; and (5) a speculative generation strategy that expands tokens quadratically in the future while maintaining high fidelity. Our method achieves significant speedups through supervised fine-tuning on pretrained models. For example, it generates code and math nearly 5x faster, and improves general chat and knowledge tasks by almost 2.5x. These gains come without any loss in quality.