VAE: вероятностный латентный мир

Вариационный автокодировщик кодирует объект не точкой, а облаком. Реконструкция требует сохранить детали, расстояние Кульбака — Лейблера требует уложить облака в простое опорное распределение. Их спор и делает скрытое пространство пригодным сразу для двух дел: восстанавливать старое и порождать новое.

Код, из которого нечего доставать

Возьмём 1797 рукописных цифр $8\times8$ из классического набора digits — по 1400 на обучение и 397 на проверку, яркость каждого из 64 пикселей поделена на 16 и превращена в число от нуля до единицы. Обучим на них обычный автокодировщик с двумерным узким местом: кодировщик $f(x)=z$ сжимает картинку до двух чисел, декодировщик $g(z)\approx x$ разворачивает её обратно. На обучающей выборке он работает хорошо: 19,6 ната ошибки на объект.

А теперь попробуем сделать то, ради чего вообще заводят порождающую модель: взять случайную точку $z\sim\mathcal N(0,I)$ и посмотреть, что нарисует декодировщик. Результат обескураживает — и причина видна сразу, стоит нарисовать облако кодов. Коды разъехались по квадрату со стороной почти 57: максимальная координата равна 28,3. А в круге $\|z\|\le2$, куда стандартная двумерная нормаль кладёт 86,5% своей массы, лежит ровно 1,6% кодов.

Рисунок шире экрана — проведите по немуОткрыть целиком ↗

Никто не просил автокодировщик держать коды около нуля — вот он и не держит. Величина $\tfrac12\sum_j(\mu_j^2+\sigma_j^2-\ln\sigma_j^2-1)$, которая ниже окажется центральной, равна у него 93,6 ната против 2,62 у вариационной версии. Дело не в том, что коды «плохие»: линейный классификатор по этим двум числам угадывает цифру в 74,1% случаев против 61,5% у VAE. Дело в том, что мы не знаем, откуда брать новые $z$. Пространство кодов без заданного распределения — это карта без масштабной линейки.

Кодировщик, который выдаёт распределение

Вариационный автокодировщик заменяет точку облаком. Кодировщик с параметрами $\phi$ выдаёт для объекта $x$ не вектор, а нормальное распределение с диагональной ковариацией:

$$q_\phi(z\mid x)=\mathcal N\bigl(\mu_\phi(x),\ \operatorname{diag}\sigma_\phi^2(x)\bigr).$$

Опорное (априорное) распределение фиксировано и просто:

$$p(z)=\mathcal N(0,I),$$

а декодировщик задаёт правдоподобие наблюдения:

$$p_\theta(x\mid z)=\prod_{k=1}^{64}\operatorname{Bern}\bigl(x_k\mid g_\theta(z)_k\bigr).$$

Оговоримся честно: яркость пикселя — не двоичная величина, мы трактуем её как вероятность «зажечь пиксель». Это предположение о наблюдении, и оно влияет на всё остальное; к нему мы ещё вернёмся.

Хочется, как в уроке 45, максимизировать $\log p_\theta(x)$. Но интеграл по всем $z$ не берётся: декодировщик — нейросеть, никакой формулы для $\int g_\theta(z)p(z)\,dz$ нет. Можно ли оценить интеграл выборкой из $p(z)$? Формально да:

$$p_\theta(x)\approx\frac1M\sum_{m=1}^Mp_\theta(x\mid z^{(m)}),\qquad z^{(m)}\sim p(z).$$

Практически — нет. Почти все случайные $z$ дают ничтожное $p_\theta(x\mid z)$ для конкретного $x$, и оценка почти всегда близка к нулю с редкими огромными выбросами. Коды нужно брать не откуда попало, а оттуда, где они правдоподобны для данного $x$. То есть нужен приближённый апостериор — тот же ход, что в уроке 47, только вместо честной формулы Байеса мы будем учить приближение.

Нижняя граница вместо неберущегося интеграла

Возьмём произвольное распределение $q(z\mid x)$ и вставим его под логарифм:

