Вариационные автоэнкодеры

Заметка об автоэнкодерах закончилась открытым вопросом: обычный автоэнкодер хорошо сжимает и реконструирует, но генератором не является. Мы не можем выбрать точку в его латентном пространстве и ожидать осмысленную картинку. Вариационный автоэнкодер (variational autoencoder, VAE), предложенный Kingma and Welling (2014) и независимо Rezende, Mohamed and Wierstra (2014), исправляет это, упорядочивая латентное пространство. Каждый вход кодируется не в одну точку, а в гауссиану, и эти гауссианы притягиваются к одному фиксированному априорному распределению. Кроме того, VAE — первая нейросетевая генеративная модель курса, которая обучается максимизацией правдоподобия (точнее, его нижней оценки), то есть по той целевой функции, которую заметка об энтропии назвала естественной.

Обозначения. $x \in \mathbb R^D$ — точка данных (изображение из $D$ значений пикселей), $z \in \mathbb R^d$ — её латентный код, $\hat x$ — реконструкция. $p(z)$ — априорное распределение кодов. $\theta$ — параметры декодера, $\varphi$ — параметры энкодера, как в заметке об автоэнкодерах. $\mathcal N(\mu, \sigma^2)$ всегда означает распределение со средним $\mu$ и дисперсией $\sigma^2$. Для векторов $\mathcal N(\mu, \sigma^2)$ — сокращённая запись $\mathcal N(\mu, \operatorname{diag}(\sigma^2))$, а $\mathcal N(0, I)$ — стандартное нормальное распределение в $d$ измерениях. $\odot$ — поэлементное произведение.

Порядок чтения. Основной путь идёт от модели через ELBO и трюк репараметризации к тому, что даёт упорядоченное латентное пространство, а затем к KL в замкнутом виде в разделе «Что фиксировано». Вывод ELBO, матричная форма KL и разложение слагаемого KL по агрегированному апостериорному распределению — дополнительный материал: они обосновывают и уточняют основное, и их можно прочитать позже.

От автоэнкодеров к генерации#

Напомним конвейер «энкодер — декодер»:

$$\underset{\text{данные}}{x} \;\xrightarrow{\ \text{энкодер}\ }\; z \;\xrightarrow{\ \text{декодер}\ }\; \underset{\text{реконструкция}}{\hat x} .$$

Код $z$ — это представление меньшей размерности (lower-dimensional representation) (512 чисел у свёрточного автоэнкодера из предыдущей заметки), а задача декодера — реконструкция. Как обученный автоэнкодер мог бы генерировать новые изображения? Есть два очевидных способа, и оба не работают.

1. Сэмплировать случайный $z$ и декодировать его. Возьмём, например, $z$ из стандартного нормального распределения. Коды обычного автоэнкодера ничем не привязаны к $\mathcal N(0, I)$, поэтому такой $z$ может оказаться далеко от всех обучающих кодов, там, где декодер никогда не обучался, и нет никаких причин, чтобы он декодировался в цифру: латентное пространство не упорядочено. Обучение требует лишь, чтобы коды обучающих изображений декодировались обратно в эти изображения. Ничто не говорит, что лежит между кодами и даже где в $\mathbb R^d$ они находятся. Подгонка гауссианы к кодам каждого класса цифр, как в заметке об автоэнкодерах, работает лучше, но требует меток и статистик по классам, часть её сэмплов всё равно сильно искажена, а интерполированные коды декодируются в наложение двух цифр.
2. Сэмплировать латентные векторы рядом с опорным. Закодируем реальное изображение в $z_{\text{ref}}$ и декодируем соседние точки $z_1, z_2, z_3$. Соседняя точка тоже не обязана быть осмысленным кодом: если она попадает в промежуток между кодами обучающих изображений, декодер на ней никогда не обучался, и декодированный сосед «7» вовсе не обязан походить на цифру.

Рисунок

Обычный автоэнкодер не умеет генерировать

Решение — сделать распределение кодов частью модели. Для этого нужен язык моделей с латентными переменными (latent-variable models).

Совместное, маргинальное и условное распределения#

Возьмём две случайные величины: данные $x$ и латентную переменную $z$. Их совместная плотность $p(x, z)$ — облако на плоскости $(x, z)$. Проекции облака на каждую из осей дают маргинальные распределения $p(x)$ и $p(z)$. Маргинализация вычисляет маргинальное распределение, интегрируя по другой переменной:

$$p(x) = \int p(x, z)\, dz .$$

Условная вероятность — это совместная плотность, ограниченная на фиксированное значение условия и заново нормированная:

$$p(x \mid z) = \frac{p(x, z)}{p(z)} .$$

Например, $p(x \mid z = 3)$ — срез $p(x,z)$ вдоль прямой $z = 3$, делённый на $p(z = 3)$, чтобы интеграл по $x$ был равен единице.

Рисунок

Совместное, маргинальные и условное распределения

Модель с латентными переменными и её апостериорное распределение#

Модель с латентными переменными генерирует данные в два шага: сначала берёт код из априорного распределения (prior) $p(z)$, затем изображение из декодера (правдоподобия) $p_\theta(x \mid z)$. Априорное распределение выбирается, а не обучается; стандартный выбор — $p(z) = \mathcal N(0, I)$. Распределение модели на изображениях — это маргинальное распределение

