Variational Lossy Autoencoder: когда VAE перестаёт учить латент

Проблема, ради которой придумали VLAE

Вы обучаете вариационный автоэнкодер, подключаете к нему сильный авторегрессионный декодер вроде PixelCNN — и обнаруживаете, что латентный код z почти не используется. Модель генерирует нормальные картинки, значение регуляризатора KL падает почти в ноль, а всё представление данных живёт в весах декодера. Латент, ради которого VAE и затевался, оказался мёртвым грузом.

Этот эффект называют posterior collapse — коллапс апостериорного распределения. Работа Variational Lossy Autoencoder (Chen et al., 2016) разбирает его причину через теорию информации и предлагает не бороться с ней костылями, а спроектировать модель так, чтобы латент осознанно кодировал одни свойства данных и игнорировал другие. Отсюда слово lossy — с потерями: латент хранит только ту часть информации, которую вы решили в него поместить.

Почему латент схлопывается

ELBO — целевая функция VAE — состоит из двух слагаемых: правдоподобия реконструкции и KL-дивергенции между приближённым апостериором q(z|x) и приором p(z). Оптимизатор минимизирует суммарную стоимость описания данных в битах.

Если декодер p(x|z) настолько выразителен, что способен смоделировать распределение данных сам по себе, без всякого z, то оптимальная стратегия — вообще не тратить биты на латент. KL-слагаемое штрафует любое отклонение q(z|x) от приора, и модель обнуляет его, делая апостериор равным приору. Латент становится шумом, не зависящим от входа.

Коллапс латента — не баг обучения, а рациональный выбор модели: если декодер справляется сам, платить биты за латентный код невыгодно.

Идея Bits-Back и что она объясняет

Авторы формулируют ключевое наблюдение через принцип Bits-Back Coding¹. Стоимость кодирования одного примера через латентную переменную складывается так: вы платите за передачу z, но часть битов возвращается обратно, потому что z сам по себе содержит информацию, которую можно «вернуть» получателю.

Вывод из этого анализа: латент будет использоваться только для той информации, которую дешевле закодировать через z, чем оставить на откуп декодеру. Если авторегрессионный декодер моделирует локальные зависимости между соседними пикселями практически бесплатно, латент не станет их дублировать. А значит, если вы хотите, чтобы z хранил глобальную структуру, нужно намеренно ограничить декодер так, чтобы локальные детали он ловил, а глобальные — нет.

Lossy code как инструмент, а не побочный эффект

Именно здесь VLAE переворачивает постановку. Вместо того чтобы заставлять латент кодировать всё подряд, авторы предлагают:

• Выбрать декодер с ограниченным рецептивным полем — например, PixelCNN, который видит только небольшую окрестность вокруг текущего пикселя.
• Локальные, «текстурные» детали изображения такой декодер моделирует сам.
• Глобальную структуру — что за объект, его расположение, общую композицию — декодер локально ухватить не может, поэтому эта информация вынуждена уходить в латент.

Результат: латент получает осмысленное, глобальное представление, а не схлопывается. Потеря информации (lossy) становится проектным решением — вы контролируете, что именно теряется в латентном описании.

Автогрессивный приор вместо простого гауссиана

Второй вклад работы — замена стандартного приора N(0, I) на обучаемый авторегрессионный приор. Обычно в VAE приор фиксирован и прост, а вся выразительность ложится на апостериор. Авторы показывают, что перенос части сложности в приор через инверсный авторегрессионный поток (IAF) даёт более гибкое латентное пространство без роста стоимости KL.

Формально приор моделируется так, что случайные величины связаны авторегрессионной зависимостью — распределение каждой компоненты z зависит от предыдущих. Это позволяет апостериору оставаться относительно простым, а совокупной модели — точнее приближать истинное распределение данных.

Сравнение подходов к коллапсу латента

Почему это до сих пор актуально

VLAE появился до бума диффузионных моделей, но его идея пережила смену поколений. Мысль «спроектируй информационное бутылочное горлышко под то, что тебе нужно» лежит в основе многих современных подходов к обучению представлений. Когда вы обучаете латентную диффузионную модель и решаете, что кодировать в латент VAE, а что оставить декодеру-U-Net, вы решаете ровно ту же задачу, что сформулировал VLAE: как распределить информацию между стохастическим кодом и мощной детерминированной сетью.

Практический вывод для инженера прост. Если ваш VAE игнорирует латент, первым делом посмотрите не на гиперпараметры обучения, а на баланс мощности: слишком сильный декодер съедает латент независимо от расписания KL. Иногда правильнее ослабить или структурно ограничить декодер, чем крутить веса регуляризатора.

¹ Bits-Back Coding — схема сжатия, при которой отправитель кодирует данные через латентную переменную и «возвращает» себе часть битов за счёт информации, содержащейся в самой латентной переменной. Даёт теоретическую нижнюю границу стоимости описания через ELBO.