Classifier-free guidance

Задача

Условная модель $p(x\mid c)$ обучается так же, как безусловная, — просто с $c$ на входе. Проблема в том, что получается: сэмплы соответствуют условию слабее, чем хотелось бы. Модель честно воспроизводит $p(x\mid c)$, а хочется чего-то более выраженного.

Formально это признание в том, что мы хотим сэмплировать не из того распределения, которое выучили. Стоит сказать это прямо, потому что дальше вся техника — про то, из какого именно.

Формула

$$\tilde\varepsilon = \htmlData{k=uncond}{\varepsilon_\theta(x_t)} + \htmlData{k=w}{w}\,\htmlData{k=delta}{\big(\varepsilon_\theta(x_t, c) - \varepsilon_\theta(x_t)\big)}$$

Обучение — один трюк: во время обучения условие с вероятностью около $10\%$ заменяется пустым. Одна и та же сеть выучивает обе функции, и на генерации её вызывают дважды. Отсюда и цена: два прогона на шаг вместо одного.

Почему разность в скобках — это направление условия? В score-виде

$\nabla\log p(x\mid c) - \nabla\log p(x) = \nabla \log \frac{p(x\mid c)}{p(x)} = \nabla\log p(c\mid x)$

по правилу Байеса, поскольку $\log p(c)$ от $x$ не зависит. То есть разность — это градиент логарифма правдоподобия классификатора, взятый без всякого классификатора. В этом весь смысл названия.

Из какого распределения мы сэмплируем

Подставим и соберём:

$\tilde{s} = \nabla\log p(x) + w\,\nabla\log p(c\mid x) = \nabla\log\Big[p(x)\,p(c\mid x)^{w}\Big]$

То есть guidance сэмплирует из распределения, где правдоподобие условия возведено в степень $w$. При $w = 1$ это ровно $p(x\mid c)$; при $w > 1$ — что-то более сосредоточенное.

Оценивать это стоит трезво: это не более точная выборка из $p(x\mid c)$, а точная выборка из другого распределения. Метрики соответствия условию растут, метрики разнообразия падают, и компромисс здесь не артефакт реализации, а определение метода.

Что именно происходит с распределением

Здесь стоит поправить популярное «guidance делает распределение острее». Проверим два случая.

Случай гауссиан. Пусть $p(x\mid c) = \mathcal{N}(\mu, \sigma^2)$ и $p(x) = \mathcal{N}(0, \sigma^2)$. Тогда

$\tilde{s} = -\frac{x}{\sigma^2} + w\frac{\mu}{\sigma^2} = -\frac{x - w\mu}{\sigma^2}$

то есть направляемое распределение — это $\mathcal{N}(w\mu, \sigma^2)$. Среднее уехало в $w$ раз дальше, дисперсия не изменилась вовсе. Никакого «острее» здесь нет.

Случай смеси. Возьмём безусловную смесь двух мод с весами $0.5/0.5$ и условную с $0.8/0.2$ (обе с $\sigma = 0.6$). Численно:

$w$	среднее	ст. откл.	масса на левой моде
$0$	$0.000$	$1.616$	$0.500$
$1$	$-0.900$	$1.342$	$0.796$
$2$	$-1.330$	$0.924$	$0.938$
$3$	$-1.468$	$0.693$	$0.983$
$5$	$-1.520$	$0.583$	$0.999$
$10$	$-1.535$	$0.566$	$1.000$

Вот теперь разнообразие действительно падает — но конкретным образом: guidance выбирает моду, а не сжимает её. Стандартное отклонение падает с $1.62$ до $0.57$ и дальше не идёт: $0.57$ — это почти ширина одной компоненты ($\sigma = 0.6$). Всё, что guidance убрал, — это разнообразие между модами.

Точная формулировка, которую стоит унести: guidance перераспределяет массу к тем областям, где условие выполняется лучше. Выглядит это как «резкость», если моды соответствуют разным вариантам ответа, и не выглядит никак, если распределение унимодально.

Покрутите $\gamma$ в виджете — там та же операция, применённая к смеси трёх гауссиан.

блуждание

сила guidance γ

1

шаг ε

0.05

температура T

1

макс. |score|

17.82

длина пути

98.2

доля времени у мод

25.3%

Практика

$w$	что получается
$0$	безусловная генерация
$1$	честная условная модель, соответствие слабое
$3{-}8$	рабочий диапазон для изображений
$> 15$	пересыщенные цвета, артефакты, потеря разнообразия

Верхняя строка объясняет типичный артефакт. Сдвиг среднего (гауссов случай выше) в $w$ раз означает выход за пределы области, где модель обучалась: при больших $w$ значения пикселей вылезают за допустимый диапазон, и результат приходится обрезать. Отсюда и приёмы вроде динамического ограничения, и вариант с изменением $w$ по ходу сэмплирования — большой на ранних шагах, где решается композиция, и малый на поздних, где решаются детали.

Итог

• Guidance — одна линейная комбинация двух предсказаний, обучаемая случайным обнулением условия.
• Разность предсказаний равна $\nabla\log p(c\mid x)$ — классификатор без классификатора.
• Сэмплирование идёт из $p(x)p(c\mid x)^w$, то есть из другого распределения, а не более точно из условного.
• Гауссианы равной ширины при guidance сдвигаются, но не сужаются; сужение возникает из выбора моды в мультимодальном распределении.
• Цена: два прогона сети на шаг и потеря разнообразия, растущая с $w$.