ELBO

Задача, которую нельзя решить прямо

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

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

и этот интеграл в общем случае не берётся. Для смеси из блока 3 сумма по $K$ компонентам ещё считалась; для VAE, где $z$ — вектор из сотни непрерывных чисел, интеграл неподъёмен.

Обходной путь — не считать $p(x)$, а ограничить его снизу величиной, которую считать умеем.

Вывод в три строки

Возьмём любое распределение $q(z)$ с подходящим носителем и вставим его так, чтобы получилось ожидание:

$\log p(x) = \log \int p(x,z)\,dz = \log \int q(z)\frac{p(x,z)}{q(z)}\,dz = \log \mathbb{E}_q\!\left[\frac{p(x,z)}{q(z)}\right]$

Это importance sampling из блока 3, урок 150 — тот же вес $p/q$. Теперь по неравенству Йенсена ($\log$ вогнут, поэтому знак противоположен тому, что был в уроке про KL):

$$\htmlData{k=evidence}{\log p(x)} \ \ge \ \htmlData{k=elbo}{\mathbb{E}_{q}\!\left[\log \frac{p(x, z)}{q(z)}\right]} \ = \ \text{ELBO}(q)$$

Всё. Один шаг Йенсена — и вместо интеграла у нас ожидание по $q$, которое считается методом Монте-Карло, потому что из $q$ мы умеем сэмплировать (мы его сами выбрали).

Зазор — это ровно KL

Насколько граница неточна? Посчитаем разность честно:

$\log p(x) - \text{ELBO}(q) = \mathbb{E}_q\!\left[\log \frac{q(z)}{p(z \mid x)}\right] = \text{KL}\big(q(z) \,\|\, p(z \mid x)\big)$

Отсюда центральное тождество блока:

$\boxed{\ \log p(x) = \text{ELBO}(q) + \text{KL}\big(q \,\|\, p(z\mid x)\big)\ }$

Левая часть от $q$ не зависит вообще — это свойство модели и данных. Значит, максимизируя ELBO по $q$, вы автоматически минимизируете KL до истинного posterior. Одно действие, две цели, и никакого компромисса между ними.

Подвигайте $q$ в виджете. Потолок стоит на месте, а зелёная и красная части обмениваются высотой.

веса q(z)

z = 1

0.25

z = 2

0.25

z = 3

0.25

z = 4

0.25

взять q = истинный posterior сделать q равномерной

log p(x)

-0.99425

ELBO(q)

-1.31865

KL(q‖posterior)

0.3244

Кнопка «взять $q$ = истинный posterior» закрывает зазор ровно: KL становится нулём, а ELBO садится на $\log p(x) = -0.9943$. Равномерная $q$ даёт ELBO $-1.3187$ и зазор $0.3244$ ната. Сумма в обоих случаях одна и та же — потому что это тождество, а не приближение.

Второе прочтение: реконструкция и регуляризация

Разложим ELBO иначе, подставив $p(x,z) = p(x \mid z)p(z)$:

$\text{ELBO} = \underbrace{\mathbb{E}_q\big[\log p(x \mid z)\big]}_{\text{реконструкция}} \ - \ \underbrace{\text{KL}\big(q(z\mid x) \,\|\, p(z)\big)}_{\text{регуляризация}}$

Это та форма, в которой ELBO появляется в коде VAE, и в ней читается напряжение между двумя целями:

• первый член хочет, чтобы $z$ содержал как можно больше информации об $x$ — иначе восстановить $x$ не выйдет;
• второй член хочет, чтобы $q(z \mid x)$ был похож на prior — то есть чтобы $z$ содержал как можно меньше информации об $x$.

Крайние случаи объясняют обе типичные поломки. Если второй член побеждает, наступает posterior collapse: $q(z\mid x) = p(z)$ для всех $x$, KL равна нулю, латент не несёт ничего, а декодер выучивает безусловное распределение. Если побеждает первый — латентное пространство становится дырявым, и сэмплы из prior декодируются в мусор.

Отсюда и $\beta$-VAE: коэффициент перед KL, которым это равновесие сдвигают вручную. Заметьте, что при $\beta \ne 1$ вы оптимизируете уже не нижнюю границу правдоподобия — тождество перестаёт работать, и это осознанный обмен, а не улучшение.

EM — это тот же ELBO

Обещание из блока 3 (урок 200) можно теперь выполнить. EM — координатный подъём по ELBO:

шаг EM	что делает с ELBO
E-шаг	максимизирует по $q$ при фиксированных параметрах
M-шаг	максимизирует по параметрам при фиксированном $q$

E-шаг доступен в закрытой форме именно потому, что оптимальное $q$ — это истинный posterior $p(z \mid x)$, который для смеси считается правилом Байеса. После него зазор равен нулю, и ELBO совпадает с $\log p(x)$.

Отсюда монотонность EM становится очевидной, а не удивительной: каждый из двух шагов не уменьшает ELBO, а после E-шага ELBO равна правдоподобию — значит правдоподобие тоже не уменьшается.

И отсюда же видно, чем VAE отличается от EM: там posterior недоступен, поэтому E-шаг заменяется на «подвинуть $q$ градиентом в сторону posterior», зазор никогда не закрывается полностью, и монотонности нет.

Что ещё стоит знать

Reparametrisation trick. Чтобы взять градиент ELBO по параметрам $q$, нужно дифференцировать через операцию сэмплирования. Приём из блока 3 (урок 110): $z = \mu + \sigma\varepsilon$, где $\varepsilon \sim \mathcal{N}(0, I)$ не зависит от параметров. После этого $\mathbb{E}_q$ становится обычным ожиданием по $\varepsilon$, и градиент проходит.

ELBO — плохая оценка правдоподобия. Это нижняя граница с неизвестным зазором, поэтому сравнивать по ней модели рискованно: модель с более высокой ELBO может иметь более низкое $\log p(x)$, если её зазор меньше. Для честного сравнения используют importance-weighted оценки (IWAE), которые дают более тугую границу за счёт нескольких сэмплов.

Один сэмпл достаточно. На практике $\mathbb{E}_q$ оценивают одним сэмплом $z$ на пример. Оценка получается шумной, но несмещённой — и этого хватает, ровно по тому же аргументу, по которому хватает одного минибатча в SGD (блок 3, урок 150).