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

Обычный автокодировщик переводит вход $x$ в одну точку $z$ и обратно. Реконструирует он хорошо, но генератором не является: возьмите случайную точку латентного пространства — и у декодера нет причин выдать что-то осмысленное.

Вариационный автокодировщик (Kingma and Welling, 2014) упорядочивает латентное пространство. Каждый вход кодируется в гауссиану, а не в точку:

и каждая такая гауссиана подтягивается к одному фиксированному априорному распределению $p(z) = \mathcal N(0, I)$. В коде энкодер заканчивается двумя «головами»: одна для $\mu$, другая для $\log\sigma^2$.

Обучение максимизирует ELBO, а его слагаемое реконструкции — это матожидание по кодам:

Оцениваем его одним сэмплом на вход: берём $z \sim \mathcal N(\mu, \sigma^2)$, декодируем, сравниваем с $x$. Функция потерь явно зависит от $\mu$ и $\sigma$ — сдвинете $\mu$, и коды сдвинутся вместе с ним.

Но в вычислительном графе дыра. Шаг «вытащить случайное число из $\mathcal N(\mu, \sigma^2)$» возвращает значение без производной: градиент функции потерь доходит до $z$ и останавливается. Декодер учится, а энкодер, задача которого — выдавать $\mu$ и $\sigma$, от слагаемого реконструкции не получает ничего.

Решение — перенести случайность во вход графа. Сэмплируем шум, который не имеет отношения к сети, а затем сдвигаем и масштабируем его:

Распределение то же, что и раньше: сдвиг стандартной гауссианы на $\mu$ и масштаб на $\sigma$ дают ровно $\mathcal N(\mu, \sigma^2)$. Модель не изменилась — изменилось только место, где живёт случайность.

Энкодер выдаёт $\log\sigma^2$, а не $\sigma$, поэтому перед сдвигом и масштабом $\sigma$ нужно из него восстановить. В PyTorch шум даёт torch.randn_like — он сэмплирует $\varepsilon$ той же формы, что и аргумент.

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

Градиенты слагаемого реконструкции проходят через $z$ обратно в энкодер (Kingma and Welling, 2014, §2.4). Autograd сделает это сам — важно, что теперь ему есть что дифференцировать.

Есть и менее заметный выигрыш. Универсальная альтернатива — оценка через градиент логарифма плотности (score-function, REINFORCE) — тоже даёт несмещённые градиенты без репараметризации, но её дисперсия обычно гораздо выше. Репараметризованная оценка использует производную самой функции потерь, а она несёт куда больше информации на один сэмпл.

Почему энкодер предсказывает $\log\sigma^2$, а не $\sigma$? Стандартное отклонение обязано быть положительным, а последний линейный слой сети может выдать любое число. Предсказание логарифма снимает ограничение бесплатно:

Заодно очень маленькие дисперсии остаются в удобном численном масштабе. Частая альтернатива — softplus; сам трюк работает с любой положительной $\sigma$.

Вторая половина функции потерь, $D_{\mathrm{KL}}\big(q_\varphi(z \mid x) \,\|\, p(z)\big)$, вообще не требует сэмплирования. Между двумя гауссианами у неё есть замкнутая форма; в сумме по $d$ латентным координатам

Она равна нулю ровно тогда, когда $\mu_j = 0$ и $\sigma_j = 1$ во всех координатах, то есть когда апостериорное распределение совпадает с априорным. В коде это одна строка, и она дифференцируется по $\mu$ и $\log\sigma^2$ напрямую.

Всё вместе — один шаг обучения:

mu, logvar = encoder(x)
std = torch.exp(0.5 * logvar)
eps = torch.randn_like(std)
z = mu + std * eps
x_hat = decoder(z)
 
recon = ((x - x_hat) ** 2).sum(dim=1)
kl = -0.5 * (1 + logvar - mu**2 - logvar.exp()).sum(dim=1)
loss = (recon + kl).mean()
loss.backward()

Квадратичная ошибка соответствует гауссовскому декодеру с фиксированной дисперсией $s^2 = \tfrac12$ — тогда это отрицательный ELBO с точностью до константы. Другая фиксированная дисперсия меняет вес KL-слагаемого — на этой идее построен β-VAE.