Better & Faster Large Language Models via Multi-token Prediction
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.
6631a41ae97d286a1d1b3cbd
Fabian Gloeckle
extracted_confirmed
2024-05-01T09:23:37.597Z
64eca847971abfd26fe965be
/avatars/c53e119c6319768ffedc10442d4830ab.svg
Fabian Gloeckle
user
faaabian
6631a41ae97d286a1d1b3cbe
Badr Youbi Idrissi
admin_assigned
2024-05-02T08:09:57.392Z
632df0e4860318fab98f8804
/avatars/94b5810ec4a167d9e2b523c0626a23b4.svg
Badr Youbi Idrissi
user
cetosignis
6631a41ae97d286a1d1b3cbf
Baptiste Rozière
admin_assigned
2024-05-02T08:10:06.134Z
62d063dac375d0c84255b9a1
/avatars/de0fc34bad8c761210c0895ebfa4feba.svg
Baptiste Roziere
user
broz
6631a41ae97d286a1d1b3cc0
David Lopez-Paz
6631a41ae97d286a1d1b3cc1
Gabriel Synnaeve
admin_assigned
2024-05-02T08:10:19.482Z
630eac7931970d1cd4fbacf2
/avatars/b7ccbddfa745db854dc342be1327cd53.svg
Gabriel Synnaeve
user
gsynnaeve
Better & Faster Large Language Models via Multi-token Prediction
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.
Training large language models to predict multiple future tokens simultaneously enhances their sample efficiency and performance, particularly on generative benchmarks and small algorithmic tasks, with decreased inference time.
next-token prediction loss multi-token prediction output heads shared model trunk auxiliary training task HumanEval MBPP induction heads algorithmic reasoning capabilities