Modelos de Lenguaje Más Eficientes y Rápidos mediante Predicción Multi-token

Better & Faster Large Language Models via Multi-token Prediction

April 30, 2024

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

cs.AI

Resumen

Los grandes modelos de lenguaje como GPT y Llama se entrenan con una pérdida de predicción del siguiente token. En este trabajo, sugerimos que entrenar modelos de lenguaje para predecir múltiples tokens futuros simultáneamente resulta en una mayor eficiencia muestral. Más específicamente, en cada posición del corpus de entrenamiento, solicitamos al modelo que prediga los siguientes n tokens utilizando n cabezales de salida independientes, operando sobre un tronco de modelo compartido. Al considerar la predicción de múltiples tokens como una tarea de entrenamiento auxiliar, medimos capacidades mejoradas en tareas posteriores sin sobrecarga en el tiempo de entrenamiento, tanto para modelos de código como de lenguaje natural. El método es cada vez más útil para tamaños de modelo más grandes y mantiene su atractivo al entrenar durante múltiples épocas. Las mejoras son especialmente notables en benchmarks generativos como la codificación, donde nuestros modelos superan consistentemente líneas base sólidas por varios puntos porcentuales. Nuestros modelos de 13B parámetros resuelven un 12 % más de problemas en HumanEval y un 17 % más en MBPP en comparación con modelos de predicción de un solo token. Experimentos en tareas algorítmicas pequeñas demuestran que la predicción de múltiples tokens favorece el desarrollo de cabezales de inducción y capacidades de razonamiento algorítmico. Como beneficio adicional, los modelos entrenados con predicción de 4 tokens son hasta 3 veces más rápidos en inferencia, incluso con tamaños de lote grandes.

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.