Лучшие и быстрее большие языковые модели через множественное предсказание токенов.

Better & Faster Large Language Models via Multi-token Prediction

April 30, 2024

Авторы: Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Rozière, David Lopez-Paz, Gabriel Synnaeve

cs.AI

Аннотация

Большие языковые модели, такие как GPT и Llama, обучаются с использованием потерь предсказания следующего токена. В данной работе мы предлагаем, что обучение языковых моделей на предсказание нескольких будущих токенов одновременно приводит к более высокой эффективности выборки. Более конкретно, на каждой позиции в обучающем корпусе мы просим модель предсказать следующие n токенов, используя n независимых выходных головок, работающих поверх общей модели. Рассматривая многотоковое предсказание как вспомогательную задачу обучения, мы измеряем улучшенные возможности на следующем уровне без дополнительных затрат времени на обучение как для кодовых, так и для естественных языковых моделей. Этот метод становится все более полезным для более крупных моделей и сохраняет свою привлекательность при обучении на протяжении нескольких эпох. Приросты особенно заметны на генеративных бенчмарках, таких как кодирование, где наши модели последовательно превосходят сильные базовые значения на несколько процентных пунктов. Наши модели с 13 миллиардами параметров решают на 12 % больше задач в HumanEval и на 17 % больше в MBPP, чем сравнимые модели с предсказанием следующего токена. Эксперименты на небольших алгоритмических задачах показывают, что многотоковое предсказание благоприятно для развития индуктивных голов и алгоритмических способностей рассуждения. Как дополнительное преимущество, модели, обученные с предсказанием 4 токенов, работают в 3 раза быстрее при выводе, даже с большими размерами пакетов.

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.