Denoising Diffusion Probabilistic Models

Разбор статьи

Проблема и мотивация

Хотя диффузионные вероятностные модели изначально были предложены Sohl-Dickstein et al. (2015) на основе неравновесной термодинамики, их ранние итерации существенно уступали GAN и авторегрессионным архитектурам по визуальной четкости. Они воспринимались как медленные, сложные в настройке и непроверенные на сложных естественных датасетах высокого разрешения. Авторы решили переосмыслить диффузионные модели, установив строгую связь между их вариационными целевыми функциями и score matching с динамикой Ланжевена, показав, что при правильной параметризации и функции потерь диффузионные модели способны достигать наилучшего визуального качества сэмплов.

Ключевые идеи

• Параметризация через предсказание шума: Вместо прямого предсказания апостериорного среднего ${\tilde{\mathbf{\mu }}}_t$, нейросеть обратного процесса ${\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t)$ обучается предсказывать немасштабированный стандартный вектор гауссовского шума $\mathbf{ϵ}$, добавленный на прямом шаге.
• Связь с Denoising Score Matching: Denoising diffusion эквивалентен denoising score matching на нескольких масштабах шума, а обучаемая обратная цепь действует как сэмплер отожженной динамики Ланжевена, полученный через вариационный вывод.
• Упрощенная функция потерь ($L_{\mathrm{simple}}$): Отказ от строгого, зависящего от дисперсии весового коэффициента в вариационной нижней оценке дает невзвешенную функцию потерь MSE, которая снижает вес низких уровней шума ($t\approx 1$), заставляя модель фокусироваться на перцептивно значимых грубых структурах.
• Взгляд с точки зрения прогрессивного сжатия с потерями: Анализ показывает, что более половины длины кода модели без потерь приходится на неразличимый глазом высокочастотный пиксельный шум. Это означает, что диффузионные модели неявно оптимизируют функцию «скорость-искажение», сначала декодируя глобальное содержание и лишь затем добавляя мелкую текстуру.

Метод

Прямой процесс последовательно искажает данные ${\mathbf{x}}_0$ на протяжении $T=1000$ шагов с фиксированным расписанием дисперсий ${\beta }_1,\dots ,{\beta }_T$: $$q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})=\mathcal{N}({\mathbf{x}}_t;\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1},{\beta }_t\mathbf{I})$$ Используя ${\alpha }_t=1-{\beta }_t$ и ${\bar{\alpha }}_t={\prod }_{s=1}^t{\alpha }_s$, промежуточные латентные переменные можно сэмплировать напрямую в замкнутой форме: $$q({\mathbf{x}}_t|{\mathbf{x}}_0)=\mathcal{N}({\mathbf{x}}_t;\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0,(1-{\bar{\alpha }}_t)\mathbf{I})$$

Обратная цепь $p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)=\mathcal{N}({\mathbf{x}}_{t-1};{\mathbf{\mu }}_{\theta }({\mathbf{x}}_t,t),{\sigma }_t^2\mathbf{I})$ параметризует среднее как: $${\mathbf{\mu }}_{\theta }({\mathbf{x}}_t,t)=\frac{1}{\sqrt{{\alpha }_t}}\left({\mathbf{x}}_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t)\right)$$ где ${\sigma }_t^2$ зафиксирована либо на ${\beta }_t$, либо на ${\tilde{\beta }}_t=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t$.

Сеть оптимизируется с использованием упрощенной суррогатной функции потерь: $$L_{\mathrm{simple}}(\theta )={\mathbb{E}}_{t,{\mathbf{x}}_0,\mathbf{ϵ}}\left[{\Vert \mathbf{ϵ}-{\mathbf{ϵ}}_{\theta }(\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}\mathbf{ϵ},t)\Vert }^2\right]$$ где $t\sim \text{Uniform}(\{1,\dots ,T\})$ и $\mathbf{ϵ}\sim \mathcal{N}(0,\mathbf{I})$.

Генерация (сэмплирование) выполняется в обратном направлении, начиная с ${\mathbf{x}}_T\sim \mathcal{N}(0,\mathbf{I})$, по формуле: $${\mathbf{x}}_{t-1}=\frac{1}{\sqrt{{\alpha }_t}}\left({\mathbf{x}}_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t)\right)+{\sigma }_t\mathbf{z}$$ где $\mathbf{z}\sim \mathcal{N}(0,\mathbf{I})$ при $t>1$, и $\mathbf{z}=0$ при $t=1$.

Результаты

• CIFAR-10: Достигнуты показатели безусловного Inception Score (IS) $9.46\pm 0.11$ и FID $3.17$ с использованием $L_{\mathrm{simple}}$, что превзошло аналоги того времени, такие как SNGAN (FID 21.7) и NCSN (FID 25.32).
• Генерация в высоком разрешении: Успешно сгенерированы реалистичные изображения $256\times 256$ на датасетах LSUN (Bedroom FID 4.90, Church FID 7.89, Cat FID 19.75) и CelebA-HQ.
• Компромисс между правдоподобием и качеством: Хотя обучение по точной вариационной границе $L$ обеспечило лучшее отрицательное логарифмическое правдоподобие ($\le 3.70$ бит/изм. на тесте), чем $L_{\mathrm{simple}}$ ($\le 3.75$ бит/изм. на тесте), функция потерь $L_{\mathrm{simple}}$ показала значительно превосходящий результат по FID (3.17 против 13.51).

Ограничения и открытые вопросы

• Скорость сэмплирования: Генерация требует последовательного применения нейронной сети $T=1000$ раз для каждого батча изображений, что делает инференс на порядки медленнее по сравнению с GAN или однопроходными моделями потоков нормализации (flow models).
• Логарифмическое правдоподобие: Модели не достигают чистых значений тестового правдоподобия без потерь, характерных для чисто авторегрессионных моделей, таких как Sparse Transformers (достигших $2.80$ бит/изм. на CIFAR-10).
• Фиксированное расписание дисперсий: Расписание прямой дисперсии ${\beta }_t$ задавалось фиксированным линейным образом, а не обучалось и не масштабировалось адаптивно; вопрос о том, могут ли нелинейные расписания или непрерывная обучаемая динамика дать дополнительное улучшение, остался открытым.