Score Distillation Sampling в задаче text-to-3D

Введение

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

И после успеха в 2D, диффузионные модели решили начать применять в 3D-задачах генерации. Оказалось, что это не так уж просто.

Сегодня мы рассмотрим основные проблемы использования 3D-диффузионных моделей, а также подход, позволяющий обойти ограничения и применить для генерации в 3D привычную text-to-image диффузионную модель.

Диффузия в 3D

Авторы многих работ пытались обучить диффузию на генерацию 3D-объектов. Однако у такого подхода есть ряд ограничений:

1. Сложно сделать обучение устойчивым — возникают проблемы с 3D-репрезентацией
2. 3D-данных гораздо меньше, чем 2D-данных

Именно поэтому появилась следующая идея — использовать хорошо предобученную text-to-image модель и как-то дистиллировать из неё знания в нашу 3D-модельку.

Диффузия в 2D

Напомним, что представляет собой диффузионная модель, и как она учится.

Наша основная задача — научиться генерировать изображения, стартуя из нормального шума и постепенно его убирая, оставляя лишь само изображение.

Диффузионный процесс состоит из двух частей:

• forward process — постепенное добавление шума к распределению объектов до тех пор, пока итоговое распределение не станет неотличимым от случайного шума (чаще всего работают с гауссовским шумом, но недавно была представлена работа, показывающая, как можно работать и с другими типами шумов). Важный фактор — процесс марковский, и каждый следующий этап зашумления зависит только от предыдущего.
• backward process — зеркальное отражение прямого процесса. Начав с чистого шума, модель постепенно удаляет шум и приближает данные к исходным чистым данным.

Прямой процесс (forward process)

Пусть у нас есть сэмпл \(x_0 \sim q(x)\) из исходного распределения данных. Тогда прямой процесс диффузии — постепенное добавление небольшого количества шума, чаще всего Гауссовского (далее работаем с ним), в течение \(T\) шагов, в результате которых получаем последовательность зашумленных сэмплов \(x_1, x_2 …, x_T\) и в итоге — гауссовский шум с нулевым средним и единичной дисперсией. Количество добавляемого шума регулируется variance scheduler \(\{\beta_t \in (0,1)\}^T_t\). Он показывает, как много мы оставляем от исходной картинки и берём именно шума на каждом шаге.

Тогда зашумлённое распределение на шаге \(t\) можно записать следующим образом:

\(q(x_t|x_{t-1}) = \mathcal{N}(x_t; \sqrt{1-\beta_t}x_{t-1}, \beta_t \mathbf{I})\)

\(q(x_{1:T}|x_0) = \prod_{t=1}^{T}q(x_t|x_{t-1})\)

Помимо распределений нам нужно смотреть или выполнять операции с элементами этих распределений, то есть сэмплировать \(x_t\).

Используем трюк репараметризации:

\(x_t = \sqrt{1-\beta_t}x_{t-1} + \sqrt{\beta_t}\epsilon_{t-1}\)

Обратный процесс (backward process)

Если мы сможем обернуть процесс — получим генеративную модель, которая возьмёт сэмпл из Гауссовского шума с нулевым средним и единичной дисперсией \(x_T \sim \mathcal{N}(\mathbf{0}, \mathbf{I})\) и восстановит элемент из исходного распределения.

Важная особенность диффузионного процесса: если \(\beta_t\) — достаточно маленькое значение, то \(q(x_{t-1}|x_t)\) тоже будет Гауссовским. Однако мы не можем так просто оценить это распределение, поскольку для оценки нам нужно всё распределение исходного датасета. Поэтому нашей задачей становится аппроксимация с помощью какой-либо модели \(p_{\theta}\):

\(p_{\theta}(x_{0:T}) = p(x_T) \prod_{t=1}^{T}p_{\theta}(x_{t-1}|x_{t})\)

\( p_{\theta}(x_{t-1} \mid x_t) = \mathcal{N}(x_{t-1}; \mu_{\theta}(x_t, t), \Sigma_{\theta}(x_t, t)) \)

Лосс функция

Обучается диффузионная модель за счёт минимизации следующего лосса:
\( \mathcal{L}_{\text{Diff}}(\phi, x) = \mathbb{E}{\substack{t \sim \mathcal{U}(0,1) \ \epsilon \sim \mathcal{N}(0, \mathbf{I})}} \left[ w(t) \left| \epsilon_{\phi}(\alpha_t x + \sigma_t \epsilon;\, t) — \epsilon \right||_2^2 \right] \)

Таким образом, работая в пиксельном (латентном) пространстве, диффузионный процесс во время генерации выглядит как постепенное обновление пиксельного пространства за счёт уменьшения количества шума.

Но как теперь это использовать для обучения 3D-моделей?

Dreamfusion и Score Distillation Sampling

Интуиция

Напомним, наша проблема заключается в том, что обучать модель напрямую на 3D-данных сложно — их мало, нужны специальные трюки для работы с различными 3D-представлениями. Вместо этого DreamFusion использует уже обученную 2D-диффузионную модель, которая умеет генерировать 2D-изображения, и «переносит» её знания в 3D-пространство.

