Математика Stable Diffusion: от стохастических уравнений до генерации изображений

СправочникТема: Продвинутая математика

⚠️ ВНИМАНИЕ: ОЧЕНЬ СЛОЖНЫЙ МАТЕРИАЛ!

Этот материал требует знания высшей математики уровня 3-4 курса университета:

  • Теория вероятностей и математическая статистика
  • Дифференциальные уравнения
  • Математический анализ (несколько переменных)
  • Линейная алгебра
  • Базовое понимание машинного обучения

Не для школьной программы! Этот материал здесь для демонстрации того, насколько глубокой может быть математика 🚀

Введение

Stable Diffusion — это не просто "нейросеть для картинок", а математически изящная система, основанная на теории стохастических дифференциальных уравнений (SDE) и вариационном исчислении.

Основная идея простыми словами

Представьте, что вы берёте фотографию и постепенно добавляете к ней всё больше визуального "шума" (как помехи на старом телевизоре), пока картинка не превратится в полный хаос. Stable Diffusion учится делать обратное — превращать хаос обратно в осмысленное изображение!

Это как научиться восстанавливать разбитую вазу по её осколкам, причём делать это так хорошо, что можно даже "восстановить" вазы, которых никогда не существовало (но которые выглядят правдоподобно).

1. Прямой диффузионный процесс

Простыми словами: Берём чистую картинку и постепенно "портим" её, добавляя всё больше случайного шума на каждом шаге.

Начнём с основы: постепенного добавления шума к данным. Пусть x0∼q(x0)\mathbf{x}_0 \sim q(\mathbf{x}_0) — исходное изображение.

Марковская цепь с гауссовским шумом

Что это значит? На каждом шаге мы берём предыдущую картинку, чуть-чуть уменьшаем её яркость (1−βt\sqrt{1 - \beta_t}) и добавляем случайный шум (βt\beta_t). Повторяя это 1000 раз, получаем полный хаос.

Определим последовательность латентных переменных:

q(xt∣xt−1)=N(xt;1−βtxt−1,βtI)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})

где βt∈(0,1)\beta_t \in (0, 1) — schedule шума для шага tt (насколько сильно "портим" картинку на этом шаге).

Закрытая форма для произвольного t

Зачем это нужно? Вместо того чтобы делать 1000 шагов последовательно, мы можем сразу перейти к любому промежуточному состоянию! Это как телепортация в нужный момент времени.

Благодаря свойствам нормального распределения:

q(xt∣x0)=N(xt;αˉtx0,(1−αˉt)I)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})

где αt=1−βt\alpha_t = 1 - \beta_t и αˉt=∏s=1tαs\bar{\alpha}_t = \prod_{s=1}^t \alpha_s (это произведение показывает, сколько от оригинала осталось после tt шагов).

Reparameterization trick

Можем выразить явно:

xt=αˉtx0+1−αˉtϵ,ϵ∼N(0,I)\mathbf{x}_t = \sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\boldsymbol{\epsilon}, \quad \boldsymbol{\epsilon} \sim \mathcal{N}(0, \mathbf{I})

При t→∞t \to \infty: αˉt→0\bar{\alpha}_t \to 0, следовательно xT∼N(0,I)\mathbf{x}_T \sim \mathcal{N}(0, \mathbf{I}).

2. Обратный диффузионный процесс

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

Цель: научиться восстанавливать x0\mathbf{x}_0 из шума xT\mathbf{x}_T.

Обратное распределение

Истинное обратное распределение:

q(xt−1∣xt,x0)=N(xt−1;μ~t(xt,x0),β~tI)q(\mathbf{x}_{t-1} | \mathbf{x}_t, \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_{t-1}; \tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t, \mathbf{x}_0), \tilde{\beta}_t \mathbf{I})

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

μ~t(xt,x0)=αˉt−1βt1−αˉtx0+αt(1−αˉt−1)1−αˉtxt\tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t, \mathbf{x}_0) = \frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1 - \bar{\alpha}_t}\mathbf{x}_0 + \frac{\sqrt{\alpha_t}(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t}\mathbf{x}_t

и дисперсия:

β~t=1−αˉt−11−αˉtβt\tilde{\beta}_t = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t}\beta_t

Параметризация модели

Обучаем нейросеть pθ(xt−1∣xt)p_\theta(\mathbf{x}_{t-1} | \mathbf{x}_t) для аппроксимации q(xt−1∣xt,x0)q(\mathbf{x}_{t-1} | \mathbf{x}_t, \mathbf{x}_0):

pθ(xt−1∣xt)=N(xt−1;μθ(xt,t),Σθ(xt,t))p_\theta(\mathbf{x}_{t-1} | \mathbf{x}_t) = \mathcal{N}(\mathbf{x}_{t-1}; \boldsymbol{\mu}_\theta(\mathbf{x}_t, t), \boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t))