$$p_\theta(x) = \int p_\theta(x \mid z)\, p(z)\, dz ,$$

которое называют обоснованностью (evidence). Обучение методом максимального правдоподобия означает максимизацию $\log p_\theta(x)$ на данных.

Обратное направление, «какие коды объясняют это изображение», — это апостериорное распределение (posterior). По теореме Байеса

$$p_\theta(z \mid x) = \frac{p_\theta(x \mid z)\, p(z)}{p_\theta(x)} ,$$

и обоснованность в знаменателе — интеграл по всему латентному пространству. Для декодера-нейросети замкнутой формы у него нет. В высокой размерности квадратурные формулы безнадёжны, а наивное усреднение $p_\theta(x \mid z)$ методом Монте-Карло по $z \sim p(z)$ имеет слишком большую дисперсию, чтобы применять его на каждом шаге обучения: почти все сэмплы из априорного распределения плохо объясняют данный $x$. Поэтому и правдоподобие, которое мы хотим максимизировать, и апостериорное распределение, которое понадобилось бы для кодирования, вычислительно неразрешимы (intractable).

Рисунок

Распределение данных и латентное распределение

Вариационная аппроксимация и ELBO#

Выход в том, чтобы приблизить апостериорное распределение членом простого семейства распределений. VAE использует диагональную гауссиану, среднее и дисперсии которой выдаёт сеть-энкодер с параметрами $\varphi$:

$$q_\varphi(z \mid x) = \mathcal N\big(\mu_\varphi(x), \sigma_\varphi^2(x)\big) \approx p_\theta(z \mid x).$$

Замена неразрешимого апостериорного распределения ближайшим к нему по $D_{\mathrm{KL}}\big(q \,\|\, p_\theta(z \mid x)\big)$ членом простого семейства — это вариационный вывод (variational inference) (Blei, Kucukelbir and McAuliffe, 2017). Поскольку одна сеть строит приближение для каждого $x$ вместо отдельной оптимизации для каждой точки данных, такой вывод называют амортизированным (amortized) (Kingma and Welling, 2019, §2.1). Максимизируемая величина — нижняя оценка обоснованности (evidence lower bound, ELBO):

$$\begin{aligned} \mathcal L(x) ={}& \underbrace{\mathbb E_{q_\varphi(z\mid x)}\big[\log p_\theta(x \mid z)\big]}_{\text{реконструкция}} \\ &-\underbrace{D_{\mathrm{KL}}\big(q_\varphi(z \mid x)\,\|\,p(z)\big)}_{\text{регуляризация}} \\ \le{}& \log p_\theta(x). \end{aligned}$$

Неравенство выводится ниже. Два его слагаемых тянут в противоположные стороны.

Реконструкция измеряет, насколько хорошо изображение переживает путь «изображение → код → изображение». С гауссовским декодером она сводится к функции потерь реконструкции $L_2$. Если $p_\theta(x \mid z) = \mathcal N\big(x;\, \hat x(z),\, s^2 I\big)$, где сеть-декодер выдаёт среднее $\hat x(z)$, а $s^2$ — фиксированная дисперсия, то

$$\begin{aligned} -\log p_\theta(x \mid z) ={}& \frac{1}{2s^2}\,\big\|x - \hat x(z)\big\|_2^2 \\ &+ \frac D2 \log(2\pi s^2), \end{aligned}$$

и последнее слагаемое не зависит от сети. При $s^2 = \tfrac12$ вес квадратичной ошибки равен ровно 1, поэтому квадратичная ошибка, просуммированная по $D$ пикселям, плюс слагаемое KL-дивергенции (KL divergence) — это отрицательная ELBO с точностью до константы. (Среднему по пикселям соответствовало бы $s^2 = D/2$.) Другая фиксированная $s^2$ меняет баланс между двумя слагаемыми: умножение функции потерь на $2s^2$ даёт $\|x - \hat x(z)\|_2^2 + 2s^2\, D_{\mathrm{KL}}$, то есть вес KL $\beta = 2s^2$ (см. β-VAE ниже).

Регуляризация измеряет, насколько каждое апостериорное распределение $q_\varphi(z \mid x)$ далеко от априорного. Она тянет каждое распределение кода к $\mathcal N(0, I)$, так что все коды живут в одной хорошо известной области латентного пространства.

Рисунок

Два слагаемых ELBO

Итак, вместо того чтобы отображать вход в одну точку, VAE отображает его в гауссиану. Энкодер выдаёт $\mu$ и $\log\sigma^2$, мы сэмплируем $z \sim \mathcal N(\mu, \sigma^2)$, и декодер реконструирует вход по $z$. Функция потерь при обучении — отрицательная ELBO с $s^2 = \tfrac12$, оценённая по одному сэмплу $z \sim \mathcal N(\mu, \sigma^2)$ на каждый вход:

$$\begin{aligned} \mathcal L_{\text{VAE}}(x) ={}& \underbrace{\big\|x - \hat x(z)\big\|_2^2}_{\text{реконструкция}} \\ &+ D_{\mathrm{KL}}\big(\mathcal N(\mu, \sigma^2)\,\|\,\mathcal N(0, I)\big). \end{aligned}$$

