🔬 Метод

За основу берётся фреймворк flow matching со стохастическим интерполянтом на непрерывных данных.

Однако вместо стандартного v-prediction предлагают перейти к x-prediction, мотивируя это аналогично работе JiT из картиночного мира: мол, оно лежит на многообразии, а скорость — нет.

Потом они получают лемму о том, что оптимальный денойзер есть потокенное апостериорное распределение, откуда следует, что оптимальный денойзер есть минимум кросс-энтропии (т.е. получают совпадение оптимизации диффузионного лосса и правдоподобия).

Далее в статье вводят формулировку Flow Map — отображения, которое переводит из заданной точки в любую другую. Целевое отображение параметризуют в виде two-time денойзера, обладающего следующими свойствами:

🔹 Выход денойзера лежит на вероятностном симплексе.
🔹 Если оба времени совпадают, получаем обычный денойзер для x0-prediction.
🔹 Полугрупповое свойство: попасть в точку t из s — то же самое, что из s в u посередине, а затем из u в t.

Авторы показывают, что подобная формулировка для дискретной диффузии нереализуема.

⚠️ Ещё из важных аспектов следует отметить следующее

🔸 Равномерные шаги сэмплирования работают плохо. Авторы замеряют decoding error rate (долю неправильно предсказанных токенов) в зависимости от времени и замечают, что она держится в районе 1 (всегда ошибаемся) большую часть траектории, а в конце резко падает в 0. Поэтому шаги подбираются так, чтобы равномерно улучшать данную метрику. То есть шаги сконцентрированы в районе переходной зоны.

🔸 В отличие от дискретной диффузии, на непрерывную легко переносятся известные техники гайденса. Автогайденс с плохой моделью или гайденс на некоторый реворд.

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

Обучают модель на 179M параметров на LM1B и OpenWebText. Сравнивают с RDLM, CANDI, MDLM, Duo.

На большом числе шагов сэмплирования FLM (без дистилляции) достигает более низкой перплексии, чем конкурентные подходы. Но при малом числе шагов уступает Duo. Энтропия примерно у всех подходов на одном уровне.

А если сравнивать дистиллированную версию FMLM, то на малом числе шагов (1–8) она оказывается заметно лучше конкурентов.

Затем сравнивают guidance. MDLM и Duo, начиная с какого-то момента (силы guidance), разваливаются, и перплексия уходит в стратосферу. У FLM перплексия только убывает с ростом гайденса, но энтропия тоже убывает.

🎯 Гайденс на реворд улучшает целевую метрику.

Затем проводят ablation метода:

🔹 x0-prediction гораздо лучше, чем v-prediction.
🔹 CE loss с softmax на логитах даёт лучший результат.
🔹 Диффузия в евклидовом пространстве лучше различных римановых и симплексных вариантов.

Для FMLM:

🔹 two-time параметризация оптимальна.
🔹 Лучше стартовать с предобученной диффузии, а не с нуля.
🔹 Полугрупповой дистилляционный objective лучше, чем эйлеров или лагранжев (что есть что прекрасно разобрано у Sander Dieleman).

📝 Выводы

Небезынтересный подход с неплохой теоретической подоплёкой. Но, как и везде, встаёт вопрос масштабирования и обобщаемости.