Несоответствия в Моделях Консистенции: Лучшее Решение ODE Не Обязательно Улучшает Качество Генерации

В последние годы модели диффузии (DMs) стали стандартом для генеративных моделей в обработке перцептивных данных, таких как изображения, видео и аудио. Несмотря на их успех, основной недостаток этих моделей заключается в дорогой итеративной процедуре семплирования, требующей сотен или даже тысяч вызовов функции для генерации высококачественных образцов. Это ограничивает их применение в условиях низкой задержки.

Одним из подходов к улучшению эффективности семплирования является дистилляция моделей диффузии в модели, способные генерировать образцы за несколько шагов. Среди таких методов, модели консистенции (CMs) привлекают внимание благодаря своей простой идее и способности генерировать образцы всего за несколько итераций. CMs используют формулировку обыкновенного дифференциального уравнения (ODE) моделей диффузии, известную как поток вероятности (PF) ODE, для создания детерминированного отображения между шумом и данными.

Основные Понятия и Связанные Работы

Модели Диффузии

Основная цель моделей диффузии - научиться обратить процесс зашумления, который итеративно преобразует данные в шум. Этот процесс можно формализовать как стохастическое дифференциальное уравнение (SDE), называемое прямым SDE:

[ d x_t = \mu(x_t, t) dt + \sigma(t) dW_t, \quad x_0 \sim p_0 ]

где ( t \in [0, T] ), ( \mu ) и ( \sigma ) - гиперпараметры, а ( W_t ) - многомерное броуновское движение. Распределение ( p_t ) обозначает маргинальное распределение ( x_t ), и при правильном выборе гиперпараметров ( p_T ) становится почти чистым шумом.

Song et al. [48] показали, что следующий ODE, называемый PF ODE, имеет те же маргинальные распределения, что и прямой SDE:

[ d x_t = \left( \mu(x_t, t) - \frac{\sigma^2(t)}{2} \nabla \log p_t(x_t) \right) dt ]

где ( \nabla \log p_t ) - это функция оценки (Stein score function). Таким образом, если PF ODE запускается с ( x_0 \sim p_0 ), то ( x_t \sim p_t ).

Модели Консистенции

CMs используют PF ODE для обеспечения генерации образцов за несколько шагов. Они могут быть использованы как для дистилляции DM, так и для обучения с нуля. В нашем исследовании мы сосредоточимся на дистилляции, поскольку функция оценки предварительно обученных DM дает нам инструмент для прямого изучения влияния решения ODE на CMs.

Идея дистилляции модели консистенции заключается в обучении нейронной сети ( f_\theta ), чтобы ( f_\theta(x_{t_n}, t_n) \approx f_{\text{solver}}(x_{t_n}, t_n, 0) ) для всех ( n \in {1, \ldots, N} ). Другими словами, CMs учатся имитировать решатель ODE, тем самым избегая необходимости многократно вызывать функцию оценки во время семплирования.

CMs обучаются путем минимизации функции потерь консистентной дистилляции:

[ L_{\text{CD}} := \mathbb{E}{x_0 \sim p_0, n \sim \mathcal{U}{1, N}, x{t_n} \sim p_{t_n|0}(\cdot|x_0)} \left[ \lambda(t_n) d \left( f_\theta(x_{t_n}, t_n), \bar{f}\theta(\bar{x}{t_{n-1}}, t_{n-1}) \right) \right] ]

где ( p_{t|0} ) - ядро перехода, соответствующее уравнению (1), ( \lambda > 0 ) - весовая функция, ( d ) - любое расстояние, ( \bar{\theta} ) - замороженная версия ( \theta ), и ( \bar{x}{t{n-1}} = \Phi(x_{t_n}, t_n, t_{n-1}) ).

Прямые Модели Консистенции

В уравнении (4) ( x_0 ) и ( x_{t_n} ) не принадлежат к одной траектории ODE, поскольку для получения ( x_{t_n} ) из ( x_0 ) добавляется шум через прямой SDE. Таким образом, не имеет смысла требовать консистентности путем минимизации ( d(f_\theta(x_{t_n}, t_n), x_0) ), и вместо этого используется уравнение (4).

