Варианты VAE: β-VAE и VQ-VAE

VAE закончился двумя претензиями. Его сэмплы размыты, а фиксированное гауссово априорное распределение лишь приблизительно заполняется кодами, которые на самом деле выдаёт энкодер, так что сэмплы из априорного распределения могут попасть туда, где декодер никогда не обучался. Эта заметка разбирает три ответа. Первый сохраняет модель и меняет целевую функцию: вес $\beta$ при слагаемом KL решает, сколько информации может нести код. Второй сохраняет автоэнкодер и меняет код: сетка дискретных символов из обучаемого словаря, а априорное распределение подгоняется потом моделью, достаточно сильной, чтобы им соответствовать. Третий сохраняет гауссов код и меняет правдоподобие декодера: точное правдоподобие, вычисленное нормализующим потоком (normalizing flow), вместо квадратичной ошибки, просуммированной по пикселям.

Обозначения. Как в заметке о VAE: изображение $x$, код $z$, энкодер $q_\varphi(z \mid x)$, декодер $p_\theta(x \mid z)$, априорное распределение $p(z)$. $\mathrm{sg}[\cdot]$ — оператор остановки градиента (stop-gradient): при прямом проходе возвращает аргумент, а при обратном имеет нулевую производную. Словарь кодов (codebook) — набор из $K$ векторов $e_1, \dots, e_K \in \mathbb R^d$.

Три способа изменить VAE#

ELBO одного изображения состоит из двух слагаемых:

$$\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{KL к априорному}} .$$

Первое хочет, чтобы код сообщал об $x$ как можно больше; второе берёт плату за каждый сообщённый бит. Изменить их баланс — это β-VAE. Изменить то, чем является код, — непрерывным гауссовым или одним символом из $K$, — это VQ-VAE. Одно не зависит от другого: первое двигает точку вдоль компромисса, второе меняет то, где находится трудность генерации. Третье изменение, дополнительное, в конце заметки, сохраняет оба слагаемых и заменяет вид самого слагаемого реконструкции, $p_\theta(x \mid z)$.

β-VAE: вес при слагаемом KL#

β-VAE (Higgins et al., 2017) максимизирует

$$\mathcal L_\beta(x) = \mathbb E_{q_\varphi(z \mid x)}\big[\log p_\theta(x \mid z)\big] - \beta\, D_{\mathrm{KL}}\big(q_\varphi(z \mid x)\,\|\,p(z)\big) .$$

$\beta = 1$ — обычный VAE. При $\beta > 1$ целевая функция по-прежнему нижняя оценка $\log p_\theta(x)$, потому что вычитает неотрицательной величины больше, чем ELBO; при $\beta < 1$ она может превысить $\log p_\theta(x)$ и оценкой уже не является. В обоих случаях $\beta$ меняет то, какую модель предпочитает обучение, а не только масштаб функции потерь.

Скорость и искажение#

После усреднения по данным у двух слагаемых появляются имена из теории информации (Alemi et al., 2018). Скорость (rate) — средняя KL,

$$R = \mathbb E_{x}\, D_{\mathrm{KL}}\big(q_\varphi(z \mid x)\,\|\,p(z)\big),$$

измеряемая в натах (nats). Она ограничивает сверху взаимную информацию между изображениями и их кодами при данном энкодере: тождество, на котором это держится, — скорость равна взаимной информации плюс KL от агрегированного апостериорного распределения к априорному, — выведено в заметке о VAE. Искажение (distortion) $D = -\mathbb E_x \mathbb E_{q_\varphi(z \mid x)} \log p_\theta(x \mid z)$ — средняя функция потерь реконструкции. Отрицательный ELBO равен $D + R$, а β-VAE минимизирует $D + \beta R$ — лагранжиан задачи «реконструировать как можно лучше при скорости не выше заданной». Каждое $\beta$ выбирает одну точку компромисса; $\beta = 1$ — та, что лучше всего оценивает правдоподобие.

Разобранный пример: одна латентная координата#

Чтобы увидеть компромисс в числах, возьмём одну латентную координату, стандартное нормальное априорное распределение $p(z) = \mathcal N(0, 1)$ и фиксированный линейный декодер $p(x \mid z) = \mathcal N(w z, \sigma^2)$ с $w = 1$ и $\sigma^2 = 0.1$. Данные порождены этой же моделью, так что $x \sim \mathcal N(0, w^2 + \sigma^2) = \mathcal N(0, 1.1)$. Фиксированный декодер позволяет увидеть, что $\beta$ делает именно с энкодером. Пусть энкодер — $q(z \mid x) = \mathcal N(m, v)$. Для одного $x$ β-целевая функция с точностью до констант равна

$$J(m, v) = -\frac{(x - w m)^2 + w^2 v}{2\sigma^2} - \frac{\beta}{2}\big(v + m^2 - 1 - \ln v\big),$$

где использованы $\mathbb E_q (x - w z)^2 = (x - wm)^2 + w^2 v$ и KL к стандартному нормальному распределению в замкнутой форме из заметки о VAE. Приравнивая производные к нулю,

$$\frac{\partial J}{\partial v} = -\frac{w^2}{2\sigma^2} - \frac{\beta}{2} + \frac{\beta}{2v} = 0, \qquad \frac{\partial J}{\partial m} = \frac{w(x - wm)}{\sigma^2} - \beta m = 0,$$

получаем оптимальный энкодер

$$m = \frac{w}{w^2 + \beta\sigma^2}\, x, \qquad v = \frac{\beta\sigma^2}{w^2 + \beta\sigma^2} .$$