Слагаемое KL теперь входит со знаком плюс, потому что мы минимизируем.

Трюк репараметризации#

Функция потерь содержит сэмпл $z \sim \mathcal N(\mu, \sigma^2)$, а через сэмплирование нельзя провести обратное распространение ошибки: случайный сэмпл не является дифференцируемой функцией $\mu$ и $\sigma$. Решение — трюк репараметризации (reparameterization trick): перенести случайность во вход, ровно как в заметке о гауссианах:

$$z = \mu + \sigma \odot \varepsilon, \qquad \varepsilon \sim \mathcal N(0, I).$$

$\varepsilon$ сэмплируется независимо от сети, затем масштабируется на $\sigma$ и сдвигается на $\mu$. Распределение остаётся тем же, $z \sim \mathcal N(\mu, \sigma^2)$, но теперь $z$ — дифференцируемая функция $\mu$ и $\sigma$: $\partial z_j / \partial \mu_j = 1$ и $\partial z_j / \partial \sigma_j = \varepsilon_j$. Градиенты слагаемого реконструкции проходят обратно в энкодер (Kingma and Welling, 2014, §2.4). Полученная оценка градиента обычно имеет гораздо меньшую дисперсию, чем универсальная оценка REINFORCE через градиент логарифма плотности (score-function estimator), дисперсию которой Kingma и Welling (§2.2) называют очень высокой.

Рисунок

Зачем нужен трюк репараметризации

KL-часть функции потерь вообще не требует сэмплирования: между двумя гауссианами KL-дивергенция имеет замкнутую форму, выведенную ниже. Суммируя по $d$ латентным измерениям, получаем

$$\mathcal L_{\mathrm{KL}} = -\frac12 \sum_{j=1}^{d} \big(1 + \log \sigma_j^2 - \mu_j^2 - \sigma_j^2\big) .$$

Что даёт упорядоченное латентное пространство#

• Мы можем сэмплировать $z \sim \mathcal N(0, I)$ и декодировать его: слагаемое KL тянет коды к априорному распределению, поэтому сэмплы из него в основном попадают туда, где декодер обучался (насколько хорошо это работает, обсуждается ниже). Можно также сэмплировать рядом с кодом опорного изображения и получать его осмысленные вариации.
• Мы можем интерполировать: декодировать $z_\lambda = (1-\lambda)\, z_a + \lambda\, z_b$ для $\lambda$ от 0 до 1 и плавно переходить от одного изображения или класса к другому — пока отрезок остаётся там, где коды хорошо представлены. Это не гарантировано: коды лишь приближённо согласуются с априорным распределением (ниже), а в высокой размерности прямой отрезок может пересекать области, где нет ни одного обучающего кода.
• Мы можем выполнять арифметику атрибутов. Усредним коды изображений, у которых есть некоторый атрибут, например улыбка, и коды изображений, у которых его нет. Разность двух средних — вектор атрибута, и прибавление его кратного к коду другого изображения добавляет этот атрибут; Carter and Nielsen (2017) делают так с VAE, обученным на шрифтах, и вектором «жирности» (bolding vector). Liu et al. (2019) строят визуальный инструмент, который позволяет задавать такие направления в латентном пространстве и проверять, что они кодируют.

Рисунок

Интерполяция в упорядоченном латентном пространстве

Недостатки.

• Размытые изображения. С факторизованным гауссовским декодером мы показываем его среднее $\hat x(z)$, а не сэмпл из $p_\theta(x \mid z)$. Там, где апостериорные распределения разных изображений перекрываются, один код может объясняться несколькими изображениями, и оптимум квадратичной ошибки $\hat x(z)$ — среднее по ним. Одно из объяснений размытости на уровне всей модели — то, что максимальное правдоподобие (прямая KL, forward KL) покрывает моды (mode-covering), см. заметку Подгонка распределений с помощью KL: модель, которая не может представить данные, размазывает вероятность по всему, что выглядит возможным (Theis, van den Oord and Bethge, 2016).
• Результаты хуже, чем у других моделей, основанных на правдоподобии, на протяжении нескольких лет: Vahdat and Kautz (2020) отмечают, что VAE уступали нормализующим потокам (normalizing flows) и авторегрессионным моделям. Их иерархическая модель NVAE достигла наилучших на тот момент результатов среди неавторегрессионных моделей, основанных на правдоподобии, на CelebA 64 и CelebA HQ и стала первым VAE, применённым к естественным изображениям размером $256 \times 256$ пикселей.
• Нет обусловливания в базовой модели: ей нельзя указать, что генерировать.

На эти недостатки нацелены расширения. Условный VAE (conditional VAE, Sohn, Lee and Yan, 2015) подаёт метку и в энкодер, и в декодер, так что ему можно сказать, что генерировать; он выводится ниже. β-VAE (Higgins et al., 2017) перевзвешивает слагаемое KL, VQ-VAE (van den Oord, Vinyals and Kavukcuoglu, 2017) заменяет гауссов код дискретным, а VAPNEV (Agrawal and Dukkipati, 2016) заменяет попиксельное слагаемое реконструкции точным правдоподобием нормализующего потока, зависящего от кода; всем трём посвящена следующая заметка.

Откуда берётся ELBO#

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