Song et al. [49] теоретически показали, что идеальная минимизация уравнения (4) с произвольно гибким ( f_\theta ) приводит к ( f_\theta(x_{t_n}, t_n) = f_{\text{solver}}(x_{t_n}, t_n, 0) ). Однако на практике было замечено, что CMs могут быть сложными для оптимизации, с медленной сходимостью или даже расходимостью.

Мы приписываем это поведение "слабому надзору" в потере CM, а именно тому, что ( f_\theta ) не обучается напрямую отображать ( x_{t_n} ) в начало его траектории ODE. Ограничение, что CM должна отображать любую точку на траектории ODE в начало этой траектории, лишь слабо соблюдается через граничное условие.

Для оценки влияния решения ODE на CMs, мы предложили более интуитивную и интерпретируемую вариацию его потери:

[ L_{\text{Direct CD}} := \mathbb{E}{x_0 \sim p_0, n \sim \mathcal{U}{1, N}, x{t_n} \sim p_{t_n|0}(\cdot|x_0)} \left[ \lambda(t_n) d \left( f_\theta(x_{t_n}, t_n), f_{\text{solver}}(x_{t_n}, t_n, 0) \right) \right] ]

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

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

Обучение

Для всех наших экспериментов мы стремились сравнить CMs и Direct CMs, используя крупномасштабные и современные DMs, обученные на интернет-данных, чтобы лучше отразить производительность этих моделей в реальных условиях. Мы выбрали SDXL, текст-картинку латентную модель диффузии с 2.6 миллиардами параметров U-Net, способную генерировать изображения с разрешением 1024px.

Оценка

Мы провели количественные сравнения, используя метрики, измеряющие качество решения ODE, а также качество изображений. Для оценки качества решения ODE мы использовали ошибку ( E ) (уравнение 6, меньше - лучше), которая применима только для генерации за один шаг. Для оценки качества изображений мы использовали такие метрики, как FID, FD-DINO, CLIP score и эстетический балл.

Анализ

Мы предоставили количественную оценку CMs и Direct CMs в таблице 1. Мы показали результаты для трех различных выборов численных решателей ODE, а именно DDIM, Euler и Heun. По всем метрикам, связанным с изображениями, CMs значительно превосходили Direct CMs, что указывает на то, что обучение с уравнением (4) приводит к значительно лучшему качеству изображений, чем обучение с уравнением (5).

Однако, в плане точности решения PF ODE, Direct CMs последовательно показывали лучшие результаты. Это парадоксально, поскольку цель CMs, как представлено Song et al. [49], заключается в верном решении PF ODE, и удивительно, что более точное решение может привести к худшему качеству изображений.

Заключение и Будущие Исследования

Несмотря на успехи моделей консистенции в дистилляции моделей диффузии в генераторы с несколькими шагами, мы обнаружили разрыв между их теорией и практикой. Решение PF ODE является центральным для теоретической мотивации CMs, но мы показали, что можно решать тот же PF ODE более точно с помощью Direct CMs, при этом генерируя образцы заметно худшего качества.

Мы предлагаем несколько возможных объяснений этому феномену:

1. Пространство латентов: Поскольку наши эксперименты проводились с латентными моделями диффузии, ODE определены в соответствующем латентном пространстве, и близость к решениям решателя, наблюдаемая в Direct CMs, может быть утрачена после декодирования в пространство пикселей.

2. Приближение истинной функции оценки: Если предварительно обученная модель диффузии не смогла точно аппроксимировать истинную функцию оценки, то даже если модель близко аппроксимирует решатель и эмпирический PF ODE, она не обязательно будет правильно аппроксимировать истинный PF ODE.

3. Индуктивное смещение: Хотя обе цели (уравнение 4 и уравнение 5) предназначены для имитации решателя на оптимуме, на практике этот оптимум никогда не достигается идеально, и цель CM может непреднамеренно предоставить полезное индуктивное смещение, улучшающее качество образцов.

Мы призываем сообщество к дополнительным исследованиям для выяснения этого кажущегося парадоксального поведения дистилляции моделей диффузии на основе ODE и его влияния на качество генерации образцов.