$$\log p_\theta(x)=\log\int q(z\mid x)\,\frac{p_\theta(x,z)}{q(z\mid x)}\,dz =\log\mathbb E_{q}\frac{p_\theta(x,z)}{q(z\mid x)}.$$

Логарифм вогнут, поэтому по неравенству Йенсена

$$\log\mathbb E_q\frac{p_\theta(x,z)}{q(z\mid x)}\ \ge\ \mathbb E_q\log\frac{p_\theta(x,z)}{q(z\mid x)}.$$

Правую часть называют нижней границей обоснованности, ELBO. Раскроем совместную плотность и получим рабочую форму:

$$\mathcal L(\theta,\phi;x)=\mathbb E_{q_\phi(z\mid x)}\log p_\theta(x\mid z) -D_{\mathrm{KL}}\bigl(q_\phi(z\mid x)\,\|\,p(z)\bigr).$$

Есть и вторая дорога к тому же результату, и она объясняет размер зазора. Подставив $p_\theta(z\mid x)=p_\theta(x,z)/p_\theta(x)$ в определение KL, получаем тождество:

$$\log p_\theta(x)=\mathcal L(\theta,\phi;x) +D_{\mathrm{KL}}\bigl(q_\phi(z\mid x)\,\|\,p_\theta(z\mid x)\bigr).$$

Расстояние Кульбака — Лейблера неотрицательно, поэтому ELBO действительно не превосходит логарифма правдоподобия, а зазор равен ровно тому, насколько наш приближённый апостериор отличается от истинного. Максимизируя $\mathcal L$ по $\theta$, мы улучшаем порождающую модель; максимизируя по $\phi$ — поджимаем зазор. Одна формула тянет обе задачи.

KL в закрытой форме: за что именно штрафуют

Для нормальных распределений с диагональными ковариациями расстояние считается вручную. Для одной координаты:

$$D_{\mathrm{KL}}\bigl(\mathcal N(\mu,\sigma^2)\,\|\,\mathcal N(0,1)\bigr) =\frac12\bigl(\mu^2+\sigma^2-\ln\sigma^2-1\bigr),$$

а для $d$ независимых координат слагаемые складываются:

$$D_{\mathrm{KL}}\bigl(q_\phi(z\mid x)\|p(z)\bigr)=\frac12\sum_{j=1}^d \bigl(\mu_j^2+\sigma_j^2-\ln\sigma_j^2-1\bigr).$$

Формула читается как счёт за две провинности. Слагаемое $\mu^2$ штрафует за уход от начала координат: сдвиг на единицу при $\sigma=1$ стоит 0,5 ната. Слагаемое $\sigma^2-\ln\sigma^2-1$ штрафует за любую ширину, отличную от единичной, причём несимметрично: при $\sigma=2$ штраф равен 0,807 ната, а при вдвое меньшей $\sigma=0{,}5$ — только 0,318. Слишком широкое облако наказывается сильнее слишком узкого — но наказание за узость растёт неограниченно: при $\sigma\to0$ член $-\ln\sigma^2$ уходит в бесконечность. Именно это и запрещает вырождение в точку, то самое вырождение, которое мы наблюдали у обычного автокодировщика.

Пинскер: почему малый KL вообще о чём-то говорит

Мы штрафуем $D_{\mathrm{KL}}$ — величину, у которой нет наглядного смысла: она не метрика (не симметрична и не подчиняется неравенству треугольника) и измеряется в натах. Почему её малость означает, что распределения близки в человеческом смысле?

Ответ дал Марк Семёнович Пинскер, один из создателей советской школы теории информации в Институте проблем передачи информации. В монографии 1960 года «Информация и информационная устойчивость случайных величин и процессов» он доказал неравенство, связывающее информационное уклонение с честной метрикой — расстоянием по вариации:

$$\sup_A\bigl|P(A)-Q(A)\bigr|\ \le\ \sqrt{\tfrac12\,D_{\mathrm{KL}}(P\|Q)}.$$

Левая часть — максимальная разница вероятностей одного и того же события, величина от 0 до 1, понятная кому угодно. Правая — то, что мы оптимизируем. Малый KL гарантирует, что никакое событие не различает распределения сильно. Обратное неверно, и это тоже важно: KL может быть огромным при крошечном расстоянии по вариации.

