🔬🧪 Метод и эксперименты

Известно, что естественные изображения заметают не все возможные комбинации пикселей, а лежат на некотором подпространстве.

При обучении диффузионных моделей используются обыкновенно 3 типа предсказания:

📌Предсказание расшумленного сэмпла - x0
📌Предсказание шума epsilon
📌Предсказание скорости - v (взвешенная комбинация x0 и eps)

Математически вроде бы все формулировки эквивалентны (с точностью до изменения коэффициентов перед лоссом). Однако на подпространстве лежит лишь x0, в то время как шум и, соответственно, скорость заметают все пространство.

Отсюда авторы делают предположение, что учиться на x0 должно быть проще, чем на \epsilon или скорости.

Для валидации гипотезы сэмплируют данные в виде 2-мерной спирали, вложенной в пространство большой размерности. И оказывается, что начиная с какого-то момента v-prediction и eps-prediction работает плохо, а x0 - хорошо.

Затем гипотезу проверяют уже на большем масштабе. Берут Vision Transformer, который называют просто Just Image Transformer (JiT, не путать с JiT-компиляцией), нарезают картинку на большие патчи (16x16, 32x32), и гоняют в таком пространстве диффузию на ImageNet-256/512.

Все варианты, кроме x0-prediction работают из ряда вон плохо, и тюнинг уровня шума не помогает. x0-prediction же работает сносно. Оказывается, что добавление дополнительного боттлнека после патчевалки даже немного улучшает качество.

Далее в модель накидывают ряд архитектурных модификаций, отходя от классического DiT - SwiGLU, RMSNorm-ы, RoPE и 32 in-context класс токена, что немного улучшает метрики.

Итоговый результат не SOTA 🏆, но вполне пристойный.

💡 Выводы

Интерпретация через manifold learning выглядит интересно. Однако вход в модель-то все равно (зашумленный сэмпл) остается чем-то лежащим в пространстве высокой размерности. Интересно, перенесутся ли полученные результаты на class-conditional случай.