Практика: от кросс-энтропии до VAE

Блок на одной странице

Шесть уроков связаны одним неравенством и одним тождеством. Стоит увидеть их вместе, прежде чем писать код.

$$\htmlData{k=jensen}{\mathbb{E}[\varphi(X)] \ge \varphi(\mathbb{E}[X])} \quad (\varphi \text{ выпукла})$$

результат	как получается из Йенсена
$\text{KL}(p\|q) \ge 0$	$\varphi = -\log$
$H(p, q) \ge H(p)$	$= H(p) + \text{KL}$
$H(p) \le \log n$	$\text{KL}(p\|\text{uniform}) \ge 0$
$I(X;Y) \ge 0$	KL совместного от произведения
$\text{ELBO} \le \log p(x)$	$\varphi = \log$, вогнута, знак обратный

И одно тождество, из которого выводится вся практика:

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

Кросс-энтропия из MLE, на бумаге

Первое упражнение — без кода. Начните с MLE и дойдите до nn.CrossEntropyLoss:

1. запишите правдоподобие категориального: $p(y \mid x) = \prod_c q_c^{[y = c]}$;
2. возьмите логарифм и минус: получится $-\log q_y$, то есть только член истинного класса;
3. усредните по выборке: это $H(\hat p, q)$ для one-hot $\hat p$;
4. подставьте $q = \text{softmax}(z)$ и убедитесь, что градиент равен $q - \hat{p}$.

Четвёртый шаг — тот, что был в блоке 2 (урок 080). Смысл упражнения в том, чтобы увидеть: «кросс-энтропия», «минус log-правдоподобие категориального» и «NLL после log_softmax» — три названия одного объекта, и в PyTorch они буквально живут в разных функциях с одним результатом.

Проверьте на виджете, что при one-hot $p$ кросс-энтропия равна $-\log_2 q_y$ и от остальных $q_c$ не зависит вовсе — выключите все веса $p$, кроме одного, и подвигайте $q$.

цель p (one-hot)

класс 0

1

класс 1

0

класс 2

0

предсказание q

класс 0

0.5

класс 1

0.3

класс 2

0.2

сделать q равномерной подогнать q под p

H(p)

0 бит

H(p, q) — лосс

1 бит

KL(p‖q)

1 бит

KL(q‖p)

бесконечна

H(q)

1.4855 бит

Обратите внимание, что при one-hot цели $H(p) = 0$, поэтому кросс-энтропия равна KL — весь лосс оказывается «зазором». Это единственный случай, когда «минимум лосса равен нулю» действительно верно, и он же объясняет, почему label smoothing меняет картину: сделав цель не жёсткой, вы поднимаете $H(p)$ и вместе с ним минимум лосса.

ELBO на бумаге

Второе упражнение — вывести ELBO двумя путями и убедиться, что получается одно.

Путь 1, через Йенсена. $\log p(x) = \log \mathbb{E}_q\left[\frac{p(x,z)}{q(z)}\right] \ge \mathbb{E}_q\left[\log\frac{p(x,z)}{q(z)}\right]$.

Путь 2, через KL. Начните с $\text{KL}(q \| p(z\mid x)) \ge 0$, подставьте $p(z \mid x) = \frac{p(x,z)}{p(x)}$ и раскройте. Получите то же неравенство и вдобавок точное выражение для зазора.

Второй путь сильнее: он даёт не только границу, но и её погрешность. Проверьте численно на виджете — сумма ELBO и зазора не меняется, куда бы вы ни увели $q$.

веса q(z)

z = a

0.333

z = b

0.333

z = c

0.333

q = posterior q равномерна

log p(x)

-0.91629

ELBO(q)

-1.23909

зазор KL

0.3228

Перплексия руками

Третье упражнение — считать перплексию, не пользуясь библиотекой, на игрушечном корпусе:

1. возьмите текст из сотни слов, постройте биграммную модель со сглаживанием $\alpha = 1$;
2. посчитайте среднюю кросс-энтропию на отложенном тексте в натах, затем перплексию как $e^H$;
3. пересчитайте в биты и убедитесь, что $2^{H_{\text{бит}}} = e^{H_{\text{наты}}}$;
4. уберите сглаживание и получите бесконечность на первом же незнакомом биграмме;
5. сравните перплексию юниграммной и биграммной моделей на одном тексте.

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

VAE на MNIST

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

Что реализовать:

• энкодер выдаёт $\mu(x)$ и $\log \sigma^2(x)$ — логарифм дисперсии, чтобы не требовать положительности;
• репараметризация $z = \mu + \sigma \varepsilon$, $\varepsilon \sim \mathcal{N}(0, I)$;
• KL к стандартному prior в закрытой форме — её стоит вывести самому:

$\text{KL}\big(\mathcal{N}(\mu, \sigma^2) \,\|\, \mathcal{N}(0,1)\big) = \tfrac12\big(\mu^2 + \sigma^2 - 1 - \log \sigma^2\big)$

• реконструкционный член как бинарная кросс-энтропия по пикселям;
• один сэмпл $z$ на пример: оценка шумная, но несмещённая.

Что посмотреть после обучения:

1. нарисуйте латентное пространство. Двумерный $z$, раскрашенный по цифрам, — классы должны разделиться, хотя метки в обучении не участвовали;
2. пройдите по сетке в латенте и декодируйте. Переходы должны быть плавными — это то, чего автоэнкодер без KL-члена не даёт;
3. следите за двумя членами лосса порознь. Их сумма — не то, что нужно смотреть; интересно именно соотношение;
4. воспроизведите posterior collapse. Умножьте KL-член на $\beta = 10$ и убедитесь, что KL уходит в нуль, латент перестаёт зависеть от $x$, а реконструкции становятся одинаковыми;
5. сравните с $\beta = 0$. Это обычный автоэнкодер: реконструкции резче, а сэмплы из prior — мусор, потому что латентное пространство «дырявое».

Пункты 4 и 5 важнее самого обучения. Они показывают, что ELBO — не произвольная сумма двух слагаемых, а равновесие, и что оба его нарушения выглядят как «модель обучилась», если смотреть только на суммарный лосс.

Что должно сойтись

Числовые ориентиры, по которым видно, что всё правильно:

проверка	ожидание
KL-член в начале обучения	близок к нулю ($\mu \approx 0$, $\sigma \approx 1$ при малой инициализации)
KL-член после обучения	десятки нат, растёт с размерностью латента
ELBO	монотонно растёт, но не гладко — оценка по одному сэмплу шумная
$2^{H_{\text{бит}}}$ против $e^{H_{\text{наты}}}$	совпадают до машинной точности
$\beta = 10$	KL $\to 0$, реконструкции вырождаются

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