В этом разделе мы опускаем $\theta$ и $\varphi$. Начнём с обоснованности и умножим подынтегральное выражение на $1 = q(z\mid x)/q(z\mid x)$. Такая запись требует, чтобы $q(z \mid x) > 0$ везде, где $p(x, z) > 0$, и для гауссовской $q$ это выполняется:

$$\begin{aligned} \log p(x) &= \log \int p(x, z)\, dz \\ &= \log \int \frac{q(z \mid x)}{q(z \mid x)}\, p(x, z)\, dz \\ &= \log \mathbb E_{q(z \mid x)}\!\left[\frac{p(x, z)}{q(z \mid x)}\right] \\ &\ge \mathbb E_{q(z \mid x)}\!\left[\log \frac{p(x, z)}{q(z \mid x)}\right] = \mathcal L(x). \end{aligned}$$

Неравенство здесь — это неравенство Йенсена. Для выпуклой $f$ оно имеет вид $\mathbb E[f(Y)] \ge f(\mathbb E[Y])$. Логарифм вогнут, поэтому знак меняется на противоположный:

$$\log \mathbb E[Y] \;\ge\; \mathbb E[\log Y].$$

Рисунок

Неравенство Йенсена для вогнутого логарифма

Запись $p(x, z) = p(x \mid z)\, p(z)$ разбивает оценку на два знакомых слагаемых:

$$\begin{aligned} \mathcal L(x) &= \mathbb E_{q(z\mid x)}\!\left[\log p(x \mid z) + \log \frac{p(z)}{q(z \mid x)}\right] \\ &= \mathbb E_{q(z\mid x)}\big[\log p(x \mid z)\big] \\ &\quad - D_{\mathrm{KL}}\big(q(z \mid x)\,\|\,p(z)\big). \end{aligned}$$

Обоснованность никогда не меньше ELBO, поэтому, максимизируя ELBO, мы максимизируем гарантированную нижнюю границу для $\log p(x)$.

Насколько оценка отстаёт от правдоподобия? Ровно на KL-дивергенцию. Поскольку $p(x, z) = p(z \mid x)\, p(x)$,

$$\begin{aligned} &\log p(x) - \mathcal L(x) \\ &\quad= \mathbb E_{q(z\mid x)}\!\left[\log p(x) - \log\frac{p(z\mid x)\,p(x)}{q(z \mid x)}\right] \\ &\quad= \mathbb E_{q(z\mid x)}\big[\log q(z \mid x) - \log p(z \mid x)\big] \\ &\quad= D_{\mathrm{KL}}\big(q(z \mid x)\,\|\,p(z \mid x)\big) \;\ge\; 0 . \end{aligned}$$

Во второй строке $\log p(x)$ сокращается, а на последнем шаге используется неотрицательность KL (см. заметку об энтропии). С параметрами это записывается как $\mathcal L = \log p_\theta(x) - D_{\mathrm{KL}}\big(q_\varphi(z \mid x)\,\|\,p_\theta(z \mid x)\big)$. Шаг по $\theta$ увеличивает эту сумму: обычно за счёт роста $\log p_\theta(x)$, но, возможно, и за счёт того, что истинное апостериорное распределение становится легче приблизить с помощью $q$, а это смещает модель в сторону таких апостериорных распределений. Поэтому максимизация ELBO лишь приближённо максимизирует правдоподобие (Kingma and Welling, 2019, §2.2.1; Hoffman and Johnson, 2016, формула 5). По параметрам энкодера $\log p_\theta(x)$ не меняется, поэтому энкодер может лишь сокращать зазор: он делает $q(z \mid x)$ лучшим приближением истинного апостериорного распределения в смысле обратной KL (reverse KL). Оценка точна ровно тогда, когда $q(z\mid x) = p(z\mid x)$.

Что фиксировано: априорное распределение, а не апостериорное#

Обязано ли апостериорное распределение быть $\mathcal N(0, I)$? Нет. Фиксировано априорное распределение $p(z) = \mathcal N(0, I)$. Апостериорное не фиксировано: для каждого входа энкодер выдаёт два вектора, $\mu(x)$ и $\log \sigma^2(x)$.

Почему $\log \sigma^2$, а не $\sigma^2$?

• $\sigma^2$ — дисперсия и должна быть положительной. Линейный слой может выдать любое вещественное число, а $\sigma^2 = e^{\log \sigma^2} > 0$ всегда.
• Так удобнее численно: крошечная дисперсия вроде $10^{-4}$ превращается в умеренное число, $\log 10^{-4} \approx -9.2$.
• Замкнутая форма KL между гауссианами записывается непосредственно через $\log\sigma^2$.

Постоянное апостериорное распределение бесполезно. Предположим, что энкодер игнорирует вход и выдаёт $q(z \mid x) = \mathcal N(0, I)$ для любого $x$. Слагаемое KL было бы нулевым, но $z$ не нёс бы никакой информации об $x$: декодер получал бы чистый шум. Декодер вроде нашего — факторизованная гауссиана с фиксированной дисперсией, сеть которой выдаёт среднее изображение, — в лучшем случае выдавал бы тогда среднее изображение. Более выразительный декодер, например авторегрессионный, может и без $z$ моделировать нетривиальное распределение изображений, и тогда ничто в ELBO не заставляет использовать код. Этот сбой известен как коллапс апостериорного распределения (posterior collapse), и на практике он случается, когда декодер достаточно мощный, чтобы моделировать данные без кода (Bowman et al., 2016). Вход должен оставаться восстановимым по своему коду.

