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

VAE кодирует объект не точкой, а распределением в latent space. Реконструкция просит сохранить детали, KL-дивергенция — организовать коды около простого prior. Именно их спор делает пространство пригодным и для восстановления, и для новых samples.

Почему обычный autoencoder плохо сэмплируется

Encoder $f(x)=z$ сжимает изображение, decoder $g(z)\approx x$ восстанавливает. Если обучать только reconstruction loss, коды training objects могут образовать разрозненные острова. Сэмпл из промежутка попадёт в область, которую decoder не видел.

Variational autoencoder задаёт вероятностный encoder:

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

prior

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

и likelihood decoder $p_\theta(x\mid z)$.

Мы хотим максимизировать $\log p_\theta(x)$, интегрируя неизвестный latent:

$$p_\theta(x)=\int p_\theta(x\mid z)p(z)\,dz.$$

Интеграл обычно недоступен, поэтому появляется нижняя граница.

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

ELBO по шагам

Вставим произвольный $q_\phi(z\mid x)$:

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

KL неотрицательна, значит

$$\log p_\theta(x)\ge \underbrace{ \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) }_{\mathrm{ELBO}}.$$

Первое слагаемое — качество реконструкции в вероятностной модели. Второе штрафует отличие posterior-кода от prior. Максимизация ELBO одновременно обучает decoder и приближённый inference.

Для diagonal Gaussian KL имеет закрытую форму:

$$D_{\mathrm{KL}}(q\|p) =\frac12\sum_j \left( \mu_j^2+\sigma_j^2-\log\sigma_j^2-1 \right).$$

Reparameterization trick

Нельзя обычным backprop пройти через случайную операцию $z\sim\mathcal N(\mu,\sigma^2)$, если sampling скрывает зависимость от параметров. Перепишем:

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

Случайность теперь находится в независимом $\varepsilon$, а $z$ дифференцируем по $\mu,\sigma$. Это pathwise gradient estimator.

На inference для реконструкции можно использовать $\mu$, для генерации — $z\sim p(z)$. Эти режимы проверяют разные свойства.

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

Лаборатория latent space

Меняйте вес KL. При нуле коды становятся островами и reconstruction резкая, но prior samples плохи. При слишком большом весе коды похожи на prior, а изображения теряют детали.

Выберите два объекта и интерполируйте $z(t)=(1-t)z_1+tz_2$. Смотрите не только на плавность картинки, но и на плотность prior вдоль пути. В высокой размерности линейный отрезок может проходить через нетипичную область; spherical interpolation иногда лучше.

$\beta$-VAE и rate–distortion

Обобщённый objective:

$$\mathcal L_\beta= -\mathbb E_q\log p_\theta(x\mid z) +\beta D_{\mathrm{KL}}(q\|p).$$

Reconstruction term — distortion, KL — rate: сколько информации код несёт об $x$ относительно prior. Большая $\beta$ ограничивает канал и иногда делает факторы более disentangled, но ухудшает детали.

Disentanglement оценивают на данных с известными generative factors: угол, размер, положение. Красивое плавное изменение одной координаты не достаточно; нужно измерять независимость факторов и устойчивость к seed.

Posterior collapse

Сильный autoregressive decoder может предсказывать $x$ без $z$. Тогда

$$q_\phi(z\mid x)\approx p(z), \qquad D_{\mathrm{KL}}\approx0,$$

и latent ничего не кодирует. ELBO может быть хорошей, а representation бесполезным.

Диагностика: KL по координатам, mutual-information proxy, изменение output при перестановке $z$, число active units. KL annealing постепенно увеличивает вес, free bits разрешают каждой координате небольшую бесплатную информацию.

Это пример того, как оптимизатор честно находит обходной путь. Архитектура decoder и расписание objective входят в постановку.

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

Fashion-MNIST и честное сравнение

На Fashion-MNIST можно взять latent dimension 2 для прямой карты, затем 16 для качества. Для Bernoulli likelihood пиксели должны трактоваться соответственно; для непрерывных grayscale данных Gaussian или discretized likelihood может быть осмысленнее. MSE — это предположение о шуме, не нейтральная мера.

Сравните:

• reconstruction test objects;
• random prior samples;
• latent traversal;
• KL по координатам;
• nearest train neighbors;
• linear probe класса по $\mu$;
• negative ELBO на test.

Не выбирайте только удачные samples. Используйте фиксированную сетку quantiles prior.

VAE, GAN и diffusion

