Локальная MAP-выборка для моделей диффузии

Основная информация

• ID статьи: 2510.07343
• Название: Local MAP Sampling for Diffusion Models
• Авторы: Shaorong Zhang (UC Riverside), Rob Brekelmans (Vector Institute), Greg Ver Steeg (UC Riverside)
• Классификация: cs.GR cs.AI eess.IV
• Время публикации/конференция: Препринт (На рецензии)
• Ссылка на статью: https://arxiv.org/abs/2510.07343

Аннотация

Диффузионная апостериорная выборка (DPS) предоставляет принципиальный байесовский подход к обратным задачам путём выборки из $p(x_0 \mid y)$. Однако на практике цель решения обратной задачи состоит не в покрытии апостериорного распределения, а в восстановлении наиболее точного результата реконструкции; методы диффузии на основе оптимизации обычно превосходят в этом отношении, несмотря на отсутствие чёткой вероятностной основы. В данной работе представлена локальная MAP-выборка (LMAPS) — новая схема вывода, которая итеративно решает локальные подзадачи MAP вдоль траектории диффузии. Этот подход проясняет связь с глобальной MAP-оценкой и DPS, обеспечивая единую вероятностную интерпретацию методов на основе оптимизации. На этой основе мы разработали практические алгоритмы с вероятностно интерпретируемыми аппроксимациями ковариации, переформулировкой целевой функции для стабильности и интерпретируемости, а также аппроксимацией градиентов для недифференцируемых операторов.

Исследовательский контекст и мотивация

Постановка проблемы

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

1. Несоответствие целей: DPS направлена на выборку из апостериорного распределения $p(x_0|y)$, но практическая цель решения обратной задачи — получить наиболее точную реконструкцию, а не разнообразную выборку
2. Отсутствие теоретической основы: Методы диффузии на основе оптимизации (такие как Resample, DiffPIR, DCDP) демонстрируют отличную производительность, но лишены чёткой вероятностной теоретической основы

Исследовательская мотивация

• Практическая ориентация: Протоколы оценки обратных задач обычно сравниваются с единственным истинным эталоном и не предоставляют вознаграждения за полноту или разнообразие
• Теоретическое единство: Необходимо обеспечить вероятностную интерпретацию методов оптимизации и прояснить их связь с MAP-оценкой и DPS
• Повышение производительности: Достичь лучшей производительности реконструкции при сохранении теоретической основы

Основные вклады

1. Теоретический вклад: Предложена схема локальной MAP-выборки (LMAPS), которая итеративно решает локальные подзадачи MAP вдоль траектории диффузии, анализируется её связь с глобальной MAP и DPS, обеспечивается единая интерпретация TMPD и методов оптимизации для обратных задач
2. Методологический вклад:
   • Предоставлены вероятностно интерпретируемые аппроксимации ковариации, заменяющие эвристические выборы в существующих решателях
   • Введена переформулировка целевой функции для достижения интерпретируемых параметров и улучшения стабильности
   • Разработана стратегия аппроксимации градиентов для недифференцируемых операторов
3. Экспериментальный вклад: Верификация на 10 задачах восстановления изображений и 3 научных обратных задачах; достижение лучших результатов в 46 из 60 случаев FFHQ/ImageNet; получение улучшения ≥2 дБ PSNR на задачах деблюринга движения, восстановления JPEG и квантования

Подробное описание метода

Определение задачи

Обратная задача направлена на восстановление неизвестного изображения или сигнала $x_0 \in \mathbb{R}^n$ из априорного распределения $\pi(x_0)$ и зашумленного измерения $y \in \mathbb{R}^m$: $y = H(x_0) + z$ где $H(\cdot): \mathbb{R}^n \to \mathbb{R}^m$ — прямой оператор, $z \sim \mathcal{N}(0, \sigma_y^2 I)$ — шум измерения.

Основная теоретическая схема

Глобальная MAP vs локальная MAP

Глобальная MAP непосредственно оптимизирует полное апостериорное распределение: $x_0^{\text{MAP}} := \arg\max_{x_0} p(x_0|y)$

Локальная MAP решает условную задачу оптимизации на каждом временном шаге $t$: $x_0^*(t, x_t, y) := \arg\max p(x_0|x_t, y)$$x_{t-\Delta t} = g(x_0^*, x_t, \epsilon), \quad \epsilon \sim \mathcal{N}(0,I)$