Замкнутая форма. Для одного латентного измерения

$$\begin{aligned} &D_{\mathrm{KL}}\big(\mathcal N(\mu, \sigma^2)\,\|\,\mathcal N(0, 1)\big) \\ &\quad= \tfrac12\big(\mu^2 + \sigma^2 - 1 - \log \sigma^2\big), \end{aligned}$$

а для диагональной гауссианы вклады измерений складываются (Kingma and Welling, 2014, приложение B). При $q = \mathcal N(\mu, \sigma^2)$ и $p = \mathcal N(0,1)$

$$\begin{aligned} \log q(z) &= -\tfrac12 \log(2\pi\sigma^2) - \frac{(z-\mu)^2}{2\sigma^2}, \\ \log p(z) &= -\tfrac12 \log(2\pi) - \frac{z^2}{2}. \end{aligned}$$

При $z \sim q$ имеем $\mathbb E[(z-\mu)^2] = \sigma^2$ и $\mathbb E[z^2] = \mu^2 + \sigma^2$, поэтому

$$\begin{aligned} D_{\mathrm{KL}}(q\,\|\,p) &= \mathbb E_q[\log q - \log p] \\ &= -\tfrac12\log \sigma^2 - \tfrac12 + \tfrac12(\mu^2 + \sigma^2) \\ &= \tfrac12\big(\mu^2 + \sigma^2 - 1 - \log\sigma^2\big). \end{aligned}$$

Если внести минус в скобку, получится формула для $\mathcal L_{\mathrm{KL}}$ выше. Значение равно нулю только при $\mu = 0$, $\sigma = 1$, потому что $\sigma^2 - 1 - \log\sigma^2 \ge 0$ с равенством только при $\sigma^2 = 1$ (это следует из $\log t \le t - 1$ при $t > 0$ с равенством только при $t = 1$).

В матричном виде. Для $d$-мерного $q = \mathcal N(\mu, \Sigma)$ с положительно определённой ковариационной матрицей $\Sigma$

$$\begin{aligned} &D_{\mathrm{KL}}\big(\mathcal N(\mu, \Sigma)\,\|\,\mathcal N(0, I)\big) \\ &\quad= \tfrac12\big(\operatorname{tr}\Sigma + \mu^\top\mu - d - \log\det\Sigma\big). \end{aligned}$$

Шаги те же, что выше. Логарифмы плотностей равны $\log q(z) = -\tfrac12\log\det(2\pi\Sigma) - \tfrac12(z - \mu)^\top\Sigma^{-1}(z - \mu)$ и $\log p(z) = -\tfrac d2\log(2\pi) - \tfrac12 z^\top z$. При $z \sim q$ математическое ожидание квадратичной формы равно $\mathbb E\big[(z - \mu)^\top\Sigma^{-1}(z - \mu)\big] = \operatorname{tr}\big(\Sigma^{-1}\,\mathbb E[(z - \mu)(z - \mu)^\top]\big) = \operatorname{tr} I = d$, поскольку скаляр равен своему следу и $\operatorname{tr}(AB) = \operatorname{tr}(BA)$. Точно так же $\mathbb E[z^\top z] = \operatorname{tr}\Sigma + \mu^\top\mu$, а слагаемые с $2\pi$ сокращаются. Для диагональной матрицы энкодера $\Sigma = \operatorname{diag}(\sigma^2)$ имеем $\operatorname{tr}\Sigma = \sum_j \sigma_j^2$ и $\log\det\Sigma = \sum_j \log\sigma_j^2$, так что формула распадается на сумму одномерных слагаемых. Например, $\mu = (1, 0)$ и $\sigma = (1, 0.5)$ дают $0.5 + 0.318 = 0.818$ ната (nats) — сумму второй и четвёртой строк таблицы ниже. Полная ковариационная матрица позволила бы координатам кода ещё и коррелировать, и у корреляции своя цена: при $\mu = 0$, единичных дисперсиях и коэффициенте корреляции $\rho$ имеем $\det\Sigma = 1 - \rho^2$, а KL равна $-\tfrac12\log(1 - \rho^2)$, около $0.144$ ната при $\rho = 0.5$. Диагональное апостериорное распределение VAE такие корреляции выразить вообще не может.

Как это читать: модель платит штраф, если $\mu \ne 0$ или $\sigma \ne 1$, и у штрафа вполне определённая форма.

• Сдвиг среднего стоит $\tfrac12\mu^2$: 0.5 при $\mu = 1$ и 2 при $\mu = 2$.
• Уменьшение дисперсии стоит $\tfrac12\sigma^2 - \log\sigma - \tfrac12 \approx -\log\sigma - \tfrac12$ при малых $\sigma$, поэтому штраф растёт лишь логарифмически: точные значения — 1.81 при $\sigma = 0.1$ и 4.11 при $\sigma = 0.01$. При $\sigma \to 0$ штраф неограниченно растёт, так что детерминированный точечный код, как в обычном автоэнкодере, заплатил бы бесконечный штраф. Точечная масса не абсолютно непрерывна относительно гауссианы, поэтому её KL бесконечна; см. случай «атом против плотности» в заметке Подгонка распределений с помощью KL.
• Увеличение дисперсии стоит $\tfrac12\sigma^2 - \log\sigma - \tfrac12$, что растёт как $\tfrac12\sigma^2$.