GAN оптимизирует adversarial качество samples без явного encoder-а и likelihood. VAE даёт encoder и probabilistic objective, но pixel likelihood часто ведёт к размытию. Diffusion тоже использует Gaussian noise, однако строит длинную цепь denoising в пространстве данных.

VAE часто служит latent compressor внутри diffusion systems. Тогда качество reconstruction определяет потолок деталей, а меньший latent удешевляет denoising.

Мини-исследование: prior sample и aggregate posterior

Хотя каждый $q(z\mid x)$ штрафуется к $\mathcal N(0,I)$, смесь

$$q(z)=\frac1N\sum_iq(z\mid x_i)$$

не обязана точно совпадать с prior. Возьмите двумерный VAE, нанесите $\mu(x)$, samples из aggregate posterior и независимые samples prior. Сравните radial distribution, covariance и occupancy сетки.

Декодируйте три группы точек: typical prior, области высокой $q(z)$ и «дыры», где prior mass есть, а encoded data почти нет. Если в дырах outputs плохи, средний KL скрывает mismatch. Можно fit более гибкий prior или увеличить регуляризацию, но это снова меняет rate–distortion.

Проведите two-sample test classifier-ом, различающим $z\sim q(z)$ и $z\sim p(z)$. Accuracy около 0,5 означает, что простой classifier не нашёл различия, а не доказательство равенства. Сопоставьте с качеством generated samples.

Мини-исследование: что именно делает изображение размытым

На synthetic dataset создайте две возможные позиции тонкой линии при одинаковом входном условии. Обучите deterministic MSE decoder: optimum проведёт две слабые линии или одну среднюю. Затем probabilistic decoder с latent может выбрать одну из мод.

Измерьте pixel MSE, sharpness и coverage обеих позиций. MSE-модель может победить по своей метрике и проиграть по реалистичности. Это пример различия loss и предметной метрики.

Повторите с Bernoulli и Gaussian likelihood. Формула observation model задаёт, что считается естественной ошибкой. Когда VAE называют «размытым», полезно уточнить: ограничена ли архитектура, слишком силён KL или likelihood усредняет мультимодальное будущее.

Представление для downstream-задачи

Заморозьте encoder и обучите linear probe класса на $\mu$. Затем обучите probe на nuisance-признак, например положение. Хороший generative latent не обязан быть class-invariant. При разных $\beta$ class signal может сначала очищаться, затем исчезать вместе со всей информацией.

Сравните с контрастивным embedding, где objective явно выбирает invariances через пары. VAE сохраняет то, что нужно реконструкции, а не то, что удобно classifier-у.

Мини-исследование: uncertainty реконструкции

Для одного $x$ сэмплируйте сто $z\sim q(z\mid x)$ и декодируйте. Посчитайте pixel-wise variance и variation предметных признаков. Высокая variance на границе объекта может означать uncertainty положения, но diagonal Gaussian posterior не гарантирует калибровку.

Сравните reconstruction из $\mu$ и среднее ста stochastic reconstructions. Из-за нелинейного decoder

$$G(\mathbb E z)\ne\mathbb E G(z).$$

Покажите разницу численно. «Средний код» и «средняя картинка» являются разными объектами.

Мини-исследование: latent arithmetic под контролем

Найдите среднее направление между двумя размеченными группами, например рукав/без рукава, и примените его к независимым объектам. Измерьте target classifier и сохранение других признаков. Если направление работает только на training examples, это переобученная геометрия.

Повторите после случайного orthogonal rotation latent и переобучения только линейного направления. Функция генератора та же при согласованном преобразовании, а координатная история меняется. Это хороший антидот к буквальным названиям отдельных нейронов.

Для conditional VAE публикуйте две сетки samples: с равным числом объектов каждого класса и с естественными частотами. Они отвечают на разные вопросы о качестве и покрытии.

Проверьте prior predictive до обучения decoder: samples из prior через случайную сеть бессмысленны, а после обучения должны соответствовать диапазону данных. Но хороший prior predictive не гарантирует хорошие reconstructions. Две панели отвечают на противоположные направления модели.

Для условного VAE держите label balance при sampling. Иначе редкий класс может выглядеть слабым из-за малого evaluation sample, а не posterior. Публикуйте равные сетки и естественные доли отдельно.

Полезный latent живёт между крайностями

VAE превращает compression в probabilistic inference. ELBO раскладывает задачу на reconstruction и согласование с prior, reparameterization проводит gradient через sampling, $\beta$ управляет rate–distortion. Успех нельзя измерить одной реконструкцией: нужны prior samples, активность latent и проверка collapse.

Задачи