Modelli Linguistici di Grandi Dimensioni Migliori e Più Veloci tramite Predizione Multi-token

Better & Faster Large Language Models via Multi-token Prediction

April 30, 2024

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

cs.AI

Abstract

I grandi modelli linguistici come GPT e Llama vengono addestrati con una funzione di perdita basata sulla previsione del token successivo. In questo lavoro, suggeriamo che addestrare i modelli linguistici a prevedere più token futuri contemporaneamente porti a una maggiore efficienza campionaria. Più specificamente, in ogni posizione del corpus di addestramento, chiediamo al modello di prevedere i successivi n token utilizzando n teste di output indipendenti, che operano su un tronco condiviso del modello. Considerando la previsione multi-token come un'attività di addestramento ausiliaria, misuriamo miglioramenti nelle capacità downstream senza sovraccarico nel tempo di addestramento, sia per i modelli di codice che per quelli di linguaggio naturale. Il metodo è sempre più utile per dimensioni maggiori del modello e mantiene il suo appeal quando si addestra per più epoche. I guadagni sono particolarmente evidenti su benchmark generativi come la codifica, dove i nostri modelli superano costantemente baseline robuste di diversi punti percentuali. I nostri modelli da 13B parametri risolvono il 12% in più di problemi su HumanEval e il 17% in più su MBPP rispetto a modelli comparabili basati sulla previsione del token successivo. Esperimenti su piccoli compiti algoritmici dimostrano che la previsione multi-token favorisce lo sviluppo di teste di induzione e capacità di ragionamento algoritmico. Come ulteriore vantaggio, i modelli addestrati con la previsione di 4 token sono fino a 3 volte più veloci nell'inferenza, anche con dimensioni di batch elevate.

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.