Des modèles de langage plus performants et plus rapides grâce à la prédiction multi-tokens

Better & Faster Large Language Models via Multi-token Prediction

April 30, 2024

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

cs.AI

Résumé

Les grands modèles de langage tels que GPT et Llama sont entraînés avec une fonction de perte de prédiction du prochain token. Dans ce travail, nous suggérons que l'entraînement des modèles de langage à prédire plusieurs tokens futurs simultanément améliore l'efficacité en termes d'échantillons. Plus précisément, à chaque position du corpus d'entraînement, nous demandons au modèle de prédire les n tokens suivants en utilisant n têtes de sortie indépendantes, opérant sur un tronc de modèle partagé. En considérant la prédiction multi-tokens comme une tâche d'entraînement auxiliaire, nous observons une amélioration des capacités en aval sans surcoût en temps d'entraînement, aussi bien pour les modèles de code que pour les modèles de langage naturel. Cette méthode est d'autant plus utile pour les modèles de grande taille et conserve son intérêt lors d'un entraînement sur plusieurs époques. Les gains sont particulièrement marqués sur les benchmarks génératifs comme le codage, où nos modèles surpassent systématiquement les modèles de référence de plusieurs points de pourcentage. Nos modèles de 13 milliards de paramètres résolvent 12 % de problèmes supplémentaires sur HumanEval et 17 % de plus sur MBPP par rapport aux modèles de prédiction du prochain token. Les expériences sur de petites tâches algorithmiques démontrent que la prédiction multi-tokens favorise le développement de têtes d'induction et de capacités de raisonnement algorithmique. Comme avantage supplémentaire, les modèles entraînés avec une prédiction de 4 tokens sont jusqu'à 3 fois plus rapides à l'inférence, même avec de grandes tailles de lots.

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.