То есть мы хотим обучить нашу 3D-модель так, чтобы рендеры картинок с разных ракурсов лежали в пространстве, выученном 2D-диффузионной моделью.

Итак, мы хотим, чтобы наша диффузионная модель, обученная на огромном количестве 2D-изображений, выступала в роли критика. Она знает, как выглядит «белка на мотоцикле» на 2D-картинке, и может сказать, насколько рендер 3D-модели похож на нужный нам рендер.

Математика

Теперь давайте разберёмся с математикой процесса.

Пусть у нас есть генератор с параметрами \(\theta\), который может выдать нам некоторый рендер с определённого угла \(x\): \(x = g(\theta)\).

Мы зашумляем его до конкретного шага \(t\) и делаем один шаг расшумления с кондишнингом диффиузии на заданный текстовый промпт, а затем считаем обычный диффузионный лосс:

\(\mathcal{L}_{\text{Diff}}(\phi, x) = \mathbb{E}{t \sim \mathcal{U}(0,1), \epsilon \sim \mathcal{N}(0, \mathbf{I})} \left[ w(t) \| \epsilon_{\phi}(\alpha_t x + \sigma_t \epsilon; t) — \epsilon \|_2^2 \right]\)

Далее минимизируем этот лосс по параметрам нашего 3D-генератора:

\(\theta^* = \arg\min_{\theta} \mathcal{L}_{\text{Diff}}(\phi, x = g(\theta))\)

Весь пайплайн выглядит следующим образом:

1. Мы инициализируем нашу 3D-модель, делая рендер под углом:

2. Зашумляем этот рендер до определённого шага t: \(z_t = \alpha_t x + \sigma_t \epsilon\)

3. Делаем шаг расшумления с кондишнингом на промпт с помощью UNet диффузионной модели text-to-image — предсказываем \(\epsilon_{\phi}(\alpha_t x + \sigma_t \epsilon; t)\)

4. Считаем диффузионный лосс:

\(\mathcal{L}_{\text{Diff}}(\phi, x) = \mathbb{E}{t \sim \mathcal{U}(0,1), \epsilon \sim \mathcal{N}(0, \mathbf{I})} \left[ w(t) \| \epsilon_{\phi}(\alpha_t x + \sigma_t \epsilon; t) — \epsilon \|_2^2 \right]\)

Всё вместе это выглядит так:

Подробности о 3D-модели из статьи

DreamFusion использует NeRF (Neural Radiance Fields) для представления 3D-сцены. NeRF — это нейронная сеть с параметрами \(\theta\), которая параметризует 3D-сцену.

Она принимает на вход координаты точки в пространстве (x,y,z) и направление взгляда (d), а на выходе выдаёт цвет (RGB) и плотность (density) этой точки.

То есть вместо явного 3D-объекта NeRF запоминает, как свет проходит через сцену.

Как создаётся изображение?

1. Из камеры выпускаются лучи: \(\mathbf{r}(t) = \mathbf{o} + t \mathbf{d}\), где:
   \(\mathbf{o}\) — начало луча (точка камеры);
   \(\mathbf{d}\) — направление луча;
   \(t\) — расстояние вдоль луча.
   Часто направление взгляда параметризуется углами \((\theta, \phi)\), где \(\theta\) — угол от оси \(z\), \(\phi\) — угол в плоскости \((x, y)\).
2. На пути луча мы сэмплируем множество точек (на разной глубине проверяем цвет и плотность): \(F_{\Theta}(\mathbf{x}, \mathbf{d}) = (\mathbf{c}, \sigma)\)
3. Далее мы с помощью рендерига агрегируем полученный результат в картинку-рендер: \(C(\mathbf{r}) = \int_{t_n}^{t_f} T(t) \sigma(\mathbf{r}(t)) \mathbf{c}(\mathbf{r}(t), \mathbf{d}) dt\) где: \(T(t) = \exp\left(- \int_{t_n}^{t} \sigma(\mathbf{r}(s)) ds \right)\) — прозрачность (количество света, проходимого до точки \(t\)). На практике этот интеграл аппроксимируется суммированием по дискретным точкам вдоль луча: \(\hat{C}(\mathbf{r}) = \sum_{i=1}^{N} T_i (1 — e^{-\sigma_i \delta_i}) \mathbf{c}_i\) где: \(T_i = \exp\left(-\sum_{j=1}^{i-1} \sigma_j \delta_j \right)\)
   

Лосс

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

\(\mathcal{L}_{\text{Diff}}(\phi, x) = \mathbb{E}{t \sim \mathcal{U}(0,1), \epsilon \sim \mathcal{N}(0, \mathbf{I})} \left[ w(t) \| \epsilon_{\phi}(\alpha_t x + \sigma_t \epsilon; t) — \epsilon \|_2^2 \right]\)

Для выполнения шага градиентного спуска считаем градиент диффузионного лосса по параметрам 3D-генератора:

