Better & Faster Large Language Models via Multi-token Prediction

TL;DR

Training language models to predict multiple future tokens at once improves sample efficiency and inference speed.

Problem

Traditional next-token prediction only predicts one token at each position, limiting long-range dependency capture and efficient learning, especially for structured outputs like code.

Approach

Place n independent output heads on top of a shared model trunk to predict the next n tokens simultaneously at each position. This is used as an auxiliary training task; during inference, only the first head is used, or multiple heads can be leveraged for parallel decoding to speed up inference.

Results & Contribution

For 13B parameter models, performance improves by 12% on HumanEval and 17% on MBPP. Models trained with 4-token prediction are up to 3 times faster at inference. On small algorithmic tasks, multi-token prediction fosters the development of induction heads and algorithmic reasoning capabilities.