При $\beta = 1$ это среднее и дисперсия точного апостериорного распределения $p(z \mid x)$, так что ELBO точен. Усреднение по $x \sim \mathcal N(0, 1.1)$ даёт скорость и искажение в виде квадратичной ошибки $\mathbb E(x - wz)^2$:

$\beta$	$m / x$	$v$	скорость $R$, наты	искажение $\mathbb E(x - wz)^2$
0.25	0.976	0.024	1.89	0.025
1	0.909	0.091	1.20	0.100
4	0.714	0.286	0.55	0.376
16	0.385	0.615	0.13	1.032

Рост $\beta$ стягивает код к априорному распределению: среднее тянется к 0, а дисперсия растёт к 1, так что код сообщает об $x$ меньше, и реконструкция ухудшается. При $\beta = 16$ искажение 1.03 близко к дисперсии данных 1.1: код почти ничего не несёт. Уменьшение $\beta$ делает обратное и покупает более чёткие реконструкции ценой большей скорости.

Высокая скорость сама по себе не означает, что код далёк от априорного распределения. Скорость раскладывается на взаимную информацию $I(X; Z)$ между изображениями и кодами и KL от агрегированного апостериорного распределения $q(z) = \mathbb E_x\, q(z \mid x)$ к априорному. Здесь $q(z) = \mathcal N\big(0, (m/x)^2 \cdot 1.1 + v\big)$: при $\beta = 0.25$ его дисперсия равна 1.071, а рассогласование $D_{\mathrm{KL}}\big(q(z)\,\|\,p(z)\big)$ составляет 0.001 из 1.89 ната, так что почти вся скорость — информация об $x$. Сэмплы из априорного распределения плохи, когда агрегированное апостериорное распределение уходит от априорного или когда декодер ошибается там, куда априорное распределение кладёт массу. В этом игрушечном примере декодер зафиксирован и совпадает с истинной моделью, поэтому $\beta$ меняет только энкодер, а сэмплы, декодированные из априорного распределения, остаются точными. В VAE, где всё обучается совместно, малое $\beta$ ослабляет и штраф за рассогласование — так оно и может ухудшить сэмплы из априорного распределения.

Рисунок

β-VAE передвигает точку по кривой «скорость–искажение»

Что даёт β, а что нет#

Higgins et al. (2017) сообщают, что при $\beta > 1$ отдельные латентные координаты начинают следить за отдельными факторами изменчивости (положение, масштаб, поворот спрайта), и называют это распутыванием (disentanglement). Распутанное представление — такое, в котором отдельная латентная координата чувствительна к изменениям одного порождающего фактора и относительно инвариантна к изменениям остальных (раздел 1). Важны обе половины определения: координата должна следовать за своим фактором и мало меняться, когда меняется другой фактор. На лицах CelebA при $\beta = 250$ отдельные координаты поворачивают голову, меняют выражение лица в сторону улыбки или меняют причёску, а остальное лицо остаётся прежним; в обычном VAE ($\beta = 1$) координата, которая поворачивает голову, меняет заодно выражение лица, очки и пол (рисунок 1). Такой код позволил бы редактировать по одному признаку за раз, переиспользовать факторы в новых задачах, рассуждать о не встречавшихся сочетаниях известных факторов и замечать входы, которые не подходят ни под одно известное сочетание (раздел 1).

Механизм — давление тратить мало скорости: факторизованное априорное распределение дешевле всего подогнать апостериорным с независимыми координатами. $\beta < 1$ ослабляет это давление: такое значение покупает более чёткие реконструкции, и целевая функция поощряет распутанный код ещё меньше. Эффект $\beta > 1$ тоже не гарантирован. Locatello et al. (2019) доказывают, что без учителя или индуктивных смещений (inductive biases) распутывание нельзя идентифицировать по данным, и в исследовании более чем 12 000 обученных моделей находят, что начальные значения генератора случайных чисел (seeds) и гиперпараметры значат больше, чем выбор между распутывающими целевыми функциями. Burgess et al. (2018) заменяют фиксированный вес на целевую скорость $C$, которая растёт по ходу обучения, $\mathbb E_q \log p_\theta(x \mid z) - \gamma\,\lvert D_{\mathrm{KL}} - C \rvert$, так что коду разрешается по одному нату больше за раз. Полезное прочтение β-VAE — то, что выше: ручка регулировки скорости, а не выключатель интерпретируемости.

VQ-VAE: дискретный латентный код#

VQ-VAE (van den Oord, Vinyals and Kavukcuoglu, 2017) отказывается от непрерывного кода. Энкодер по-прежнему выдаёт карту векторов $z_e(x) \in \mathbb R^{h \times w \times d}$, но перед подачей в декодер каждый вектор заменяется ближайшим элементом обучаемого словаря:

$$k^*(i) = \arg\min_k \big\lVert z_e(x)_i - e_k \big\rVert^2, \qquad z_q(x)_i = e_{k^*(i)} .$$

Код изображения — сетка индексов $k^*(i)$: для изображений $32 \times 32$ ниже это 64 целых числа от 1 до 512 на нижнем уровне. В терминах ELBO апостериорное распределение $q(z = k \mid x)$ — единица в $k^*$, и при равномерном априорном распределении на $K$ кодах его KL равна $\log K$ на позицию, то есть константе. Слагаемое KL выпадает из обучения, и VQ-VAE обучается как обычный автоэнкодер с квантованным узким местом.

Рисунок

Векторное квантование: ближайший вектор словаря

Обучение через argmin#

