Карты потока: обучение интегралу диффузионной модели для ускоренной выборки

Карты потоков — это новый подход к ускорению семплирования диффузионных моделей путем обучения нейронных сетей предсказывать весь путь от шума к данным за один шаг, а не итеративно интегрировать направления касательных. В блоге Сандера Дилемана разбирается теория и таксономия.

Как работает семплирование диффузионных моделей сегодня

Стандартные диффузионные модели предсказывают направление касательной (оценку очищенного сигнала) на каждом уровне шума и делают маленькие шаги для интегрирования пути от шума к данным. Это требует много итераций (например, 50–1000 шагов в DDPM). Детерминированные семплеры, такие как DDIM, рассматривают процесс как ОДУ, но по-прежнему требуют численного интегрирования.

Что делают карты потоков по-другому

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

Подходы к обучению

Дилеман классифицирует обучение карт потоков по нескольким стратегиям, основанным на таксономии Боффи и др.:

• Модели согласованности: Обеспечивают, что выход сети одинаков для всех уровней шума вдоль траектории (самосогласованность).
• Прямая регрессия: Обучают сеть напрямую предсказывать чистые данные по зашумленному образцу на любом уровне шума, минимизируя L2-потери относительно истинных данных.
• Дистилляция: Используют предобученную диффузионную модель как учителя для создания пар (зашумленный, чистый) и обучают студента-карту потоков имитировать интеграл ОДУ учителя.

Выбор зависит от того, хотите ли вы обратное распространение через решатель ОДУ (дорого, но точно) или избежать его (дешевле, но потенциально менее стабильно).

Практические соображения

Дилеман отмечает, что карты потоков можно обучать с нуля (без учителя), используя известное отображение шума в данные в случае гауссовского диффузионного процесса. На практике карты потоков предлагают:

• Более быструю семплировку: 1–4 шага (против 10–50 для дистиллированной диффузии).
• Обучение на основе вознаграждения: Интегральная формулировка позволяет прямую градиентную оптимизацию целевых вознаграждений (например, метрик качества изображения).
• Управляемость: Условные ограничения на латентные переменные становятся проще, поскольку весь путь генерации представляет собой единую функцию.

Применения и расширения

Карты потоков применяются в генерации изображений (например, модели согласованности LCM, SDXL-Turbo), видео и 3D-данных. В статье упоминается недавняя монография Лая и др. для более глубокого математического обзора.