3. Функция потерь

Вариационная нижняя граница (ELBO)

Минимизируем отрицательное log-правдоподобие:

E[−log⁡pθ(x0)]≤Eq[−log⁡pθ(x0:T)q(x1:T∣x0)]=L\mathbb{E}[-\log p_\theta(\mathbf{x}_0)] \leq \mathbb{E}_q\left[-\log \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T}|\mathbf{x}_0)}\right] = L

Раскладываем:

L=Eq[DKL(q(xT∣x0)∥p(xT))+∑t=2TDKL(q(xt−1∣xt,x0)∥pθ(xt−1∣xt))−log⁡pθ(x0∣x1)]L = \mathbb{E}_q\left[D_{KL}(q(\mathbf{x}_T|\mathbf{x}_0) \| p(\mathbf{x}_T)) + \sum_{t=2}^T D_{KL}(q(\mathbf{x}_{t-1}|\mathbf{x}_t,\mathbf{x}_0) \| p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t)) - \log p_\theta(\mathbf{x}_0|\mathbf{x}_1)\right]

Упрощённая функция потерь (Ho et al., 2020)

Вместо предсказания среднего, предсказываем шум:

Lsimple=Et,x0,ϵ[∥ϵ−ϵθ(xt,t)∥2]L_{simple} = \mathbb{E}_{t, \mathbf{x}_0, \boldsymbol{\epsilon}}\left[\|\boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t)\|^2\right]

где xt=αˉtx0+1−αˉtϵ\mathbf{x}_t = \sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\boldsymbol{\epsilon}.

4. Score matching и уравнение Ланжевена

Score function

Определим score function:

∇xtlog⁡q(xt)=−ϵ1−αˉt\nabla_{\mathbf{x}_t} \log q(\mathbf{x}_t) = -\frac{\boldsymbol{\epsilon}}{\sqrt{1 - \bar{\alpha}_t}}

Деноising score matching

Модель ϵθ\boldsymbol{\epsilon}_\theta учится приближать:

sθ(xt,t)≈∇xtlog⁡q(xt)\boldsymbol{s}_\theta(\mathbf{x}_t, t) \approx \nabla_{\mathbf{x}_t} \log q(\mathbf{x}_t)

Annealed Langevin dynamics

Генерация происходит через дискретизацию SDE:

xt−1=xt+βt2∇xtlog⁡q(xt)+βtz,z∼N(0,I)\mathbf{x}_{t-1} = \mathbf{x}_t + \frac{\beta_t}{2}\nabla_{\mathbf{x}_t}\log q(\mathbf{x}_t) + \sqrt{\beta_t}\mathbf{z}, \quad \mathbf{z} \sim \mathcal{N}(0, \mathbf{I})

5. Latent Diffusion (Stable Diffusion)

Проблема: Работать с огромными изображениями (например, 512×512 пикселей = 786432 числа!) очень медленно и требует много памяти.

Решение: Сначала сжимаем изображение в компактное представление (как ZIP-архив), работаем с этим сжатым вариантом, а потом распаковываем обратно!

Автоэнкодер в латентном пространстве

Вместо работы с изображениями x∈RH×W×3\mathbf{x} \in \mathbb{R}^{H \times W \times 3}, используем VAE encoder E\mathcal{E} (автоэнкодер, который умеет сжимать и разжимать картинки):

z=E(x),z∈Rh×w×c\mathbf{z} = \mathcal{E}(\mathbf{x}), \quad \mathbf{z} \in \mathbb{R}^{h \times w \times c}

где h≪Hh \ll H, w≪Ww \ll W (обычно в 8 раз меньше).

Диффузия в латентном пространстве

Применяем весь диффузионный процесс к z\mathbf{z}:

LLDM=EE(x),ϵ,t[∥ϵ−ϵθ(zt,t,τθ(y))∥2]L_{LDM} = \mathbb{E}_{\mathcal{E}(\mathbf{x}), \boldsymbol{\epsilon}, t}\left[\|\boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta(\mathbf{z}_t, t, \tau_\theta(y))\|^2\right]

где τθ(y)\tau_\theta(y) — эмбеддинг текстового промпта через CLIP.

Cross-attention mechanism

Внедрение условия yy (текста) через cross-attention:

Attention(Q,K,V)=softmax(QKTd)V\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d}}\right)V

где:

  • Q=WQ⋅ϕi(zt)Q = W_Q \cdot \phi_i(\mathbf{z}_t) — запросы из изображения
  • K=WK⋅τθ(y)K = W_K \cdot \tau_\theta(y), V=WV⋅τθ(y)V = W_V \cdot \tau_\theta(y) — ключи и значения из текста