У argmin нет полезного градиента: малое изменение $z_e$ либо оставляет тот же $e_k$, либо перескакивает на другой. VQ-VAE использует прямой проброс градиента (straight-through estimator, Bengio, Léonard and Courville, 2013): при прямом проходе декодер получает $z_q$, при обратном градиент по $z_q$ без изменений копируется в $z_e$. В коде это одна строка, z_q = z_e + (z_q - z_e).detach(): значение — $z_q$, производная — как у $z_e$. Копия — заменитель, а не производная квантования. Внутри ячейки квантователь постоянен, поэтому его настоящий якобиан равен нулю; на границе ячейки он скачет. Прямой проброс подставляет вместо этого якобиана единичную матрицу, и это остаётся подстановкой даже тогда, когда $z_e$ лежит точно на $e_{k^*}$: настоящая производная там всё равно нулевая. Близость лишь делает эвристику правдоподобной: градиент декодера, вычисленный в $e_{k^*}$, тогда разумно подсказывает, куда сдвигать $z_e$. В функции потерь три слагаемых:

$$\mathcal L = -\log p_\theta\big(x \mid z_q(x)\big) + \big\lVert \mathrm{sg}[z_e(x)] - e \big\rVert^2 + \beta_c \big\lVert z_e(x) - \mathrm{sg}[e] \big\rVert^2 ,$$

просуммированных по позициям, где $e = e_{k^*}$. Слагаемое реконструкции обучает декодер и, через прямой проброс, энкодер. Слагаемое словаря двигает $e$ к выходу энкодера; энкодер оно не трогает. Слагаемое привязки (commitment loss) двигает выход энкодера к $e$; словарь оно не трогает. Без него ничто не мешает выходам энкодера уходить от словаря быстрее, чем словарь успевает за ними. В статье $\beta_c = 0.25$, и результаты к нему нечувствительны.

Разобранный пример: один шаг квантования#

Возьмём $d = 2$, $K = 3$, словарь $e_1 = (0, 0)$, $e_2 = (2, 0)$, $e_3 = (0, 2)$ и один выход энкодера $z_e = (1.4, 0.3)$ — красную точку на рисунке. Квадраты расстояний:

$$\lVert z_e - e_1\rVert^2 = 1.96 + 0.09 = 2.05, \quad \lVert z_e - e_2\rVert^2 = 0.36 + 0.09 = 0.45, \quad \lVert z_e - e_3\rVert^2 = 1.96 + 2.89 = 4.85 .$$

Ближайший вектор — $e_2$, поэтому декодер получает $z_q = (2, 0)$, а код — индекс 2. Слагаемое словаря равно $0.45$, слагаемое привязки — $0.25 \cdot 0.45 = 0.1125$: то же расстояние, но градиенты направлены в противоположные стороны. Если градиент функции потерь декодера по $z_q$ равен $g$, то энкодер получает в $z_e$ ровно $g$ плюс градиент привязки $2\beta_c(z_e - e_2) = 0.5 \cdot (-0.6, 0.3) = (-0.3, 0.15)$, направленный от $e_2$ к $z_e$, так что шаг спуска сдвигает $z_e$ к $e_2$.

Словарь как онлайн-k-means#

Слагаемое словаря двигает каждый $e_k$ к выходам энкодера, отнесённым к нему, — это шаг k-means. Приложение статьи о VQ-VAE заменяет его экспоненциальными скользящими средними. Для каждого кода хранятся счётчик $N_k$ и сумма $m_k$, которые обновляются по $n_k$ выходам, отнесённым к коду $k$ в батче,

$$N_k \leftarrow \gamma N_k + (1 - \gamma)\, n_k, \qquad m_k \leftarrow \gamma\, m_k + (1 - \gamma) \sum_{i:\, k^*(i) = k} z_e(x)_i, \qquad e_k = \frac{m_k}{N_k},$$

с $\gamma = 0.99$, а слагаемое словаря убирается из функции потерь. Каждый $e_k$ тогда — скользящее среднее своей ячейки, то есть центроид k-means, вычисляемый онлайн. При $N_2 = 10$, $m_2 = (20, 0)$ (так что $e_2 = (2, 0)$) и батче из двух выходов $(1.4, 0.3)$ и $(1.8, -0.1)$, отнесённых к нему, обновление даёт $N_2 = 9.92$, $m_2 = (19.832, 0.002)$ и $e_2 = (1.9992, 0.0002)$: маленький шаг к среднему батча $(1.6, 0.1)$, как и положено скользящему среднему с $\gamma = 0.99$.

Слабое место этой схемы — то же, что у k-means: код, который не выбирает ни один выход, не получает обновлений и остаётся там, где был инициализирован, — мёртвый код (dead code). При многих мёртвых кодах фактический словарь мал. Практические исправления перезапускают неиспользуемые коды в недавних выходах энкодера, как в Jukebox (Dhariwal et al., 2020), или ищут коды в низкоразмерном нормированном пространстве, как в ViT-VQGAN (Yu et al., 2022). Реализации часто добавляют к счётчикам малую константу (сглаживание Лапласа), чтобы не делить на ноль для неиспользуемых кодов; оживить их это не помогает.

Куда делось априорное распределение#

Поскольку слагаемое KL постоянно, обучение оставляет априорное распределение равномерным, и независимое равномерное сэмплирование кодов дало бы изображение из случайных кусков. Поэтому VQ-VAE подгоняет априорное распределение после автоэнкодера: обучающая выборка кодируется в сетки индексов, и на них обучается авторегрессионная модель — в статье PixelCNN, — предсказывающая каждый индекс по предыдущим. Генерация сэмплирует сетку из этого распределения и декодирует её. Двухэтапная структура — суть метода. Автоэнкодер отвечает за пиксели, а априорной модели нужно описать сетку в десятки раз меньше, состоящую из символов, — там авторегрессионные модели сильны; как они устроены, разбирает заметка об авторегрессионных моделях.