Для нашего дела вывод такой: добиваясь малого $D_{\mathrm{KL}}(q\|p)$, мы получаем право сэмплировать из $p$ вместо $q$ — ошибка в вероятности любого события ограничена корнем. Заодно понятно, почему в уроке 20 кросс-энтропия оказалась естественной функцией потерь: она отличается от KL на энтропию данных, то есть на константу.

Репараметризация: как провести градиент сквозь жребий

ELBO содержит математическое ожидание по $q_\phi$, а параметры $\phi$ сидят внутри самого распределения. Наивная выборка $z\sim\mathcal N(\mu,\sigma^2)$ разрывает вычислительный граф: узел «бросить жребий» не имеет производной по $\mu$. Спасает элементарная замена переменной:

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

Теперь случайность живёт в отдельном входе $\varepsilon$, который не зависит от параметров, а $z$ — гладкая функция от $\mu$ и $\sigma$. Обратное распространение из урока 25 проходит насквозь:

$$\frac{\partial z}{\partial\mu}=1,\qquad \frac{\partial z}{\partial\sigma}=\varepsilon.$$

Есть и альтернатива — оценка через логарифмическую производную (её называют score-function, или REINFORCE):

$$\nabla_\mu\mathbb E_q[f(z)]=\mathbb E_q\bigl[f(z)\,\nabla_\mu\log q(z)\bigr].$$

Обе оценки несмещённые, но неравноценные. Возьмём игрушечный пример $f(z)=z^2$, $z\sim\mathcal N(\mu,1)$, $\mu=1$; истинный градиент равен $2\mu=2$. На 20 000 жребиев обе оценки дали в среднем 2,03 и 2,04 — несмещённость видна. А вот разброс различается радикально: 4,0 против 31,45, то есть в 7,9 раза.

Рисунок шире экрана — проведите по немуОткрыть целиком ↗

Лаборатория: спор реконструкции и штрафа

Ниже — модель, которую можно решить на бумаге и потому не нужно обучать. Пусть объектов $K$ штук, у каждого есть «своя точка» $t_i$ на прямой, декодировщик — тождественное отображение $g(z)=z$, а апостериор $q(z\mid x_i)=\mathcal N(\mu_i,\sigma^2)$ с общей шириной. Ошибка реконструкции для объекта:

$$\mathbb E_{z\sim q}\bigl(z-t_i\bigr)^2=(\mu_i-t_i)^2+\sigma^2.$$

Полная цель с весом $\beta$ при штрафе:

$$J(\mu_i,\sigma)=(\mu_i-t_i)^2+\sigma^2 +\beta\cdot\tfrac12\bigl(\mu_i^2+\sigma^2-\ln\sigma^2-1\bigr).$$

Приравняем производные нулю:

$$\frac{\partial J}{\partial\mu_i}=2(\mu_i-t_i)+\beta\mu_i=0 \ \Longrightarrow\ \mu_i=\frac{2}{2+\beta}\,t_i,$$ $$\frac{\partial J}{\partial(\sigma^2)}=1+\frac\beta2\Bigl(1-\frac1{\sigma^2}\Bigr)=0 \ \Longrightarrow\ \sigma^2=\frac{\beta}{\beta+2}.$$

Две формулы описывают всё поведение VAE в миниатюре. Коды сжимаются к нулю в $2/(2+\beta)$ раз — это ровно то же стягивание, что у регуляризованной регрессии в уроке 23. Ширина облака растёт от нуля к единице. При $\beta=0$ получается детерминированный автокодировщик, при $\beta\to\infty$ — $\mu\to0$ и $\sigma\to1$: апостериор в точности совпал с опорным распределением и не несёт об объекте ничего.

