Modelos de Linguagem Maiores e Mais Rápidos via Predição Multi-token

Resumo

Modelos de linguagem de grande escala, como GPT e Llama, são treinados com uma função de perda de previsão do próximo token. Neste trabalho, sugerimos que treinar modelos de linguagem para prever múltiplos tokens futuros simultaneamente resulta em maior eficiência amostral. Mais especificamente, em cada posição do corpus de treinamento, solicitamos que o modelo preveja os próximos n tokens usando n cabeças de saída independentes, operando sobre um tronco de modelo compartilhado. Ao considerar a previsão de múltiplos tokens como uma tarefa auxiliar de treinamento, medimos melhorias nas capacidades subsequentes sem sobrecarga no tempo de treinamento, tanto para modelos de código quanto de linguagem natural. O método é cada vez mais útil para tamanhos maiores de modelos e mantém seu apelo ao treinar por múltiplas épocas. Os ganhos são especialmente evidentes em benchmarks generativos, como codificação, onde nossos modelos consistentemente superam baselines fortes por vários pontos percentuais. Nossos modelos com 13 bilhões de parâmetros resolvem 12% mais problemas no HumanEval e 17% mais no MBPP em comparação com modelos de previsão de próximo token equivalentes. Experimentos em pequenas tarefas algorítmicas demonstram que a previsão de múltiplos tokens é favorável para o desenvolvimento de cabeças de indução e capacidades de raciocínio algorítmico. Como um benefício adicional, modelos treinados com previsão de 4 tokens são até 3 vezes mais rápidos na inferência, mesmo com grandes tamanhos de lote.

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.

Modelos de Linguagem Maiores e Mais Rápidos via Predição Multi-token

Better & Faster Large Language Models via Multi-token Prediction

Resumo

Support