VQ-VAE-2: два уровня кодов#

Одна сетка должна нести и общую компоновку, и мелкую текстуру. VQ-VAE-2 (Razavi, van den Oord and Vinyals, 2019) использует две: грубую верхнюю для компоновки и более мелкую нижнюю для деталей, которая квантуется, уже видя верхнюю. Для ImageNet $256 \times 256$ сетки имеют размер $32 \times 32$ и $64 \times 64$ по 512 кодов; для FFHQ $1024 \times 1024$ уровней три. Априорное распределение тоже иерархическое: авторегрессионная модель по верхним кодам, в которой каузальное самовнимание чередуется с гейтированными свёртками, и условная модель по нижним кодам при данных верхних — без внимания, чтобы большая сетка оставалась посильной. Статья к тому же отбирает сэмплы предобученным классификатором ImageNet — шаг отбраковки, который меняет разнообразие на качество.

Рисунок

Небольшой VQ-VAE-2: две карты кодов

Сколько битов#

Рисунок показывает небольшую версию для изображений CIFAR-10 $32 \times 32$: верхняя сетка $4 \times 4$ и нижняя $8 \times 8$, по 512 кодов в каждой. Изображение занимает $32 \cdot 32 \cdot 3 = 3072$ байта, или 24 576 бит. Его код — $16 + 64 = 80$ индексов по $\log_2 512 = 9$ бит, то есть 720 бит, примерно в 34 раза меньше. Стоимость сэмплирования масштабируется так же: авторегрессионной модели по сырым субпикселям нужны 3072 последовательных шага на изображение, двум априорным моделям — $16 + 64 = 80$. В масштабе статьи, $256 \times 256$, разрыв — $196\,608$ субпикселей против $32^2 + 64^2 = 5120$ кодов.

VAPNEV: правдоподобие потока вместо попиксельной реконструкции#

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

Изменения выше не трогают правдоподобие декодера. С факторизованным гауссовским декодером $-\log p_\theta(x \mid z)$ — это квадратичная ошибка, просуммированная по пикселям, и лучший выход декодера для кода — среднее всех изображений, которые ему подходят: отсюда размытость VAE. VAPNEV (Agrawal and Dukkipati, 2016), анаграмма VAE-NVP, сохраняет гауссов код и ELBO и заменяет это слагаемое точным правдоподобием, которое вычисляется зависящим от кода потоком real NVP (Dinh, Sohl-Dickstein and Bengio, 2017). Подробно потоки разбираются позже в курсе, в заметке о нормализующих потоках; этому разделу нужны только формула замены переменных и слой связи (coupling layer), и обе вещи мы здесь повторяем.

Обозначения. $y \in \mathbb R^D$ — промежуточная переменная того же размера, что и изображение. Для каждого кода $z$ отображение $f_z: y \mapsto x$ обратимо; это поток, записанный в принятом в курсе направлении из латентного пространства в данные, с обратным отображением $f_z^{-1}: x \mapsto y$. Agrawal и Dukkipati записывают отображение в направлении из данных в латентное пространство.

Правдоподобие декодера как замена переменных#

Сеть-декодер не выдаёт изображение. Она выдаёт среднее и дисперсии гауссианы над $y$, а поток превращает $y$ в изображение:

$$y \mid z \sim \mathcal N\big(\mu_\theta(z), \sigma_\theta^2(z)\big), \qquad x = f_z(y).$$

По формуле замены переменных правдоподобие изображения вычисляется точно (раздел 3.1, формулы 10–11):

$$\begin{aligned} \log p_\theta(x \mid z) ={}& \log\mathcal N\big(f_z^{-1}(x);\, \mu_\theta(z), \sigma_\theta^2(z)\big) \\ &+ \log\left|\det\frac{\partial f_z^{-1}(x)}{\partial x}\right| . \end{aligned}$$

ELBO сохраняет свой вид, только вместо квадратичной ошибки в нём стоит это слагаемое:

$$\begin{aligned} \mathcal L(x) ={}& \mathbb E_{q_\varphi(z \mid x)}\Big[\log\mathcal N\big(f_z^{-1}(x);\, \mu_\theta(z), \sigma_\theta^2(z)\big) \\ &\qquad + \log\left|\det\frac{\partial f_z^{-1}(x)}{\partial x}\right|\Big] \\ &- D_{\mathrm{KL}}\big(q_\varphi(z \mid x)\,\|\,p(z)\big). \end{aligned}$$

Гауссиана теперь живёт в пространстве $y$, а не пикселей. Поток перемешивает пиксели, поэтому модель может представить зависимости между пикселями при данном коде, чего не может факторизованный гауссовский декодер: он считает их независимым шумом вокруг одного среднего изображения. Если $f_z$ — тождественное отображение, модель становится обычным VAE с гауссовским декодером с обучаемыми дисперсиями. Формула верна и тогда, когда поток не зависит от $z$; код тогда действует только через $\mu_\theta$ и $\sigma_\theta$ (раздел 3.1).

Условные слои связи#

Обратное отображение $f_z^{-1}$ — стопка условных слоёв связи (conditional coupling layers; раздел 3.2). Как в аффинном слое связи, вход слоя разбивается на две части $x_a$ и $x_b$: первая проходит без изменений, а вторая масштабируется и сдвигается функциями от первой. Здесь эти функции видят ещё и код (формулы 12–13):

