Исследователи (в статье упоминаются «два Ильи и Дима») перенесли MMDi-лосс, ранее показавший себя в дистилляции диффузионных картиночных моделей, на диффузионные языковые модели. Подход заметно улучшает баланс качество/скорость (Парето-фронт) при генерации текста, снижая число шагов инференса без существенной потери качества.

Идея развивает более раннюю работу тех же авторов Scale-wise Distillation of Diffusion Models, где MMD-лосс применялся для ускорения диффузионных моделей изображений. Теперь аналогичная техника адаптирована под специфику текстовых диффузионных моделей — как дискретных, так и непрерывных.

• MMD (Maximum Mean Discrepancy) — мера различия между распределениями P и Q через положительно определенный кернел k (линейный, RBF и др.), считается по формуле с тремя слагаемыми: E[k(x,x')] + E[k(y,y')] − 2E[k(x,y)].
• В кернел подаются признаки замороженной диффузионной модели, снятые с чистых (не зашумлённых) сэмплов; признаки потокенные, а не усреднённые.
• Выбор позиций зависит от постановки задачи: все позиции для безусловной генерации, только позиции ответа для условной, только маскированные токены для маскированной диффузии.
• MMD раскладывается на три члена: real–generated (приближает P к Q), generated–generated (разнообразие внутри P), real–real (признан бесполезным и выбрасывается).
• Для дискретных диффузионных моделей градиент через категориальное распределение не протекает, поэтому используется REINFORCEi с сэмплированием группы MMD-батчей для снижения дисперсии и оптимизацией policy-gradient суррогата.
• Для непрерывных моделей важен self-conditioningi (обусловливание на прошлые выходы модели), backprop делается только через последний шаг генерации; дополнительно предложен iterative refinement distillation (IRD) — студент в один шаг воспроизводит результат многошагового MMD-учителя.
• Метод валидируется на предобученных MDLM для дискретных моделей и на ELF поверх разных энкодеров — для непрерывного случая (MDLM-MMD, по данным источника, достигает улучшения, но итог эксперимента в доступном тексте не раскрыт).