VAE: вход и выход в латентное пространство

Урок разбирает первый компонент Stable Diffusion — VAE, который связывает мир пикселей и мир латентных кодов.

VAE (вариационный автоэнкодер) в Stable Diffusion — пара сетей: энкодер сжимает картинку в латентный код, декодер восстанавливает картинку из кода.

Зачем нужен VAE

Мы выяснили: диффузия в Stable Diffusion идёт не по пикселям, а по компактному латентному коду. Кто-то должен переводить картинку в этот код и обратно — это и есть работа VAE. Без него пришлось бы расшумлять сотни тысяч пикселей напрямую, что в разы дороже.

картинка 512x512x3   --энкодер VAE-->  латент 64x64x4
(~786 000 чисел)                       (~16 000 чисел)

...диффузия работает здесь, в латенте...

латент 64x64x4       --декодер VAE-->  картинка 512x512x3

Две роли в пайплайне

В text-to-image VAE нужен лишь в конце: диффузия рождает готовый латентный код, а декодер VAE превращает его в итоговую картинку. А вот в img2img (когда мы стартуем с готового изображения) VAE нужен и в начале — энкодер кодирует исходную картинку в латент, который мы потом частично зашумляем. Про img2img будет отдельный раздел.

Как работает под капотом

Энкодер — это свёрточная сеть, которая постепенно уменьшает разрешение и увеличивает число каналов смысла, сжимая картинку примерно в 8 раз по каждой стороне. Декодер делает обратное: наращивает разрешение, восстанавливая пиксели. VAE обучают отдельно, до диффузии, на задаче «сожми и восстанови как можно точнее». Поскольку сжатие с потерями, очень мелкие детали (например, ровный текст на вывесках) восстанавливаются хуже — это известное слабое место.

Частые ошибки

Считать, что VAE генерирует картинку. Нет: генерирует диффузия (U-Net), а VAE только переводит между пикселями и латентом.
Винить «модель» за кривой мелкий текст. Часто виноват именно VAE-декодер: латент не хранит достаточно деталей для чёткого шрифта.
Думать, что разные VAE взаимозаменяемы. Замена VAE-декодера может заметно изменить цвета и резкость результата.

Почему VAE обучают заранее

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

Практическое следствие — VAE-декодеры взаимозаменяемы лишь отчасти. Сообщество выпускает улучшенные декодеры, которые дают более чистые цвета или лучше справляются с мелкими деталями вроде глаз и текста. Подменив декодер, вы меняете «последний штрих» — то, как латент превращается в пиксели, — не трогая саму генерацию. Это дешёвый способ заметно поднять визуальное качество готовых картинок.

пиксели  <--VAE-декодер--  латент  <--диффузия--  шум
              (последний                (основная
               штрих цвета)             творческая работа)

Итог

VAE — мост между пикселями и латентным пространством: энкодер сжимает, декодер восстанавливает.
В text-to-image нужен лишь декодер (в конце); в img2img — ещё и энкодер (в начале).
Сжатие с потерями объясняет слабые места вроде нечёткого мелкого текста.