Поэтому каждый вход может позволить себе собственную довольно узкую гауссиану со смещённым центром, когда это помогает реконструкции, но никогда — одну точку.

Рисунок

Штраф KL растёт лишь логарифмически при σ → 0

$\mu$	$\sigma$	$D_{\mathrm{KL}}$ (наты, nats)
0	1	0
1	1	0.5
2	1	2
0	0.5	0.318
0	2	0.807
0	0.1	1.808
0	0.01	4.105

Какое распределение должно совпадать с априорным? Слагаемое KL тянет каждое $q(z \mid x)$ к априорному распределению лишь настолько, насколько позволяет реконструкция. Чтобы увидеть, чего оно требует от кодов в целом, возьмём $N$ обучающих изображений $x_1, \dots, x_N$ и определим агрегированное апостериорное распределение (aggregate posterior) — смесь всех их апостериорных распределений:

$$q(z) = \frac1N \sum_{n=1}^{N} q(z \mid x_n) .$$

Будем считать индекс $n$ случайной величиной, равномерно распределённой на $\lbrace 1, \dots, N\rbrace$, так что $q(n, z) = \tfrac1N\, q(z \mid x_n)$. Тогда среднее слагаемое KL в ELBO распадается на две части (Hoffman and Johnson, 2016, формулы 11–12):

$$\begin{aligned} &\frac1N \sum_{n=1}^{N} D_{\mathrm{KL}}\big(q(z \mid x_n)\,\|\,p(z)\big) \\ &\quad= I_q(n; z) + D_{\mathrm{KL}}\big(q(z)\,\|\,p(z)\big). \end{aligned}$$

Здесь $I_q(n; z) = \mathbb E_{q(n, z)}\big[\log \frac{q(n \mid z)}{q(n)}\big]$ — взаимная информация (mutual information) между индексом изображения и его кодом; она лежит между 0 и $\log N$. Второе слагаемое — маргинальная KL (marginal KL) между агрегированным апостериорным распределением и априорным. Чтобы проверить тождество, запишем по теореме Байеса для $q(n, z)$ отношение $q(z \mid x_n)/p(z) = \big(q(z)/p(z)\big)\cdot\big(q(n \mid z)/q(n)\big)$ и возьмём математическое ожидание его логарифма по $q(n, z)$: первый множитель даёт маргинальную KL, а второй — взаимную информацию. Хорошая реконструкция требует, чтобы код идентифицировал изображение, поэтому это слагаемое должно оставаться большим, и именно поэтому каждое отдельное апостериорное распределение узкое. Чего-то от априорного распределения требует только второе слагаемое: смесь $q(z)$ должна заполнить $\mathcal N(0, I)$, не оставляя дыр. Обучение подталкивает смесь к априорному распределению, но совпадение получается лишь приближённым. Для VAE с 10 латентными измерениями на бинаризованном MNIST Hoffman и Johnson получают взаимную информацию, близкую к её максимуму $\log N$, и маргинальную KL $D_{\mathrm{KL}}(q(z)\,\|\,p(z))$, равную 8.35 ната (§4, таблица 1). Там, где $q(z)$ оставляет дыры, сэмплы из априорного распределения попадают в области, которые декодер видел редко, и декодируются плохо.

Метка как условие: условный VAE#

Базовому VAE нельзя указать, что генерировать. Условный VAE (CVAE) добавляет наблюдаемое условие $c$, например метку класса или вектор бинарных атрибутов, в каждую часть модели (Sohn, Lee and Yan, 2015, раздел 4):

• априорное распределение $p_\theta(z \mid c)$, часто просто фиксированное $\mathcal N(0, I)$ при любом $c$;
• декодер $p_\theta(x \mid z, c)$;
• энкодер $q_\varphi(z \mid x, c) = \mathcal N\big(\mu_\varphi(x, c), \sigma_\varphi^2(x, c)\big)$.

Распределение изображений модели при заданном условии — $p_\theta(x \mid c) = \int p_\theta(x \mid z, c)\, p_\theta(z \mid c)\, dz$. Вывод ELBO проходит без изменений, если в условие каждого распределения добавить $c$ (Sohn et al., формула 4):

$$\begin{aligned} \log p_\theta(x \mid c) \ge{}& \mathbb E_{q_\varphi(z \mid x, c)}\big[\log p_\theta(x \mid z, c)\big] \\ &- D_{\mathrm{KL}}\big(q_\varphi(z \mid x, c)\,\|\,p_\theta(z \mid c)\big), \end{aligned}$$

и зазор снова равен KL-дивергенции $D_{\mathrm{KL}}\big(q_\varphi(z \mid x, c)\,\|\,p_\theta(z \mid x, c)\big)$. Обучение максимизирует эту оценку, усреднённую по парам $(x, c)$ из данных, с той же репараметризацией и, для фиксированного априорного распределения $\mathcal N(0, I)$, с той же KL в замкнутой форме. Обучаемому гауссовскому априорному распределению $p_\theta(z \mid c)$ нужна общая формула KL между двумя гауссианами; для диагональных это одномерная формула, просуммированная по координатам. В сетях $c$ обычно подаётся конкатенацией: ко входу энкодера или к его признакам и к $z$ перед декодером.

