Зацикливание слоёв сократило обучение диффузионных LLM до 3,3 раза

Диффузионные языковые модели с равномерным шумом (UDLM) долго считались дорогим, но интересным ответом на авторегрессию. Новая работа показывает, как заметно снизить эту цену. Если внутри каждого шага денойзинга рекурсивно прогонять несколько ранних слоёв трансформера, looped UDLM достигает validation NLL обычной UDLM, обученной примерно на 0,5 трлн токенов, с затратами до 2,8 раза меньше на OpenWebText и до 3,3 раза меньше на FineWeb-Edu. Цель обучения и сэмплер при этом не менялись.

Почему UDLM дорого учить и зачем это терпеть

Маскированные диффузионные модели (MDLM) работают с чистым контекстом. Часть токенов заменяется маской, остальные остаются корректными, и модель восстанавливает пропуски. Равномерные модели устроены жёстче. Любой токен последовательности может быть заменён случайным токеном обычного словаря, а затем исправлен на следующих шагах. Во время обучения модель должна очищать текст от шума, опираясь на контекст, который сам может быть недостоверным.

Задача получается сложнее, поэтому UDLM требуют больше FLOPs, чем MDLM. Взамен авторы отмечают более сильные результаты на downstream-задачах, связанных с рассуждениями. Возможность пересмотреть уже написанный токен даёт модели шанс исправить собственную ошибку, а не тащить её до конца генерации. Вопрос в цене, и именно её работа и пытается снизить.

Петля вместо новых параметров

Идея selective looping уже показала пользу для MDLM при равном бюджете FLOPs. Несколько ранних слоёв трансформера применяются рекурсивно внутри одного шага денойзинга. Вес слоёв остаётся тем же, но эффективная глубина вычислений растёт. Авторы перенесли подход на UDLM и ничего больше не трогали: ни функцию потерь, ни процедуру сэмплирования.

Основные цифры из работы:

• Обучение: looped UDLM достигает validation NLL базовой UDLM, обученной на ~0,5 трлн токенов, с затратами до 2,8x меньше на OpenWebText.
• На FineWeb-Edu выигрыш выше: до 3,3x меньше FLOPs.
• Генерация: при одинаковом бюджете NFE (числе вычислений сети на сэмпл) generation perplexity ниже до 9,8%, чем у модели без петель.
• Большая часть улучшения достигается при умеренном числе петель, дальнейшее увеличение даёт заметно меньше.
• Распределение числа петель при обучении: смещение в сторону малых значений сокращает число петель на инференсе на 12% относительно равномерной выборки, точность на downstream-задачах не падает.

Отдельно важно, что работает именно сочетание метрик. Validation NLL показывает качество обучения, а generation perplexity при фиксированном NFE показывает качество самой генерации. Выигрыш виден по обеим, а проверка на downstream-задачах подтверждает, что он не ограничивается удобными для модели метриками.

Сколько петель класть и кому это пригодится

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

Второй практический момент связан с распределением числа петель. Равномерная выборка во время обучения кажется естественным вариантом, но более частые малые значения делают модель устойчивой к урезанному режиму работы. В итоге на инференсе нужно на 12% меньше петель при той же точности. Для команд, которые строят диффузионные LLM, это готовый рецепт настройки. Он ничего не требует менять в сэмплере, значит, его можно добавить в существующий пайплайн.

Главный интерес работы в другом. UDLM были сильнее по рассуждениям, но проигрывали по стоимости обучения. Если зацикливание стабильно срезает затраты в 2,8–3,3 раза, разрыв с MDLM по цене становится меньше, а аргумент в пользу равномерного шума начинает звучать убедительнее. Вопрос масштабирования за пределы OpenWebText и FineWeb-Edu остаётся открытым. Но направление понятное: качество диффузионных языковых моделей можно покупать архитектурой, а не только новыми токенами и GPU-часами.