Крошечная политика на 856K параметров ускорила диффузионную LLaDA-8B

Диффузионные языковые модели обещают то, чего не умеют классические авторегрессионные: генерировать много токенов за один проход. На практике всё упирается в вопрос, какие из предсказанных токенов можно безопасно зафиксировать прямо сейчас, а какие лучше оставить под маской до следующего шага. Новое контролируемое исследование разложило этот вопрос на составные части и в итоге обучило политику размером всего 856 тысяч параметров. На LLaDA-8B-Instruct она дала лучшую границу «точность — число вычислений модели (NFE)» и самую высокую пропускную способность на GSM8K.

Политика как дистилляция декодирования

Обучаемые политики unmasking уже появлялись, и они действительно улучшают компромисс между скоростью и качеством. Но работы различались сразу по нескольким осям: чем размечена обучающая выборка, какие входы видит политика, насколько она велика. По итогам сравнивать подходы между собой было почти невозможно, и непонятно, какое из решений вообще играет роль.

Авторы предложили единую рамку. Обучаемая политика с учителем здесь дистиллирует некую процедуру декодирования через «учителя», то есть разметку. Дальше они изучали четыре вещи по отдельности:

• что учитель считает токеном, который можно зафиксировать;
• на каких состояниях генерации политика имитирует учителя;
• что политика наблюдает на входе;
• насколько нужно снизить NFE, чтобы окупить собственные накладные расходы.

Где ломается наивная разметка

Самый очевидный способ разметки выглядит так: токен считается готовым к фиксации, если он совпадает с тем, что окажется в финальном ответе. Исследование показало, что при низком NFE такой критерий слишком мягкий. Политика учится принимать токены, которые случайно совпали с итогом, но на самом деле тянут генерацию в сторону от эталонного пути. Когда вычислений мало, цена таких ошибок растёт.

По входам картина тоже неоднозначная. Скрытые состояния модели сами по себе дают слабый сигнал для решения о фиксации. Зато как дополнение к уверенности (confidence) предсказания они работают: вместе два источника информации дают больше, чем каждый по отдельности.

Отдельно авторы посчитали цену самой политики. Любая надстройка над моделью добавляет вычисления, и выгода от меньшего числа шагов может уйти на её обслуживание. По их оценке, сокращения NFE примерно на 10% достаточно, чтобы компенсировать накладные расходы политик до 7,5 млн параметров во всех проверенных конфигурациях.

Траектория вместо совпадения с финалом

Из этих наблюдений выросла итоговая схема. Политику обучают на целях trajectory-contraction: фиксация принимается только тогда, когда она остаётся на опорной траектории декодирования. На входе у политики уверенность предсказаний и скрытые состояния.

Ключевые результаты:

• размер политики: 856K параметров, меньше 0,02% от 8 млрд параметров LLaDA-8B-Instruct;
• GSM8K: лучшая граница точность — NFE и наибольшая пропускная способность среди сравниваемых подходов;
• MATH-500 и HumanEval: конкурентные результаты при умеренных целевых уровнях ускорения;
• Dream-7B-Instruct: политика в основном переносится на другую модель без полной переделки.

Оговорка из самих результатов: на MATH-500 и HumanEval преимущество заявлено только на умеренных целях, то есть при агрессивном сокращении шагов лидерство не гарантировано.

Что из этого забрать в практику

Для команд, которые гоняют диффузионные LLM в продакшене, главный вывод инженерный. Ускорение декодирования не обязательно требует тяжёлой надстройки: политика меньше миллиона параметров окупается, если снижает NFE хотя бы на десяток процентов. Это удобный ориентир для бюджета: если выигрыш по шагам меньше, накладные расходы политики съедят его целиком.

Вторая рекомендация касается разметки. Если вы дистиллируете стратегию декодирования, не стоит считать «совпало с финальным ответом» достаточным критерием, особенно в режиме малого числа вычислений. Привязка к опорной траектории даёт более строгие обучающие сигналы. А скрытые состояния стоит добавлять к уверенности, а не использовать как единственный вход.

Наконец, перенос на Dream-7B-Instruct намекает, что такие политики можно делать относительно универсальными внутри семейства масочных диффузионных моделей. Для индустрии, где диффузионные языковые модели борются за место рядом с авторегрессионными в задачах с жёсткими требованиями к задержке, это аргумент в пользу дешёвых и переносимых надстроек вместо переобучения самих моделей.