Какой должна быть политика размаскирования в диффузионных языковых моделях: контролируемое исследование

Маскированные диффузионные языковые модели генерируют токены параллельно, но их эффективность зависит от того, какие из предсказаний фиксируются на каждом шаге. Выбором фиксируемых токенов занимается политика размаскирования. Обученные политики улучшают баланс между скоростью и точностью, однако они отличаются сразу по нескольким признакам: способу обучения с учителем, входным данным и ёмкости. Из-за этого неясно, какие именно решения дают выигрыш.
Авторы статьи рассматривают обучение с учителем как дистилляцию процедуры декодирования через учителя и исследуют четыре вопроса независимо друг от друга: что учитель помечает как допустимое к фиксации, на каких состояниях политика имитирует учителя, что политика наблюдает на входе и насколько нужно сократить число вычислений функции (NFE), чтобы окупить накладные расходы самой политики.
Находки такие. Пометка предсказаний, совпадающих с итоговым выходом, как допустимых к фиксации оказывается слишком мягкой при малом NFE. Скрытые состояния модели сами по себе слабый вход для политики, но могут дополнять уверенность (confidence). Накладные расходы окупаются сокращением NFE примерно на 10% для политик размером до 7,5 млн параметров в протестированных режимах.
На основе этого авторы обучают политику на целях «сжатия траектории» (trajectory-contraction): они принимают только те фиксации, которые остаются на опорной траектории. Входами служат уверенность и скрытые состояния. На LLaDA-8B-Instruct получившаяся политика в 856 тыс. параметров, по заявлению авторов, даёт лучшую границу «точность, NFE» и наибольшую пропускную способность на GSM8K, конкурентоспособна при умеренных целевых значениях на MATH-500 и HumanEval и в значительной мере переносится на Dream-7B-Instruct.
Ключевые факты
- Работа разделяет факторы обучения политики размаскирования: разметка учителя, состояния для обучения, входы политики и накладные расходы.
- Пометка предсказаний, совпадающих с итоговым выходом, как допустимых к фиксации слишком мягка при малом NFE; скрытые состояния слабы сами по себе, но дополняют уверенность.
- Сокращение NFE примерно на 10% достаточно, чтобы окупить накладные расходы политик до 7,5 млн параметров.
- Политика в 856 тыс. параметров, обученная на целях сжатия траектории, по заявлению авторов, лучшая по границе «точность, NFE» и пропускной способности на GSM8K для LLaDA-8B-Instruct.
- На MATH-500 и HumanEval политика конкурентоспособна при умеренных целях, а на Dream-7B-Instruct в значительной мере переносится.
Почему это важно
Скорость генерации диффузионных языковых моделей упирается в то, сколько токенов можно безопасно зафиксировать за один шаг. Обученные политики размаскирования уже улучшают баланс скорости и точности, но разные работы меняют сразу несколько вещей, и непонятно, что именно помогает. Контролируемое сравнение по отдельным факторам даёт ориентиры для дальнейших разработок. Сама политика при этом очень мала: 856 тыс. параметров на фоне модели в 8 млрд.
Кому это важно
Исследователям и инженерам, которые работают над ускорением вывода диффузионных языковых моделей и проектируют политики выбора токенов для фиксации, в частности для LLaDA и Dream.
Как это применить
Практические выводы статьи такие: не размечать как допустимые к фиксации все предсказания, совпадающие с итоговым выходом; использовать уверенность и дополнять её скрытыми состояниями; учитывать, что накладные расходы политики до 7,5 млн параметров окупаются сокращением NFE примерно на 10%; обучать на целях сжатия траектории. О публикации кода или весов в тексте не сказано.
Можно ли доверять
Это препринт с Hugging Face Papers, оценки в пересказе, заявления самих авторов. Авторы и организации в тексте не названы. Числовые значения точности, пропускной способности и NFE для политики и базовых методов не приведены, как и список методов, с которыми шло сравнение. Поэтому величину выигрыша по описанию оценить нельзя.
Риски и подводные камни
Лучший результат заявлен только для GSM8K на LLaDA-8B-Instruct; на MATH-500 и HumanEval политика названа лишь конкурентоспособной при умеренных целях, а перенос на Dream-7B-Instruct описан словом «в значительной мере» без количественной оценки. Вывод об окупаемости накладных расходов при сокращении NFE примерно на 10% сформулирован для протестированных режимов и политик до 7,5 млн параметров.