$$y_a = x_a, \qquad y_b = x_b \odot \exp\big(s_z(x_a)\big) + t_z(x_a).$$

При фиксированном $z$ матрица Якоби по-прежнему треугольная, $\log\lvert\det\rvert = \sum_i s_z(x_a)_i$, а обратное преобразование — $x_b = \big(y_b - t_z(y_a)\big) \odot \exp\big(-s_z(y_a)\big)$ (формулы 14–15). Для обращения нужен код, и он всегда известен: при обучении — от энкодера, при генерации — из априорного распределения.

Масштаб объединяет сеть $s_1$ от неизменной части (в статье — два остаточных блока) с небольшой сетью $s_2$ из транспонированных свёрток (transposed convolutions), которая превращает низкоразмерный код в карту той же формы (формула 16):

$$\begin{aligned} s_z(x_a) ={}& \kappa \odot s_1(x_a) \odot s_2(z) + \beta_1 \odot s_1(x_a) \\ &+ \beta_2 \odot s_2(z) + b . \end{aligned}$$

Векторы $\kappa$, $\beta_1$, $\beta_2$ и $b$ обучаются отдельно для каждого канала и распространяются (broadcasting) по позициям, а сдвиг $t_z$ строится так же из собственной пары сетей. Слагаемое с произведением позволяет коду включать и выключать признаки $x_a$, а прямое слагаемое даёт каждому слою короткий путь к коду. С этими короткими путями авторы связывают сильное влияние кода и хорошие результаты потока с меньшим числом слоёв.

Разобранный пример: одно точное вычисление правдоподобия#

Игрушечный случай: $D = 2$ и скалярный код $z = 0.5$. Один условный слой связи использует $s_z(x_a) = x_a z$ и $t_z(x_a) = x_a - z$, а декодер выдаёт $\mu_\theta(z) = (z, 2z) = (0.5, 1.0)$ и $\sigma_\theta = (1, 0.5)$. Возьмём изображение $x = (0.4, 1.0)$.

1. Переход к $y$. $y_a = 0.4$. Слой вычисляет $s = 0.4 \cdot 0.5 = 0.2$ и $t = 0.4 - 0.5 = -0.1$, так что $y_b = 1.0 \cdot e^{0.2} - 0.1 = 1.1214$.
2. Гауссовское слагаемое. $\log\mathcal N(0.4;\, 0.5, 1^2) + \log\mathcal N(1.1214;\, 1.0, 0.5^2) = -0.9239 - 0.2553 = -1.1792$.
3. Логарифм определителя. $s = 0.2$: слой растягивает направление $x_b$ в $e^{0.2}$ раз, поэтому плотность $x$ в $e^{0.2}$ раз больше плотности $y$.
4. Итого. $\log p_\theta(x \mid z) = -1.1792 + 0.2 = -0.9792$ ната.

В обратную сторону $x_b = (1.1214 + 0.1)\, e^{-0.2} = 1.0$: при известном $z$ слой обращается точно.

Обучение, сэмплирование и результаты#

• Обучение. Модель максимизирует ELBO с помощью Adam. Авторы наблюдали, как слагаемое KL в начале обучения уходит в ноль, — это коллапс апостериорного распределения (posterior collapse) при декодере, достаточно сильном, чтобы моделировать изображения без кода, — и боролись с этим отжигом KL (KL annealing) из Bowman et al. (2016), раздел 3.1: вес при слагаемом KL начинается с нуля и растёт до единицы по ходу обучения (раздел 3.3).
• Непрерывные данные. Потоку нужны непрерывные входы. 8-битные изображения проходят деквантование (dequantization) равномерным шумом (см. врезку о дискретных данных), масштабируются в $[0, 1]$ и отображаются по формуле $x' = \operatorname{logit}\big(\alpha + (1 - \alpha)\,x\big)$ с $\alpha = 0.05$, а оценка поправляется на логарифм якобиана этой предобработки, чтобы она относилась к самим изображениям (раздел 4.1). Как потоки применяются к изображениям, разбирается в заметке Архитектуры потоков.
• Сэмплирование. Берём $z \sim \mathcal N(0, I)$, затем $y \sim \mathcal N(\mu_\theta(z), \sigma_\theta^2(z))$ и возвращаем $x = f_z(y)$ (раздел 3.4). Декодер VAE тоже задаёт распределение, но его сэмплы для одного кода различаются лишь независимым шумом в пикселях, поэтому обычно показывают только его среднее. Здесь шум в $y$ проходит через поток, который смешивает пиксели, поэтому один код даёт разные изображения; $y = \mu_\theta(z)$ даёт одно изображение на код. Реконструкция устроена так же, только $z$ берётся от энкодера (раздел 3.5).
• Архитектура (раздел 4.2). Восьмислойный свёрточный энкодер и восьмислойный декодер из транспонированных свёрток, код размерности 256 и многомасштабный real NVP с двумя масштабами. На каждом масштабе 3 условных слоя связи с шахматными масками и 3 с поканальными масками; в остаточных блоках первого масштаба 64 фильтра, второго — 128.

Результаты на тестовой выборке в битах на размерность (bits per dimension), таблицы 1–2; чем меньше, тем лучше, а $<$ отмечает верхнюю оценку, полученную из ELBO:

Модель	CIFAR-10	CelebA
PixelRNN	3.00	—
VAE с обратным авторегрессионным потоком (IAF)	$< 3.28$	—
Real NVP	3.49	3.02
Convolutional DRAW	$< 3.59$	—
VAPNEV	$< 3.55$	$< 2.8$