6. Classifier-free guidance

Простыми словами: Это как ручка "насыщенности" для вашего промпта. Чем выше значение, тем сильнее модель старается соответствовать вашему описанию (но если переборщить — картинка становится неестественной).

Как это работает? Модель генерирует два варианта:

  1. Без промпта (случайная картинка)
  2. С промптом (то, что вы хотите)

Затем берём второй вариант и "усиливаем" его, отдаляя от первого:

ϵ~θ(zt,t,y)=ϵθ(zt,t,∅)+s⋅(ϵθ(zt,t,y)−ϵθ(zt,t,∅))\tilde{\boldsymbol{\epsilon}}_\theta(\mathbf{z}_t, t, y) = \boldsymbol{\epsilon}_\theta(\mathbf{z}_t, t, \emptyset) + s \cdot (\boldsymbol{\epsilon}_\theta(\mathbf{z}_t, t, y) - \boldsymbol{\epsilon}_\theta(\mathbf{z}_t, t, \emptyset))

где s>1s > 1 — guidance scale (обычно 7-15). Низкие значения = креативность, высокие = точное следование промпту.

7. Samplers: от DDPM до DPM-Solver

Зачем нужны сэмплеры? Это разные способы пройти путь от шума к картинке. Одни медленнее но точнее, другие быстрее но менее детальны. Как разные маршруты из точки А в точку Б.

DDPM sampler

Что это: Оригинальный "медленный и верный" метод. Делает ВСЕ 1000 шагов по порядку.

Базовый сэмплер из оригинальной статьи:

xt−1=1αt(xt−1−αt1−αˉtϵθ(xt,t))+β~tz\mathbf{x}_{t-1} = \frac{1}{\sqrt{\alpha_t}}\left(\mathbf{x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}}\boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t)\right) + \sqrt{\tilde{\beta}_t}\mathbf{z}

DDIM sampler (детерминированный)

Что это: "Быстрый поезд" — пропускает промежуточные станции! Вместо 1000 шагов делает, например, только 50.

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

Ускоренная генерация через прыжки по шагам:

xt−1=αˉt−1(xt−1−αˉtϵθ(xt,t)αˉt)⏟предсказанный x0+1−αˉt−1ϵθ(xt,t)\mathbf{x}_{t-1} = \sqrt{\bar{\alpha}_{t-1}}\underbrace{\left(\frac{\mathbf{x}_t - \sqrt{1-\bar{\alpha}_t}\boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t)}{\sqrt{\bar{\alpha}_t}}\right)}_{\text{предсказанный } \mathbf{x}_0} + \sqrt{1 - \bar{\alpha}_{t-1}}\boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t)

DPM-Solver (2022)

Что это: Самый умный метод! Использует продвинутую математику (численные методы для дифференциальных уравнений) чтобы делать меньше, но более точные шаги.

Аналогия: Вместо того чтобы идти мелкими шажками или большими прыжками, DPM-Solver вычисляет оптимальную траекторию, как GPS прокладывает лучший маршрут.

Интерпретация диффузии как ODE и использование high-order численных методов:

dxtdt=f(t)xt+g(t)∇xtlog⁡qt(xt)\frac{d\mathbf{x}_t}{dt} = f(t)\mathbf{x}_t + g(t)\nabla_{\mathbf{x}_t}\log q_t(\mathbf{x}_t)

Заключение

Подведём итог простыми словами:

Stable Diffusion работает так:

  1. Учится портить картинки (добавлять шум) — это легко
  2. Учится делать обратное (убирать шум) — это сложно, но именно этим и занимается обучение
  3. Работает со сжатыми версиями изображений для скорости
  4. Понимает текст через CLIP и внедряет его через attention
  5. Может генерировать быстро благодаря умным сэмплерам

Вся эта математика нужна, чтобы вы могли написать "кот на скейтборде" и получить картинку! 🎨🐱🛹

С точки зрения математики

Stable Diffusion — это синтез:

  • Стохастических процессов (марковские цепи, SDE)
  • Вариационного вывода (ELBO, KL-дивергенция)
  • Score-based моделей (denoising score matching)
  • Глубокого обучения (U-Net, attention, CLIP)

Вся эта математика работает для того, чтобы мы могли написать "a cat riding a skateboard" и получить картинку! 🎨🐱🛹

Литература

  1. Ho et al. (2020) — Denoising Diffusion Probabilistic Models
  2. Song et al. (2021) — Score-Based Generative Modeling through SDEs
  3. Rombach et al. (2022) — High-Resolution Image Synthesis with Latent Diffusion Models
  4. Lu et al. (2022) — DPM-Solver: Fast ODE Solver for Diffusion Models