TraFL: новый метод пост-тренировки диффузионных языковых моделей превзошёл базовые модели

Исследователи представили новый метод пост-тренировки диффузионных языковых моделей — TraFL (Trajectory Flow baLancing). Работа опубликована на arXiv.

Диффузионные модели становятся альтернативой авторегрессионным, но существующие методы пост-тренировки ориентированы на максимизацию награды. Это приводит к «траекторной блокировке» — концентрации вероятностной массы на узком наборе путей шумоподавления, что снижает разнообразие решений при многократной выборке.

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

Метод оценили на бенчмарках математического рассуждения и генерации кода. TraFL стал единственным методом среди рассмотренных, который улучшил результаты базовой модели на всех длинах ответов, причём улучшения сохраняются с ростом бюджета выборки.

Результаты также перенеслись на отложенные тесты: TraFL остаётся выше базовой модели на Minerva Math и показывает наилучшие показатели на всех уровнях сложности LiveCodeBench.

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