Что меняет условие. Декодер теперь получает $c$ даром, а за каждый нат, который несёт код, приходится платить KL. Поэтому целевая функция поощряет не включать в $z$ то, что уже сообщает $c$, и код склонен хранить то, чего метка не описывает, например позу или освещение. Чтобы сгенерировать изображение, фиксируем $c$, берём $z \sim p_\theta(z \mid c)$ и декодируем. Если сохранить $z$ и изменить $c$, например один из 40 бинарных атрибутов лиц CelebA, мы запрашиваем то же изображение с другим значением атрибута; насколько чисто это получается, зависит от того, какая часть атрибута просочилась в $z$. Та же конструкция добавляет условие в любую модель этого семейства, включая декодер-поток условного VAPNEV, а условный GAN — её состязательный аналог.

Итоги#

• Правдоподобие $p_\theta(x)$ и апостериорное распределение $p_\theta(z \mid x)$ модели с латентными переменными вычислительно неразрешимы, потому что оба требуют интегрирования по латентному пространству.
• Энкодер строит приближённое апостериорное распределение $q_\varphi(z \mid x)$ — диагональную гауссиану.
• Обучение максимизирует ELBO, то есть реконструкцию минус KL-дивергенцию относительно априорного распределения; это нижняя оценка $\log p_\theta(x)$ с зазором $D_{\mathrm{KL}}(q\,\|\,p(z\mid x))$.
• Трюк репараметризации делает сэмплируемое слагаемое реконструкции дифференцируемым; слагаемое KL известно в замкнутой форме.
• Фиксированное априорное распределение упорядочивает латентное пространство, поэтому декодирование $z \sim \mathcal N(0, I)$ порождает новые изображения. Агрегированное апостериорное распределение совпадает с априорным лишь приближённо, и дыры в нём дают плохие сэмплы.
• Условный VAE добавляет метку $c$ в условие априорного распределения, энкодера и декодера и максимизирует ту же оценку, теперь для $\log p_\theta(x \mid c)$; ему можно указать, что генерировать.

Ноутбук#

Реализация следует ноутбуку о VAE из репозитория HSE-LAMBDA Deep Generative Models (закреплённый коммит; ячейки нумеруются с 1, считая все ячейки подряд). В ноутбуке нет сохранённых выходов, поэтому он документирует процедуру, а не результаты.

• Ячейки 4–5: лица CelebA, приведённые к размеру $64 \times 64$ RGB и нормализованные в $[-1, 1]$, размер батча 128, drop_last=True. CelebA содержит 202 599 выровненных и обрезанных лиц размером 178 × 218 пикселей (ширина × высота), у каждого 40 бинарных атрибутов; класс датасета возвращает атрибуты, а этот ноутбук их не использует. Resize((64, 64)) масштабирует две стороны независимо и меняет соотношение сторон с 178 : 218 на 1 : 1; Resize(64) с последующим CenterCrop(64) сохранили бы его.
• Ячейки 6–7: почему обычные автоэнкодеры не могут генерировать, и рассуждение через теорему Байеса, которое приводит к приближённому апостериорному распределению.
• Ячейки 10 и 12: в энкодере четыре свёртки $4 \times 4$ с шагом 2 и числом каналов $32 \to 64 \to 128 \to 256$, за каждой следуют LeakyReLU и затем BatchNorm (кроме последней), так что изображение $64 \times 64$ превращается в карту $4 \times 4 \times 256$; она вытягивается в 4096 признаков для двух линейных выходных слоёв (heads) mean_layer и logvar_layer по 100 выходов в каждом. Декодер не является буквальным зеркалом энкодера: линейный слой из 100 в 4096 нейронов с ReLU, изменение формы на $256 \times 4 \times 4$, затем четыре блока, каждый из которых состоит из повышения разрешения ($\times 2$) и свёртки $3 \times 3$ ($256 \to 128 \to 64 \to 32 \to 3$ каналов), на выходе Tanh со значениями в $[-1, 1]$. nn.Upsample(scale_factor=2) без аргумента mode повышает разрешение методом ближайшего соседа.
• Ячейка 15: класс VAE с latent_dim = 100. reparametrize вычисляет std_dev = exp(0.5 * logvar), берёт epsilon = randn_like(std_dev) и возвращает mean + epsilon * std_dev; sample(n) декодирует $z \sim \mathcal N(0, I)$.
• Ячейки 16–17: loss_function: квадратичная ошибка, просуммированная по пикселям и усреднённая по батчу, плюс kld_loss = -0.5 * sum(1 + logvar - mean² - exp(logvar)), усреднённая по батчу и умноженная на beta (по умолчанию 1). Это оценка отрицательной ELBO по одному сэмплу с $s^2 = \tfrac12$ из абзаца о реконструкции с точностью до константы. exp(logvar) внутри KL ограничена снизу значением $10^{-8}$. Стабилизирующего эффекта это ограничение не даёт: оно меняет функцию потерь, только когда $\sigma^2 < 10^{-8}$, а само слагаемое $\log\sigma^2$ не ограничено. Текстовая ячейка 16 правильно приводит KL в матричном виде, $\tfrac12(\operatorname{tr}\Sigma + \mu^\top\mu - d - \log\det\Sigma)$, и её диагональный случай, но затем переписывает $\sum_i \log\sigma_i^2$ как $\lVert\log\boldsymbol\sigma\rVert_1$, что неверно: $\sum_i \log\sigma_i^2 = 2\sum_i \log\sigma_i$, а логарифмы могут быть отрицательными, поэтому в норме не сохраняются ни множитель 2, ни знак. Код в ячейке 17 использует правильную сумму.
• Ячейки 18–21: обучение с Adam при шаге обучения (learning rate) $10^{-3}$ в течение 10 эпох с косинусным расписанием шага обучения (CosineAnnealingLR, T_max = 10); два слагаемых функции потерь логируются по отдельности, а после каждой эпохи декодируются 64 сэмпла из априорного распределения.
• Ячейка 23: генерация из априорного распределения. Ячейка 24 (только текст) определяет функцию потерь β-VAE и распутывание на примере CelebA, где одна латентная координата отвечает за улыбку, а другая — за поворот головы (см. что даёт β); β-VAE не обучается, потому что train_vae вызывает loss_function со значением по умолчанию $\beta = 1$.
• Замечания о реализации. Логируемые и печатаемые значения функции потерь делятся на размер батча второй раз (mse_loss.item() / len(data), а суммы потерь на изображение за эпоху делятся на len(dataset)), поэтому они примерно в 128 раз меньше функции потерь на одно изображение. Модель ни разу не переключается в eval(): сэмплирование в ячейках 18 и 23 запускает BatchNorm декодера на статистиках сгенерированного батча и обновляет его скользящие средние.