Проверьте на виджете: при $\beta=1$ и пяти объектах с шагом 1,6 сжатие равно 0,667, ширина $\sigma=0{,}577$, скорость 1,354 ната, искажение 0,902, а узнаваемость объекта (доля жребиев, попадающих ближе к своей цели, чем к чужой) равна 56,3%. Поднимите $\beta$ до 8 — узнаваемость падает до 25,0% при случайном угадывании 20,0%: код почти перестал различать объекты. Сдвиньте ползунок «разнос объектов» — и увидите, что при тесно расположенных целях та же ширина облака губит узнаваемость гораздо раньше.

Скорость и искажение: чем платит $\beta$

Обобщённая цель $\beta$-VAE записывается как сумма двух слагаемых с явным весом:

$$\mathcal L_\beta=\underbrace{-\mathbb E_q\log p_\theta(x\mid z)}_{\text{искажение}} +\beta\underbrace{D_{\mathrm{KL}}\bigl(q_\phi(z\mid x)\|p(z)\bigr)}_{\text{скорость}}.$$

Слова «скорость» и «искажение» здесь не метафора: это термины теории информации. Скорость — сколько информации код несёт об объекте сверх опорного распределения; искажение — насколько плохо объект восстанавливается. Меняя $\beta$, мы движемся по границе достижимых пар, как по шенноновской кривой для источника с критерием точности.

Обучим на тех же цифрах восьмимерный VAE при восьми значениях $\beta$ (обучение с нуля, один и тот же seed, 300 эпох). Числа на тестовой выборке:

$$\begin{array}{c|c|c|c|c} \beta & \text{искажение, нат} & \text{скорость, нат} & \text{активных координат} & \text{линейный зонд}\\\hline 0 & 16{,}5 & 133{,}89 & 8 & 92{,}9\%\\ 0{,}25 & 16{,}7 & 9{,}72 & 8 & 93{,}7\%\\ 1 & 18{,}9 & 3{,}46 & 6 & 90{,}2\%\\ 2 & 22{,}1 & 1{,}27 & 4 & 83{,}6\%\\ 8 & 27{,}1 & 0{,}00 & 0 & 22{,}9\% \end{array}$$

Самая поучительная строка — вторая. Переход от $\beta=0$ к $\beta=0{,}25$ уменьшает скорость в 13,8 раза (со 133,89 до 9,72 ната), а платит за это ростом искажения на 0,18 ната. Почти вся «информация», которую нёс код обычного автокодировщика, была не нужна для восстановления: это был произвольный масштаб, который мы разбирали в начале урока.

Рисунок шире экрана — проведите по немуОткрыть целиком ↗

Коллапс апостериора: оптимизатор нашёл лазейку

При $\beta=4$ и $\beta=8$ обе колонки застывают: скорость 0,00, активных координат 0, искажение 27,1 ната. Ровно столько же — 27,1 ната — даёт тривиальный предсказатель, выдающий средний по обучающей выборке пиксель и вообще не глядящий на объект. Это и есть коллапс апостериора: модель сочла, что дешевле отказаться от кода, чем платить за него.

$$q_\phi(z\mid x)\approx p(z)\quad\Longrightarrow\quad D_{\mathrm{KL}}\approx0, \qquad p_\theta(x\mid z)\approx p_\theta(x).$$

Коварство в том, что цель при этом выглядит достойно: ELBO равна $-27{,}1$ ната и больше не меняется, обучение «сошлось». Никакая кривая обучения об ошибке не кричит. Кричит только диагностика: число активных координат $\#\{j:\operatorname{Var}_x\mu_j(x)>0{,}01\}$ и линейный зонд, упавший до 22,9%.

Сетка опорного распределения: карта, по которой можно ходить

Самая честная проверка порождающей модели — декодировать не коды объектов, а регулярную сетку квантилей опорного распределения. Возьмём 11 квантилей от 0,02 до 0,98 по каждой оси, получим 121 точку и посмотрим на картинки.

Рисунок шире экрана — проведите по немуОткрыть целиком ↗

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

Агрегат апостериоров не равен опорному распределению

Штраф действует пообъектно: каждое облако $q(z\mid x_i)$ прижимается к $\mathcal N(0,I)$. Но сэмплируем мы из общего распределения, и сравнивать надо смесь:

$$q(z)=\frac1N\sum_{i=1}^Nq(z\mid x_i).$$

