🔬 Метод

Первое наблюдение, которое делают авторы, что сила энкодера не всегда соответствует качеству генерации. Например, переход от DINOv2 к DINOv3, DINOv2-B к DINOv2-L ухудшает качество, хотя, казалось бы, внешняя модель должна была выучить лучшие репрезентации.

Потому предлагают выравнивать представления от самой модели на разных уровнях шума. Но как это правильно сделать 🤔?

Пробовали некоторые токены полностью зашумлять или зашумлять все независимо. Но это приводит к несовпадению между обучением и инференсом и только просаживает качество.

Вместо этого решили сэмплировать два времени - t и s - независимо, и часть токенов зашумлять на меньший шум, а часть - на больший.

Уже это как-то работает. Но еще лучше, если добавить лосс выравнивания (косинусную близость) между признаками ученика с входом зашумленным с разным силой и учителем (ЕМА ученика), зашумленным на меньший из шумов. Причем выравнивают более ранние (ближе к входу) признаки ученика с признаками учителя.

Итоговый лосс есть сумма flow-matching лосса и лосса выравнивания.

🧪 Эксперименты

Подход валидируют для 3 модальностей - text-2-image, text-2-video, text-2-audio и для мультимодальной генерации.

Для экспериментов на ImageNet берут SiT с SD автоэнкодером в базовой постановке. Self-Flow работает немного лучше, чем REPA и заметно лучше, чем бейзлайн SRA. Self-Flow помогает даже поверх RAE.

На text-2-image генерации Self-Flow тоже лучше SRA. REPA и SigLIP2 не помогают text-2-image генерации.

На видео 📹 генерации Self-Flow снова заметно улучшает метрики. Внешние автоэнкодеры (V-JEPA, DepthAnything) только просаживают качество.

И аналогичная картина наблюдается для аудио 🔉.

Кроме того, учат модель которая одновременно умеет в 3 модальности и video-action модель.

С увеличением размера моделей разница с REPA только растет (судя по CLIP Score).

В ablation показывают, что self-supervised лосс важен для качества. Также важно маскирование, и сэмплирование шагов на широком интервале значений, чтобы модель с высоких шумов могла видеть признаки с низких. Равномерное зашумление лучше логит-нормального.

💡 Выводы

Выглядит как универсальный и при этом достаточно простой рецепт ускорения претрейна для диффузионных моделей. Интересно, использовались ли данные наработки при обучении семейства Flux2, или додумались уже позднее?