Однако нам нужно посчитать якобиан UNet, что является довольно «дорогой» по времени и ресурсам операцией.

Поэтому авторы предлагают не учитывать этот якобиан:

Тогда итоговый градиент будет выглядеть следующим образом:

\(\nabla_{\theta} \mathcal{L}_{\text{SDS}}(\phi, \mathbf{x}=g(\theta)) \triangleq \mathbb{E}{\mathbf{t}, \epsilon} \left[ w(t) \left(\hat{\epsilon}_{\phi}(z_t; y, t)- \epsilon \right) \frac{\partial x}{\partial \theta} \right]\)

Математика

За избавлением от градиента по UNet стоит интуиция, связанная с минимизацией KL-дивергенции между диффузионным распределением и полученным от генератора.

KL-дивергенция определяется следующим образом:

\({\text{KL}}(P \| Q) = \int_{-\infty}^{\infty} p(x) \log \left( \frac{p(x)}{q(x)} \right) \mathrm{d}x\)

Интуитивно она показывает различие между двумя распределениями (чем меньше значение, тем ближе два распределения).

Тогда для распределения от 3D-генератора (\(q(\mathbf{z}t \mid \mathbf{x} = g(\theta)\)) и диффузионного распределения (\(p{\phi}(\mathbf{z}_t \mid y)\)) мы запишем:

\(\text{KL}(q(\mathbf{z}t \mid \mathbf{x} = g(\theta)) \| p{\phi}(\mathbf{z}t \mid y)) = \mathbb{E}{\epsilon} \left[ \log q(\mathbf{z}t \mid \mathbf{x} = g(\theta)) — \log p{\phi}(\mathbf{z}_t \mid y) \right]\)

Градиент в этом случае равен:

\(\nabla_{\theta} \text{KL}(q(\mathbf{z}t \mid \mathbf{x} = g(\theta)) \| p{\phi}(\mathbf{z}t \mid y)) = \mathbb{E}{\epsilon} \left[ \nabla_{\theta} \log q(\mathbf{z}t \mid \mathbf{x} = g(\theta)) — \nabla{\theta} \log p_{\phi}(\mathbf{z}_t \mid y) \right]\)

Откуда можно записать:

\(\nabla_{\theta} \mathcal{L}_{\text{SDS}}(\phi, \mathbf{x} = g(\theta)) = \nabla{\theta} \mathbb{E}{t} \left[ \sigma{t} / \alpha_{t} w(t) \text{KL} \left( q \left( \mathbf{z}{t} \mid g(\theta); y, t \right) \| p{\phi} \left( \mathbf{z}{t}; y, t \right) \right) \right] = \mathbb{E}{t, \epsilon} \left[ w(t) \left( \hat{\epsilon}(\mathbf{z}_t \mid y) — \epsilon \right) \frac{\partial \mathbf{x}}{\partial \theta} \right]\)

Более подробно этот вывод можно посмотреть в оригинальной работе DreamFusion в аппендиксе.

Итоговый алгоритм

Итоговый алгоритм text-to-3D генерации становится довольно понятным, а сама оптимизация уже не требует много времени, так как мы не используем градиент по большой модели UNet.

Результаты

Теперь самое интересное — результаты:

Выглядит неплохо! Есть и пара недостатков:

1. Рендеры приходится генерировать в маленьком разрешении, что в результате даёт плохое качество геометрии и текстуры. Обучение получается медленным.
2. Необходимо проводить оптимизацию под каждый промпт отдельно.
3. Возникает проблема — Janus (consistency) problem:

При обучении у нас может появиться проблема с 3D-консистентностью, так как на ренедерах диффузионная модель не получает полную 3D-информацию. Это, в свою очередь, приводит к нескольким головам, хвостам и так далее.

Эти проблемы решаются в следующих работах (хоть и не до конца) 🙂

Применение SDS

Кроме задач text-to-3D (image-to-3D), SDS может использоваться, например, в Image Editing и text-to-SVG задачах.

Editing

Одним из примеров использования SDS в задаче Image Editing является работа Delta Denoising Score. Авторы показали, как работает SDS loss для редактирования картинок, а затем улучшили данный подход для более качественного редактирования:

Как можно заметить, прямое использование SDS для редактирования не даёт очень качественного результата 🙂

Text-to-SVG

Векторная графика (SVG) представляется кривыми и контурами, а не пикселями.

Здесь идея использования полностью совпадает с идеей Text-to-3D, только мы используем специальный дифференцируемый рендеринг для преобразования SVG-представления в растровое изображение.

Затем к нему применяем SDS:

Пример работы, которая предлагает такой подход, — Text-to-svg by abstracting pixel-based diffusion models.

Заключение

Таким образом, мы рассмотрели Score Distillation Sampling для дистилляции знания диффузионной модели в 3D-генератор, и обсудили работу DreamFusion, которые использовали этот метод для генерации text-to-3D, основные её недостатки.

Эта идея применима не только к задаче text-to-3D, но и к задачам Image Editing и Text-to-SVG. Одним словом, её можно использовать на практике для своих кейсов и радоваться 😎