На CIFAR-10 VAE с одним стохастическим слоем и таким правдоподобием примерно наравне с real NVP, у которого 8 остаточных блоков на слой связи против 2, и его оценка немного ниже, чем у convolutional DRAW — VAE с несколькими стохастическими слоями и рекуррентностью. Он остаётся позади авторегрессионных моделей и VAE с апостериорным распределением на основе потока. На CelebA он обходит real NVP, используя 2 масштаба вместо 4 (раздел 4.4). Реконструкции сохраняют высокоуровневое содержание: позу, цвет фона, цвет волос и выражение лица (рисунки 3–4). Биты на размерность — та же мера, что в заметке об авторегрессионных моделях.

Условный VAPNEV#

Agrawal и Dukkipati предлагают условное моделирование плотности как расширение, но не реализуют его (раздел 5). Конструкция та же, что у условного VAE. Для изображения $x$ с условием $c$ — у CelebA это 40 бинарных атрибутов — сделаем условными энкодер, $q_\varphi(z \mid x, c)$, гауссиану декодера, $\mathcal N\big(\mu_\theta(z, c), \sigma_\theta^2(z, c)\big)$, и слои связи, у которых ветвь кода $s_2$ (и её аналог для сдвига) получает $z$ и $c$ вместе. Обозначим получившийся поток $f_{z,c}$ и оставим априорное распределение $p(z) = \mathcal N(0, I)$. Оценка принимает вид

$$\begin{aligned} \log p_\theta(x \mid c) \ge{}& \mathbb E_{q_\varphi(z \mid x, c)}\Big[\log\mathcal N\big(f_{z,c}^{-1}(x);\, \mu_\theta(z, c), \sigma_\theta^2(z, c)\big) \\ &\qquad + \log\left|\det\frac{\partial f_{z,c}^{-1}(x)}{\partial x}\right|\Big] \\ &- D_{\mathrm{KL}}\big(q_\varphi(z \mid x, c)\,\|\,p(z)\big). \end{aligned}$$

Сэмплирование фиксирует атрибуты, берёт $z$ из априорного распределения, берёт $y$ из гауссианы декодера $\mathcal N\big(\mu_\theta(z, c), \sigma_\theta^2(z, c)\big)$ и возвращает $x = f_{z,c}(y)$. Если изменить один атрибут, сохранив $z$, мы запрашиваем то же лицо с изменённым атрибутом — в той мере, в какой код не вобрал в себя сам атрибут.

Что изменилось#

	VAE	β-VAE	VQ-VAE / VQ-VAE-2	VAPNEV
Код	непрерывный, гауссов	непрерывный, гауссов	сетка индексов в словаре	непрерывный, гауссов
Целевая функция	ELBO	реконструкция − β · KL	реконструкция + словарь + привязка	ELBO с точным правдоподобием потока
Роль априорного распределения	фиксированное $\mathcal N(0, I)$, задаётся KL	то же, с весом $\beta$	равномерное при обучении; потом обучается авторегрессионной моделью	фиксированное $\mathcal N(0, I)$, задаётся KL
Градиент через код	репараметризация	репараметризация	прямой проброс	репараметризация
Сэмплирование	декодировать $z \sim \mathcal N(0, I)$	то же	авторегрессионно сэмплировать коды, затем декодировать	декодировать $z$ в гауссиану над $y$, сэмплировать $y$, применить поток
Типичная беда	размытость, дыры в априорном	сам компромисс «скорость–искажение»	мёртвые коды, медленная последовательная априорная модель	коллапс KL без отжига, дорогой декодер в полном разрешении

Ограничения#

• β двигает вдоль компромисса, но не избавляет от него. Меньшее искажение стоит скорости. Скорость — это в основном информация об изображении и не обязательно рассогласование с априорным распределением; плохие сэмплы из априорного распределения возникают, когда от него уходит агрегированное апостериорное распределение, а малое $\beta$ штрафует за это слабее. Распутывание — возможный побочный эффект, а не свойство целевой функции.
• Прямой проброс градиента — заменитель. Он использует единичный якобиан там, где настоящая производная квантователя равна нулю, поэтому точным градиентом он не бывает никогда; слагаемое привязки держит $z_e$ рядом с его вектором словаря, и это делает подстановку разумной, но не точной.
• Схлопывание словаря. Без перезапусков или нормированного поиска заметная доля кодов может остаться неиспользуемой.
• Двухэтапное обучение. Априорная модель подгоняется к кодам замороженного автоэнкодера; ошибки первого этапа, например сетку кодов, которую трудно предсказывать, второй исправить не может.
• Декодер-поток силён и дорог. Декодер VAPNEV может моделировать изображения без кода, поэтому слагаемое KL схлопывается, если не применять отжиг, а его поток работает в полном разрешении изображения, что обходится гораздо дороже декодера из транспонированных свёрток.
• Последовательное сэмплирование. Авторегрессионная априорная модель сэмплирует по одному индексу. Сетка меньше пикселей, но всё равно последовательна; дальнейшие работы заменяют априорную модель трансформерами по тем же кодам или диффузией в непрерывном латентном пространстве (латентная диффузия).

Итоги#