Связь с DPS и DAPS

• DPS: Использует условное математическое ожидание $E[x_0|x_t, y]$
• DAPS: Выборка из $p(x_0|x_t, y)$
• LMAPS: Использует условный режим $\arg\max p(x_0|x_t, y)$

Ключевое понимание: DPS и LMAPS эквивалентны только когда $p(x_0|x_t, y)$ является гауссовым распределением.

Практическое проектирование алгоритма

Аппроксимация ковариации

Применяется изотропная аппроксимация: $\Sigma_{0|t} \approx \frac{k}{\text{SNR}} I, \quad \text{SNR} := \frac{\alpha_t^2}{\sigma_t^2}$

Переформулировка целевой функции

Переформулировка исходной целевой функции: $x_0^* = \arg\min \left\{\frac{\text{SNR}}{k}\|x_0 - m_{0|t}\|^2 + \frac{1}{\sigma_y^2}\|y - H(x_0)\|^2\right\}$

в следующий вид: $x_0^* = \arg\min \left\{(1-\mu_t)\frac{1}{2}\|x_0 - m_{0|t}\|^2 + \mu_t k_2\|y - H(x_0)\|^2\right\}$

где $\mu_t = \frac{\sigma_t^2}{\sigma_t^2 + k_1^2} \in (0,1)$, обеспечивающая:

• Интерпретацию выпуклой комбинации: Веса $(1-\mu_t)$ и $\mu_t$
• Автоматическое отжигание: По мере уменьшения $\sigma_t^2$ переход от управления измерениями к управлению априором
• Численную стабильность: Избегание экстремального масштабирования SNR

Обработка недифференцируемых операторов

