Генерация высококачественных изображений человеческих лиц долгое время прочно ассоциировалась с генеративно-состязательными сетями (GAN). Однако исследователи из компании NVIDIA решили доказать, что классическая альтернатива в лице вариационных автоэнкодеров (VAE) незаслуженно оставалась в тени из-за давней проблемы с размытостью результатов. Известный ИИ-блогер и исследователь Янник Кильхер (Yannic Kilcher) подробно разобрал революционную статью «NVAE: A Deep Hierarchical Variational Autoencoder», авторы которой объединили стек инженерных и архитектурных находок, чтобы превзойти современное состояние ИИ-архитектур.
🎭 Преодоление проклятия размытости: В чём прорыв NVAE? 0:00
Большинство пользователей привыкли к тому, что фотореалистичные лица генерируются с помощью архитектур GAN. Вариационные автоэнкодеры (VAE) фундаментально отличаются от них по своей математической сути, однако исторически их было крайне тяжело масштабировать до разрешений высокого уровня с сохранением четких, детализированных текстур.
В обсуждаемой работе авторы Араш Вахдат (Arash Vahdat) и Ян Каутц (Jan Kautz) из NVIDIA представили глубокий иерархический вариационный автоэнкодер (NVAE). По словам Янника Кильхера, ценность публикации заключается не в каком-то одном революционном открытии, а в синергии множества известных методов и тонких инженерных оптимизаций. Каждый примененный хак наслаивается друг на друга, в результате чего модель не просто выдает невероятно четкие изображения, но и превосходит существовавший на тот момент SOTA (state-of-the-art) на популярных датасетах.
📊 От классического автоэнкодера к вероятностной модели 1:49
Чтобы понять суть прорыва, необходимо вспомнить базовые принципы работы этих моделей:
- Классический автоэнкодер (Autoencoder): принимает изображение, пропускает его через сеть-энкодер, сжимает в низкоразмерное латентное пространство (вектор $Z$), а затем декодер пытается восстановить исходную картинку. Обучение происходит за счет минимизации ошибки реконструкции (обычно L2-loss) между входом и выходом.
- Вариационный автоэнкодер (VAE): переводит эту задачу в вероятностную, байесовскую плоскость. Вместо прямой генерации фиксированного латентного кода $Z$, энкодер обучается предсказывать параметры распределения — математическое ожидание ($\mu$) и стандартное отклонение ($\sigma$) нормального распределения.
Из этого распределения случайным образом берется сэмпл, который затем отправляется в декодер для воссоздания картинки. Однако без дополнительного контроля система склонна дегенерировать обратно в обычный автоэнкодер: значение $\sigma$ начнет стремиться к нулю, превращая распределение в точечную дельта-функцию Дирака.
Чтобы сбалансировать систему, в функцию потерь вводится регуляризатор — дивергенция Кульбака — Лейблера (KL-divergence), которая заставляет энкодер удерживать предсказываемые распределения как можно ближе к стандартному нормальному распределению.
«Поскольку при обучении из распределения берется случайный сэмпл, а оптимизировать приходится жесткую функцию потерь L2, модели становится выгоднее усреднять значения, делая картинку размытой ради минимизации штрафа», — объясняет Янник Кильхер популярную гипотезу о природе размытости VAE.
🪜 Иерархическая архитектура: Как победить «кукольный» эффект 8:18
Для борьбы с потерей деталей авторы NVAE применили концепцию иерархического латентного пространства, напоминающую прогрессивное масштабирование в GAN. Процесс генерации стартует с грубого случайного шума низкого разрешения (например, размера $8 \times 8$ пикселей). Нейросеть формирует самые общие, базовые очертания объекта.
Далее архитектура начинает послойно наращивать разрешение, используя целую систему взаимосвязанных латентных переменных:
- Модель NVAE в финальной конфигурации использует 36 групп латентных переменных.
- Масштабирование структуры происходит поэтапно в диапазоне от размеров $8 \times 8$ до $128 \times 128$ пикселей.
- На каждую группу латентных переменных приходится по два остаточных (residual) блока.
Верхние слои иерархии отвечают за глобальную структуру (форму головы, ракурс), которая из-за природы VAE получается слегка размытой на своем масштабе. Но затем включаются последующие слои, которые через остаточные связи (residual connections) добавляют новые порции шума и генерируют мелкие детали поверх предыдущих данных.
Тем не менее, Янник Кильхер озвучил интересное критическое наблюдение касательно визуального стиля полученных изображений. По мнению ведущего, несмотря на потрясающую четкость линий и отсутствие грубых пространственных дефектов, лица на сэмплах NVAE кажутся избыточно сглаженными, напоминая фарфоровые маски, глянцевые вырезки из картона или «дизайнерских младенцев». Он предполагает, что это может быть остаточным проявлением работы механизмов сглаживания VAE на определенных масштабах текстуры кожи.
🛠️ Инженерные хаки: Как заставить глубокую VAE работать 14:53
Многократное последовательное сэмплирование в глубоких иерархических сетях делает процесс обучения крайне нестабильным. Главным инструментом стабилизации градиентов стали модифицированные остаточные связи (Residual Connections) как в энкодере, так и в декодере. При переходе между масштабами признаки не заменяются кардинально, а математически складываются.
Специфика внутреннего устройства вычислительных ячеек (residual cells) в NVAE включает в себя следующие элементы:
- Пакетная нормализация (Batch Normalization): применяется строго перед сверточными слоями, что авторы выделили как критически важную деталь для стабилизации.
- Свертки $1 \times 1$: используются для контролируемого увеличения количества каналов.
- Глубинно-разделяемые свертки $5 \times 5$ (Depthwise Separable Convolutions): необходимы для расширения рецептивного поля нейронов (что критично для улавливания дальнодействующих связей в изображении), но без лавинообразного роста числа параметров и потребления памяти.
- Активация Swish и блоки Squeeze-and-Excitation (SE): улучшают нелинейную выразительность сети.
Кроме того, инженерам пришлось пойти на хитрость с пакетной нормализацией. На этапе инференса (генерации картинок) стандартное использование скользящего среднего ухудшало результаты. Авторам пришлось модифицировать параметр импульса (momentum) в Batch Norm, чтобы бегущая статистика быстрее адаптировалась к особенностям батча в VAE.
💾 Оптимизация памяти и укрощение KL-дивергенции 21:21
Поскольку модель содержит огромное количество слоев, авторы внедрили продвинутые методы экономии видеопамяти. В частности, они задействовали вычисления с переменной точностью (Mixed Precision) через специализированные библиотеки NVIDIA. Также был применен метод чекпоинтинга градиентов (Gradient Checkpointing): во время прямого прохода сохраняется лишь одна карта признаков вместо двух, а значения пакетной нормализации и функции активации Swish заново пересчитываются при обратном проходе (backward pass).
«Честно говоря, будущие фреймворки глубокого обучения должны брать подобные вещи на себя автоматически, избавляя исследователей от необходимости писать такие низкоуровневые хаки вручную», — комментирует этот шаг Янник Кильхер.
Другой фундаментальный вызов глубоких HVAE — это укрощение неограниченного роста значений KL-дивергенции, способного разрушить обучение. В иерархической структуре распределение каждого последующего слоя жестко зависит от предыдущего. Если верхний слой немного промахнется, вся цепочка внизу сместится.
Чтобы исправить это, авторы изменили параметризацию: энкодер в NVAE предсказывает не абсолютные значения распределения для текущего слоя, а лишь относительное отклонение ($\Delta\mu$ и $\Delta\sigma$) от априорного распределения (prior), заданного верхним слоем. На базовом нулевом уровне это эквивалентно стандартному подходу, но на глубоких уровнях данный хак радикально снижает чувствительность к накоплению ошибок сэмплирования. Дополнительно стабильность контролируется с помощью спектральной регуляризации верхних сингулярных значений слоев.
📈 Результаты эксперимента и магия фиксации слоёв 26:30
Благодаря собранному комплексу решений, NVAE показал выдающиеся результаты, забрав статус SOTA на четырех из пяти протестированных датасетах. Модель полностью разгромила другие неавторегрессионные архитектуры и вплотную приблизилась к авторегрессионным моделям (которые генерируют изображение медленно, по одному пикселю за раз). NVAE же строит все пиксели одновременно, что дает колоссальный выигрыш в скорости генерации.
В рамках тестирования авторы провели показательный эксперимент по послойной фиксации латентных кодов, наглядно раскрывающий распределение обязанностей в иерархии сети:
- Если зафиксировать случайный вектор на самом верхнем уровне ($8 \times 8$) и изменять шум на нижних, то общая структура лица, геометрия и ракурс остаются абсолютно неизменными.
- Изменения на промежуточных слоях приводят к плавной смене глобальных атрибутов: цвета волос, оттенка кожи, выражения глаз или формы прически.
- Самые нижние слои управляют исключительно микротекстурами и деталями освещения, практически не влияя на анатомию лица.
Как заключает Янник Кильхер, исследование NVIDIA носит выраженный прикладной, инженерный характер. Оно доказывает, что при должном внимании к деталям реализации, оптимизации памяти и стабилизации градиентов, классические подходы вроде вариационных автоэнкодеров способны демонстрировать визуальные результаты, не уступающие некогда безоговорочным лидерам индустрии.