• β-VAE умножает слагаемое KL на $\beta$ и двигает модель вдоль кривой «скорость–искажение»: $\beta > 1$ даёт меньшую скорость и худшую реконструкцию, $\beta < 1$ — наоборот. $\beta = 1$ возвращает стандартную ELBO, $\beta > 1$ максимизирует более слабую нижнюю оценку, а при $\beta < 1$ целевая функция может и не оценивать правдоподобие снизу.
• VQ-VAE заменяет гауссов код ближайшим вектором словаря. KL становится константой, argmin преодолевается прямым пробросом градиента, а функция потерь добавляет слагаемые словаря и привязки — или EMA-обновление, превращающее словарь в онлайн-k-means.
• Априорное распределение обучается после автоэнкодера авторегрессионной моделью по сетке кодов. VQ-VAE-2 использует две сетки — грубую для компоновки и мелкую для деталей — и по априорной модели на каждую.
• VAPNEV сохраняет гауссов код и ELBO и заменяет попиксельную реконструкцию точным правдоподобием потока real NVP, у которого базовая гауссиана и слои связи зависят от кода. Чтобы код продолжал использоваться, нужен отжиг KL. Условная версия добавляет метку в энкодер, декодер и слои связи.

Ноутбук#

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

• Ячейка 6: Quantize вычисляет квадраты расстояний сразу до всех 512 векторов словаря через разложение $\lVert a - e\rVert^2 = \lVert a\rVert^2 - 2a^\top e + \lVert e\rVert^2$, одним матричным произведением, выбирает ближайший, обновляет словарь описанным выше EMA (затухание 0.99), возвращает расстояние привязки diff = mean((quantize.detach() - input)²) и применяет прямой проброс input + (quantize - input).detach(). Словарь — буфер, взятый из стандартного нормального распределения, а не параметр, поэтому Adam его не трогает; словарь меняется только через EMA и только в режиме обучения. Счётчики сглаживаются как $N_k \leftarrow \frac{N_k + \epsilon}{n + K\epsilon}\, n$ с $n = \sum_k N_k$ и $\epsilon = 10^{-5}$; это сохраняет их сумму и делает каждый счётчик положительным. Одна проблема: счётчики начинаются с $0$, а суммы $m_k$ — со случайного словаря. Поэтому после первого обновления код, который первый батч не выбрал, делится на сглаженный счётчик порядка $\epsilon$, и его норма вырастает примерно в $10^5$ раз. Обратно она уменьшается лишь как $0.99^t$, что занимает около $\ln 10^5 / {-\ln 0.99} \approx 1150$ шагов. В нашем собственном игрушечном прогоне этого обновления со стандартными нормальными выходами энкодера (это не выход ноутбука) между шагами 2 и 1000 выбирались только 2–3 кода из 512. Если начать счётчики с $1$, так что первое частное $m_k / N_k$ — сам случайный словарь, скачка не будет.
• Ячейки 7–10: архитектура. ResBlock — это ReLU, свёртка $3 \times 3$ из 128 в 32 канала, ReLU и свёртка $1 \times 1$ обратно в 128 каналов; результат прибавляется к входу блока. Нижний энкодер применяет две свёртки $4 \times 4$ с шагом 2 и свёртку $3 \times 3$, верхний — одну свёртку $4 \times 4$ с шагом 2 и свёртку $3 \times 3$; за каждым из них следуют остаточные блоки. Декодеры устроены зеркально: свёртка $3 \times 3$, остаточные блоки и транспонированные свёртки $4 \times 4$ с шагом 2. Настройки: channel=128, n_res_block=2, n_res_channel=32, словари из 512 векторов размерности 64. Нижний энкодер уменьшает разрешение в 4 раза (до $8 \times 8$ для CIFAR-10), верхний — ещё в 2 (до $4 \times 4$), а свёртки $1 \times 1$ перед квантованием проецируют признаки в 64-мерные векторы кода. Верхний код попадает на нижний уровень двумя отдельными путями: dec_t, небольшой декодер, возвращает его к $8 \times 8$, чтобы объединить с нижними признаками и квантовать как нижний код, а upsample_t, одна транспонированная свёртка, увеличивает его для итогового декодера, который получает его вместе с нижним кодом.
• Ячейка 12: обучение минимизирует среднюю квадратичную ошибку реконструкции плюс $0.25 \cdot$ diff, слагаемое привязки; слагаемого словаря нет, потому что его заменяет EMA. Каждые sample_interval шагов (200 в ячейке 20) сохраняется сетка: батч, а под ним его реконструкции, вычисленные в режиме оценки.
• Ячейки 13 и 22: infer_vqvae реконструирует два батча по 64 тестовых изображения CIFAR-10 в режиме оценки и сохраняет 128 реконструкций в файл.
• Ячейки 14–20: CIFAR-10 в разрешении $32 \times 32$, нормированный к $[-1, 1]$ (комментарий в ячейке 15 всё ещё говорит о $64 \times 64$), батч 32, Adam с шагом $3 \cdot 10^{-4}$, пять эпох.
• Ячейки 25–27: априорные модели. Маскированные свёртки типов A и B и две априорные модели PixelCNN из семи маскированных слоёв $7 \times 7$ по 64 канала с обычными ReLU между ними, без гейтов. На входе эмбеддинг индексов кода, Embedding(512, 64), а свёртка $1 \times 1$ даёт 512 логитов на позицию. Безусловная априорная модель описывает верхние коды. Условная описывает нижние: она строит эмбеддинги верхних кодов, увеличенных методом ближайшего соседа, через Embedding(512, 16), конкатенирует их с эмбеддингом нижних кодов и смешивает их свёрткой $1 \times 1$. Здесь две проблемы. Такая простая стопка слоёв имеет слепое пятно, и в ней нет гейтированных слоёв и внимания, как в априорных моделях статьи. А поскольку условие подмешивается в те же каналы до первого маскированного слоя типа A, маска скрывает и условие в предсказываемой позиции. Каждый нижний индекс видит верхние коды только через предыдущие позиции, поэтому левый верхний индекс каждого блока $2 \times 2$, чей верхний код не встречается ни в одной предыдущей позиции, предсказывается без него. Мы проверили это, продифференцировав собственную повторную реализацию сети: у логитов в позиции нулевой градиент по условию в той же позиции. Условие должно входить без маски и добавляться в каждом слое, как в условном Gated PixelCNN.
• Ячейки 29–32: датасеты, на лету кодирующие каждое изображение в сетки кодов, и обучение двух априорных моделей кросс-энтропией с Adam при шаге обучения $10^{-3}$, батче 32 и 10 эпохах для каждой. Кодировать нужно с VQ-VAE в режиме оценки, в котором его оставляет ячейка 22; в режиме обучения квантизатор продолжал бы обновлять словарь, пока извлекаются коды.
• Ячейки 34–35: сэмплирование — 16 последовательных шагов для верхней сетки, 64 для нижней при данной увеличенной верхней — и декодирование. Каждый шаг заново прогоняет PixelCNN по всей сетке.
• Замечание о реализации. utils.save_image(..., range=(-1, 1)) в ячейках 12, 22 и 35 использует имя аргумента, которое в torchvision переименовано в value_range; в текущих версиях эти вызовы требуют нового имени.