Для недифференцируемых задач, таких как восстановление JPEG и квантование, используется суррогатный градиент: $\nabla_{x_0}\|y - H(x_0)\|^2 \approx 2J_{H'}(x_0)^T(H(x_0) - y)$

Для квантования применяется $H'(x_0) = x_0$, упрощаясь до: $\nabla_{x_0}\|y - H(x_0)\|^2 \approx 2(H(x_0) - y)$

Экспериментальная установка

Наборы данных

• Восстановление изображений: FFHQ 256×256 и ImageNet 256×256, по 100 тестовых изображений каждый
• Научные обратные задачи: Использование набора данных InverseBench, включающего изображения флуоресцентной микроскопии (линейное обратное рассеяние), данные GRMHD (визуализация чёрных дыр), данные fastMRI коленного сустава (сжатое зондирование MRI)

Метрики оценки

• Восстановление изображений: PSNR, SSIM, LPIPS
• Научные обратные задачи: Главным образом PSNR, дополненный специфичными для задачи метриками

Методы сравнения

Включают DDNM, DDRM, ΠGDM, DPS, LGD, PnP-DM, FPS, MCG-diff, RedDiff, DAPS, DiffPIR, DCDP, DMPlug и 12 других методов

Детали реализации

• Количество шагов диффузии: обычно 200
• Количество шагов обновления градиента: 20-200 (в зависимости от задачи)
• Скорость обучения: 0.01-1.0 (зависит от задачи)
• Параметры $k_1$: 0-10, $k_2$: 0.01-30000

Результаты экспериментов

Основные результаты

Задачи восстановления изображений

В таблице 1 LMAPS достигает лучшей производительности в 49 из 60 результатов:

• Деблюринг движения: 32.62 дБ на FFHQ против 29.66 дБ DAPS (+2.96 дБ)
• Восстановление JPEG: 27.25 дБ на FFHQ против 25.04 дБ ΠGDM (+2.21 дБ)
• Квантование: 29.51 дБ на FFHQ против 25.82 дБ ΠGDM (+3.69 дБ)

Научные обратные задачи

В таблице 2 LMAPS достигает лучшего PSNR на всех задачах:

• Линейное обратное рассеяние (NR=360): 38.07 дБ против 36.56 дБ RED-diff (+1.51 дБ)
• Линейное обратное рассеяние (NR=180): 37.19 дБ против 35.41 дБ RED-diff (+1.78 дБ)
• Линейное обратное рассеяние (NR=60): 30.75 дБ против 27.07 дБ RED-diff (+3.68 дБ)

Абляционные исследования

Рисунок 4 показывает компромисс между количеством шагов оптимизации и шагов диффузии:

• Оптимальная производительность обычно наблюдается при NFE=200-500
• Увеличение количества шагов оптимизации на каждом шаге диффузии значительно улучшает производительность
• По сравнению с SITCOM (600 NFE), LMAPS достигает аналогичной производительности с меньшими вычислительными ресурсами

Вычислительная эффективность

Таблица 3 показывает время выборки LMAPS на задаче деблюринга:

• LMAPS (200 шагов диффузии, 100 шагов оптимизации): 61 сек/изображение, 30.88 дБ
• DAPS (200 шагов диффузии, 100 шагов оптимизации): 110 сек/изображение, 29.19 дБ
• SITCOM (600 шагов): 73 сек/изображение, 29.93 дБ

Связанные работы

Диффузионная апостериорная выборка

DPS и её варианты решают обратные задачи путём прямой выборки из апостериорного распределения $p(x_0|y)$, включая методы TMPD, DDNM, ΠGDM и другие.

Методы на основе оптимизации

Методы Resample, DiffPIR, DCDP, DMPlug и другие решают обратные задачи посредством чередования денойзирования, оптимизации и переиспользования выборок, демонстрируя отличную производительность, но лишены теоретической основы.

Методы MAP-оценки

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

Заключение и обсуждение

Основные выводы

1. Теоретическое единство: LMAPS обеспечивает единую вероятностную интерпретацию методов диффузии на основе оптимизации
2. Повышение производительности: Достигнуто значительное улучшение PSNR на множестве задач, особенно на сложных нелинейных и недифференцируемых задачах
3. Вычислительная эффективность: Демонстрирует лучшую вычислительную эффективность по сравнению с существующими методами

Ограничения

1. Сходимость: Последовательность локальных MAP не обязательно сходится к глобальной MAP
2. Разнообразие: По сравнению с DPS, LMAPS может производить меньше разнообразия выходных данных
3. Чувствительность гиперпараметров: Требуется настройка параметров $k_1$ и $k_2$ для различных задач

Будущие направления

Статья указывает на то, что ключевая роль глобальной MAP в байесовском выводе была значительно недооценена, и эффективное решение глобальной MAP при диффузионных априорах остаётся открытой проблемой. Хотя MAP может уменьшить разнообразие выходных данных путём концентрации на единственном режиме, он обеспечивает более высокую детерминированность и лучшее согласование с наблюдаемыми данными.

Глубокая оценка

Преимущества

1. Значительный теоретический вклад: Впервые обеспечена чёткая вероятностная интерпретация методов диффузии на основе оптимизации
2. Комплексные эксперименты: Охватывают 10 задач восстановления изображений и 3 научные обратные задачи с полной экспериментальной установкой
3. Явное улучшение производительности: Достигнуто значительное улучшение на 2 дБ и более на множестве сложных задач
4. Практичный метод: Предоставлены эффективные стратегии обработки недифференцируемых операторов
5. Ясное изложение: Теоретический анализ и описание методов хорошо структурированы

Недостатки

1. Глубина теоретического анализа: Хотя обеспечена вероятностная интерпретация, анализ сходимости и теоретических гарантий относительно ограничен
2. Сложность гиперпараметров: Требуется настройка нескольких гиперпараметров для каждой задачи, что может повлиять на обобщаемость метода
3. Вычислительные издержки: Хотя более эффективен, чем некоторые методы, всё ещё требует множественных обновлений градиентов на каждом временном шаге
4. Ограничения оценки: Основное внимание уделяется качеству реконструкции; оценка количественной оценки неопределённости недостаточна

Влияние

1. Академическая ценность: Предоставляет новую теоретическую перспективу для применения моделей диффузии в обратных задачах
2. Практическая ценность: Имеет прямое применение в области восстановления изображений и научных вычислений
3. Вдохновляющий потенциал: Может вдохновить дальнейшие исследования MAP-оценки для моделей диффузии

Применимые сценарии

• Обратные задачи, требующие высокого качества реконструкции, а не разнообразия
• Задачи с недифференцируемыми прямыми операторами (такие как восстановление JPEG, квантование)
• Решение обратных задач в научных вычислениях
• Приложения реального времени с определёнными требованиями к вычислительной эффективности

Библиография

Статья ссылается на важные работы в области моделей диффузии, решения обратных задач и байесовского вывода, включая исходные статьи ключевых методов DPS, DAPS, TMPD и других, обеспечивая хорошую литературную основу для связанных исследований.