Ноутбук с упражнениями 03_hw_vapnev.ipynb из того же репозитория продолжает эту модель. Ячейки 3–9 загружают CelebA: ячейка 5 применяет Resize(64) и CenterCrop(64) и переводит пиксели в $[-1, 1]$ преобразованием $2x - 1$, а ячейка 9 показывает 25 случайных лиц. Остальные ячейки формулируют задания и оставляют ячейки кода пустыми: ячейки 12–13 требуют обучить VAE на этих лицах и оценить его по FID; ячейки 14–15 требуют построить VAPNEV и его условную версию, описанные в следующей заметке; условная версия использует атрибуты как условие так же, как условный VAE. Единственные сохранённые выходы ноутбука — устройство, форма одного изображения и сетка лиц, так что цитировать нечего.

Дальше#

VAE обучает генератор через явную оценку правдоподобия и расплачивается за это размытостью и априорным распределением, которое его коды заполняют лишь приблизительно. Следующая заметка меняет части оценки: вес слагаемого KL, задающий, сколько может сообщать код, дискретный код, априорное распределение которого обучается после автоэнкодера, и точное правдоподобие потока вместо попиксельного слагаемого реконструкции.

Литература#

• D. P. Kingma, M. Welling. Auto-encoding variational Bayes. ICLR, 2014.
• D. J. Rezende, S. Mohamed, D. Wierstra. Stochastic backpropagation and approximate inference in deep generative models. ICML, 2014.
• D. P. Kingma, M. Welling. An introduction to variational autoencoders. Foundations and Trends in Machine Learning, 12(4), 2019.
• D. M. Blei, A. Kucukelbir, J. D. McAuliffe. Variational inference: A review for statisticians. Journal of the American Statistical Association, 112(518), 2017.
• M. D. Hoffman, M. J. Johnson. ELBO surgery: yet another way to carve up the variational evidence lower bound. NIPS Workshop on Advances in Approximate Bayesian Inference, 2016 (архивная копия).
• S. R. Bowman, L. Vilnis, O. Vinyals, A. M. Dai, R. Jozefowicz, S. Bengio. Generating sentences from a continuous space. CoNLL, 2016.
• K. Sohn, H. Lee, X. Yan. Learning structured output representation using deep conditional generative models. NIPS, 2015. Раздел 4, формула 4: условная ELBO.
• S. Agrawal, A. Dukkipati. Deep variational inference without pixel-wise reconstruction. arXiv:1611.05209, 2016. VAPNEV.
• S. Carter, M. Nielsen. Using artificial intelligence to augment human intelligence. Distill, 2017. Векторы атрибутов в латентном пространстве VAE.
• Y. Liu, E. Jun, Q. Li, J. Heer. Latent space cartography: Visual analysis of vector space embeddings. Computer Graphics Forum, 38(3), 2019.
• I. Higgins et al. β-VAE: Learning basic visual concepts with a constrained variational framework. ICLR, 2017.
• A. van den Oord, O. Vinyals, K. Kavukcuoglu. Neural discrete representation learning. NIPS, 2017.
• L. Theis, A. van den Oord, M. Bethge. A note on the evaluation of generative models. ICLR, 2016.
• A. Vahdat, J. Kautz. NVAE: A deep hierarchical variational autoencoder. NeurIPS, 2020.
• HSE-LAMBDA, репозиторий курса Deep Generative Models: VAE.ipynb.