У β-VAE нет обученного примера в ноутбуках курса: ноутбук VAE определяет функцию потерь с весом β в ячейке 24, но обучает только с $\beta = 1$.

У VAPNEV нет реализации в репозитории курса. Ноутбук с упражнениями 03_hw_vapnev.ipynb из того же репозитория формулирует задания на CelebA и оставляет ячейки кода пустыми. Его единственные сохранённые выходы — устройство (mps), форма одного изображения, $3 \times 64 \times 64$, и сетка из 25 случайных лиц.

• Ячейки 3–9: данные. CelebA через класс датасета курса, с Resize(64) и CenterCrop(64), которые сохраняют соотношение сторон, и пикселями, переведёнными из $[0, 1]$ в $[-1, 1]$ преобразованием $2x - 1$.
• Ячейки 10–11: нормализующий поток на этих лицах, оцениваемый по FID и отрицательному логарифму правдоподобия. Допускаются любые слои — слои связи, авторегрессионные или линейные, свёрточные или на вытянутом в вектор изображении; слои описаны в заметке Архитектуры потоков. Правдоподобие, которое приводится для изображений, должно учитывать деквантование и якобиан предобработки, как это делают Agrawal и Dukkipati для своего логит-преобразования.
• Ячейки 12–13: VAE на тех же лицах, оцениваемый по FID, как в ноутбуке о VAE.
• Ячейки 14–15: VAPNEV по статье и условный VAPNEV. Подробности условия ноутбук оставляет формулировке задания; естественный выбор — 40 атрибутов CelebA, которые класс датасета курса возвращает вместе с каждым изображением, а обёртка из ячейки 6 отбрасывает.
• Ячейки 16–17: анализ латентных пространств: какая модель лучше интерполирует, кодирование лица не из CelebA и поиск ближайших к нему кодов обучающих изображений, а также самый интересный сэмпл.

Дальше#

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

Литература#

• I. Higgins, L. Matthey, A. Pal, C. Burgess, X. Glorot, M. Botvinick, S. Mohamed, A. Lerchner. β-VAE: Learning basic visual concepts with a constrained variational framework. ICLR, 2017. Раздел 1: определение распутывания; рисунок 1: изменение отдельных латентных координат на CelebA (latent traversals).
• A. Alemi, B. Poole, I. Fischer, J. Dillon, R. A. Saurous, K. Murphy. Fixing a Broken ELBO. ICML, 2018. Скорость и искажение.
• C. P. Burgess, I. Higgins, A. Pal, L. Matthey, N. Watters, G. Desjardins, A. Lerchner. Understanding disentangling in β-VAE. arXiv, 2018.
• F. Locatello, S. Bauer, M. Lucic, G. Rätsch, S. Gelly, B. Schölkopf, O. Bachem. Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations. ICML, 2019.
• A. van den Oord, O. Vinyals, K. Kavukcuoglu. Neural Discrete Representation Learning. NeurIPS, 2017. VQ-VAE; EMA-обновление словаря — в приложении A.1.
• A. Razavi, A. van den Oord, O. Vinyals. Generating Diverse High-Fidelity Images with VQ-VAE-2. NeurIPS, 2019.
• Y. Bengio, N. Léonard, A. Courville. Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation. arXiv, 2013. Прямой проброс градиента.
• P. Dhariwal, H. Jun, C. Payne, J. W. Kim, A. Radford, I. Sutskever. Jukebox: A Generative Model for Music. arXiv, 2020. Случайные перезапуски неиспользуемых кодов.
• J. Yu, X. Li, J. Y. Koh, H. Zhang, R. Pang, J. Qin, A. Ku, Y. Xu, J. Baldridge, Y. Wu. Vector-quantized Image Modeling with Improved VQGAN. ICLR, 2022. Факторизованный нормированный поиск в словаре.
• S. Agrawal, A. Dukkipati. Deep Variational Inference Without Pixel-Wise Reconstruction. arXiv:1611.05209, 2016. VAPNEV; раздел 3: условное правдоподобие и условный слой связи; раздел 4: предобработка, архитектура и результаты.
• L. Dinh, J. Sohl-Dickstein, S. Bengio. Density Estimation Using Real NVP. ICLR, 2017.
• S. R. Bowman, L. Vilnis, O. Vinyals, A. M. Dai, R. Jozefowicz, S. Bengio. Generating Sentences from a Continuous Space. CoNLL, 2016. Раздел 3.1: отжиг веса KL (KL cost annealing).