Bessere und schnellere große Sprachmodelle durch Mehrfach-Token-Vorhersage.

Better & Faster Large Language Models via Multi-token Prediction

April 30, 2024

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

cs.AI

Zusammenfassung

Große Sprachmodelle wie GPT und Llama werden mit einem Verlust für die Vorhersage des nächsten Tokens trainiert. In dieser Arbeit schlagen wir vor, dass das Training von Sprachmodellen zur Vorhersage mehrerer zukünftiger Tokens auf einmal zu einer höheren Stichproben-Effizienz führt. Genauer gesagt bitten wir das Modell an jeder Position im Trainingskorpus, die nächsten n Tokens mit n unabhängigen Ausgabeköpfen vorherzusagen, die auf einem gemeinsamen Modell-Trunk arbeiten. Indem wir die Mehrfach-Token-Vorhersage als zusätzliche Trainingsaufgabe betrachten, messen wir verbesserte Fähigkeiten für nachgelagerte Anwendungen ohne zusätzlichen Zeitaufwand für das Training von Code- und natürlichen Sprachmodellen. Die Methode ist besonders nützlich für größere Modellgrößen und behält ihren Reiz bei, wenn über mehrere Epochen trainiert wird. Die Gewinne sind besonders deutlich bei generativen Benchmarks wie dem Codieren, wo unsere Modelle konsequent starke Baselines um mehrere Prozentpunkte übertreffen. Unsere 13B-Parameter-Modelle lösen auf HumanEval 12 % mehr Probleme und auf MBPP 17 % mehr als vergleichbare nächst-Token-Modelle. Experimente an kleinen algorithmischen Aufgaben zeigen, dass die Mehrfach-Token-Vorhersage für die Entwicklung von Induktionsköpfen und algorithmischen Denkfähigkeiten vorteilhaft ist. Als zusätzlicher Vorteil sind Modelle, die mit 4-Token-Vorhersage trainiert wurden, bis zu 3-mal schneller bei der Inferenz, selbst bei großen Batch-Größen.

English

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.