Малость среднего пообъектного KL не влечёт совпадения смеси с опорным распределением. Более того, из разложения

$$\frac1N\sum_iD_{\mathrm{KL}}\bigl(q(z\mid x_i)\|p(z)\bigr) =\underbrace{I(x;z)}_{\text{взаимная информация}} +\underbrace{D_{\mathrm{KL}}\bigl(q(z)\|p(z)\bigr)}_{\text{несовпадение агрегата}}$$

видно, что штраф давит сразу на две вещи: на информативность кода и на несовпадение смеси. Уменьшая его, мы не различаем, что именно уменьшилось. На наших цифрах средний радиус кода из агрегата равен 1,33, из опорного распределения — 1,24: смесь заметно шире.

Рисунок шире экрана — проведите по немуОткрыть целиком ↗

Средний код и средняя картинка — разные объекты

Реконструировать можно двумя способами: подать декодировщику $\mu(x)$ или усреднить декодирования нескольких жребиев. Это не одно и то же, потому что декодировщик нелинеен:

$$g\bigl(\mathbb E z\bigr)\ \ne\ \mathbb E\,g(z).$$

На объекте с самым широким апостериором (максимальная координата $\sigma$ равна 0,703) двести жребиев дают среднее декодирование, отличающееся от $g(\mu)$ в среднем на 0,012 яркости, а в отдельных пикселях — на 0,05. Максимальная дисперсия пикселя по жребиям равна 0,046 и сосредоточена на границах штриха: там модель и вправду не знает, где проходит линия.

Рисунок шире экрана — проведите по немуОткрыть целиком ↗

Почему картинки размыты — и почему это не лень сети

Обычное обвинение: «VAE даёт размытые изображения». Полезно уточнять, кого именно обвиняют. Возьмём две одинаково допустимые картинки $A$ и $B$ при одном и том же условии. Если модель обязана выдать одну картинку и оценивается по среднему квадрату ошибки, то оптимально выдать их среднее:

$$\arg\min_{\hat x}\ \tfrac12\|\hat x-A\|^2+\tfrac12\|\hat x-B\|^2=\frac{A+B}2.$$

На двух реальных цифрах полусумма даёт MSE 0,0173 против 0,0346 у любого из двух чётких вариантов — ровно вдвое меньше. То есть размытие — точный оптимум выбранной функции потерь, а не слабость оптимизатора. Как и в уроке 20, функция потерь есть предположение о шуме: квадрат подразумевает нормальный шум, поэлементная кросс-энтропия — независимые бернуллиевские пиксели. Ни то, ни другое не умеет сказать «здесь возможны два разных исхода».

Место VAE среди порождающих моделей

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

Полезно сравнить и с контрастивным обучением: там инварианты выбираются явно, парами «похоже — не похоже». VAE сохраняет то, что нужно для восстановления, а не то, что удобно классификатору, — отсюда и разница в поведении линейного зонда. А само разложение цели на «подгонку» и «штраф» — та же конструкция, что в уроке 23 и в байесовском обновлении: данные тянут в одну сторону, априорное предположение — в другую.

Работа Гельфанда и Яглома объясняет, почему «скорость» в нашей таблице — это взаимная информация между объектом и кодом, а неравенство Пинскера объясняет, почему её малость означает практическую неразличимость. Обе линии выросли из одной школы и ровно в те годы, когда информация только становилась математическим понятием.

Что остаётся в руках

VAE превращает сжатие в вероятностный вывод. Нижняя граница раскладывает неберущуюся задачу на два понятных слагаемых, репараметризация проводит градиент сквозь жребий с разбросом в 7,9 раза меньшим, чем у альтернативы, а вес $\beta$ переводит модель вдоль кривой «скорость — искажение»: от 133,89 ната кода при 16,5 ната ошибки до нуля нат кода при 27,1. Успех нельзя измерить одной реконструкцией: нужны сетка опорного распределения, число активных координат и явная проверка на коллапс. И главное — помнить, что размытость картинки часто говорит не о сети, а о том, что мы попросили её усреднить неопределённое будущее.

Задачи