Исследователи убрали шум из MTP-обучения LLM одной строкой кода

Почти все свежие крупные языковые модели при предобучении решают задачу чуть сложнее обычной. Они предсказывают не только следующий токен, но и несколько токенов вперёд. Приём называется multi-token prediction, MTP. Он нужен ради двух вещей: модель лучше справляется с задачами на дальний горизонт вроде математики и кода, а дополнительные «головы» можно использовать как быстрых черновщиков при спекулятивном декодировании. Черновик из нескольких токенов набрасывается дёшево, а основная модель проверяет его за один проход. Исследователи из проекта Soft-MTP выяснили, что стандартное обучение таких голов несёт в себе шумную примесь, и предложили способ её убрать. Для этого нужна одна дополнительная строка кода.

Что не так с обычным MTP-обучением

Головы MTP учат стандартной кросс-энтропией против реальных данных. Авторы показали, что градиент этой функции потерь раскладывается на две части, причём точно, без приближений.

• Первая часть — KL-самодистилляция. KL-дивергенция измеряет, насколько одно распределение вероятностей отличается от другого. Здесь она тянет каждую MTP-голову к собственному прогнозу модели для того же токена. Именно эта часть повышает долю принятых черновиков при спекулятивном декодировании: токен от головы принимается, только если он совпадает с прогнозом основной модели. • Вторая часть — остаток. Это градиент обычной кросс-энтропии следующего токена против данных, пришедший по другому пути. У него низкое отношение сигнала к шуму, и он заглушает полезный KL-сигнал.

Почему остаток шумный? Целевой токен в данных — одна конкретная выборка из реального распределения, «горячий» вектор с единицей в одной позиции. Дисперсия такой выборки не исчезает ни с ростом модели, ни с длительным обучением. Авторы измерили это напрямую. У остатка отношение сигнала к шуму заметно ниже, чем у KL-компоненты, а разрыв растёт с размером модели. Кроме того, остаток направлен против обновления, которое получает основная задача предсказания следующего токена. Выходит, он вредит и обучению MTP-голов, и обучению самой модели.

Одна строка вместо матриц размером со словарь

Чистый KL-градиент можно получить и в лоб: построить распределения учителя и ученика и минимизировать расхождение. Но при словаре в десятки тысяч токенов придётся держать в памяти огромные матрицы логитов. К тому же такой вариант не работает со стандартными «слитыми» ядрами кросс-энтропии (fused kernels), которые как раз придуманы, чтобы эти матрицы не материализовать. Поэтому передовые модели на практике учат MTP-головы обычной кросс-энтропией.

Авторы предложили приём KL gradient-recovery trick. Если у двух распределений общая линейная выходная голова, градиент их KL можно восстановить точно, скомбинировав два уже посчитанных кросс-энтропийных градиента. Технически это straight-through-маршрутизация градиента: к скрытому состоянию основной головы добавляется выражение hk.detach() - hk. Значение при этом не меняется, поэтому обучение основной задачи идёт как раньше. Зато в MTP-представление затекает дополнительный градиент, который в сумме с обычным градиентом MTP-потерь вычитает шумный остаток. В итоге в MTP-головы и общий ствол модели приходит только KL-сигнал.

Так устроен метод Soft-MTP. Новых параметров он не добавляет и заметных накладных расходов не создаёт. Он работает и с MTP в стиле DeepSeek, и с вариантом Meta, а также совместим с параллелизмом, который используется при предобучении крупных моделей. Код опубликован на GitHub: JaydenTeoh/Soft-MTP.

Что показали эксперименты от 167M до 1,2B

Авторы обучили модели размером от 167 млн до 1,2 млрд параметров и сравнили Soft-MTP со стандартным MTP.

• Перплексия при предсказании следующего токена ниже на всех размерах, и разрыв растёт с масштабом. • Перплексия MTP-голов ниже, начиная с модели на 689 млн параметров. • Зазор по точности на задачах без дообучения (zero-shot) в пользу Soft-MTP сохраняется на каждом масштабе. • Самоспекулятивное декодирование ускоряется во всех проверенных доменах. • После дообучения модель на 1,2 млрд параметров лучше решает задачи на математику и код.

Отдельно сравнили метод с прямой KL-минимизацией, которая требует явного расчёта логитов. На большинстве метрик Soft-MTP не уступает ей, а местами выигрывает. Вероятная причина в том, что Soft-MTP по-прежнему обучает общую выходную голову обычной кросс-энтропией по данным. При чистом KL она смещалась бы к собственным прогнозам модели.

Цена и ограничения

У метода есть честный компромисс. Убрав шумный остаток, Soft-MTP вносит смещение: MTP-головы тянутся к распределению самой модели, а не к реальному распределению данных. Авторы считают, что с ростом масштаба обмен становится всё выгоднее: ошибка учителя-модели уменьшается, а шум выборки остаётся. К тому же MTP-головы нужны в первую очередь для ускорения инференса, где совпадение с прогнозом основной модели как раз и нужно.

Для команд, которые предобучают собственные LLM с MTP, практический вывод простой. Правка крошечная, ничего не стоит по памяти и скорости и не ломает привычные fused-ядра. Эксперименты ограничены моделями до 1,2 млрд параметров, так что на масштабе DeepSeek-V3 результат ещё предстоит проверить. Но тенденция, при которой выигрыш растёт вместе с моделью, выглядит обнадёживающе.