K-VAE токенизатор от Сбера

Одни из самых перспективных решений среди генеративных моделей принадлежат к классу LDM (latent diffusion models): Stable Diffusion, Wan, Flux и многие другие. Превращение шума в изображение происходит в два этапа: сперва диффузионная модель создаёт латент, который затем декодируется в изображение намного большего размера. За декодирование и формирование латентного пространства в большинстве моделей этого класса отвечает токенизатор изображений. Таким образом, качество этой модели напрямую влияет на результаты генерации.

В Сбере, в Управлении базовых моделей Kandinsky были разработаны токенизаторы K-VAE как для изображений, так и для видео, превосходящие state-of-the-art аналоги как по объективным метрикам (PSNR), так и по качеству генерации.

Ниже будут даны некоторые подробности решения, которое будет полноценно представлено общественности в рамках AIJourney 2025.

2d-VAE

Большинство современных токенизаторов построены на основе вариационных автоэнкодеров (VAE). Энкодер VAE состоит из последовательно применяемых свёрточных слоёв, слоёв активации и нелинейностей.

Непосредственно сжатие происходит в downsample-блоках (разжатие в декодере устроено симметрично), где одновременно сокращается (обычно в два раза) и высота, и ширина, но увеличивается количество каналов. При этом последний слой энкодера уменьшает количество каналов в латенте до целевого.

Разработанное нами решение K-VAE, также как и Flux, формирует 16-канальный латент, у которого и высота, и ширина в 8 раз меньше исходных.

По-умолчанию в энкодере и декодере на самых глубоких уровнях присутствует хотя бы один attention-блок, но разработанное решение намеренно построено исключительно на свёрточных слоях, что объясняется желанием сохранить баланс между сложностью и качеством. И в отличие, например, от Wan, внутри декодера используются улучшающие стабильность обучения modulated-нормализации, аналогично Sber-MOVQGAN.

Стандартными функциями потерь для обучения являются MAE и LPIPS, а также GAN-loss, который можно вводить как с первых шагов, так и на более поздних этапах. Веса этих слагаемых формируют динамику обучения и определяют соотношение между попиксельными и перцептивными метриками при дальнейшем применении.

Несмотря на то, что модель обучается восстанавливать изображения, финальная цель токенизатора — упрощать работу генеративной модели после него. Так что по-настоящему хотелось бы добиться высокого качества генерации диффузионной модели в сформированном латентном пространстве. Сложность задачи заключается в том, что точных критериев «пригодности» латентного пространства нет, хотя список статей с гипотезами постоянно пополняется.

Тем не менее, регуляризация в виде KL-дивергенции на практике приводит к удовлетворительному качеству, а из-за своей простоты получила широкое распространение среди методов токенизации. Важно, что введение регуляризации заведомо портит качество реконструкций, так что на этапе обучения VAE приходится мириться с небольшой просадкой метрик ради дальнейшего обучения диффузионной модели.

Сравнение с альтернативами:

Можно видеть, что DiT, обученный с 2D-K-VAE, заметно опережает DiT, обученный с Flux-VAE по метрике CLIP, что можно интерпретировать как более качественный синтез с точки зрения семантики изображений. 

3d-VAE

В то время как 2d-VAE предназначен для работы с изображениями, 3d-VAE сжимает видео. Соответствующие архитектуры обычно заимствуют у успешных 2d-VAE, заменяя Conv2d на Conv3d, а также добавляя сжатие/разжатие по временной оси.

Техническая проблема обучения и инференса состоит в слишком большом размере тензора: 10-секундное FullHD видео при 24 FPS содержит около 1,5 миллиарда элементов (а в промежуточных представлениях внутри сети — многократно больше). Самый распространённый способ решения — tiling, когда исходные кадры нарезают на пересекающиеся сегменты (тайлы), каждый из которых обрабатывают независимо.

В текущей реализации 3d-VAE для обучения и декодинга адаптировали подход с кешированием последних кадров, предложенный CogVideoX и Wan, что увеличивает эффективный размер контекста и потенциально улучшает сжатие. Так что оказывается возможным проводить инференс и тайлами, и целиком по всему видео.

Сравнение с альтернативами:

Перспективы направления

Разработка VAE балансирует между тремя компонентами, влияющими друг на друга:

• Качество реконструкций: оно зависит от обучающего датасета, архитектуры модели и функции потерь. Распространено мнение, что качество восстановления VAE исходного изображения или видео является верхней оценкой визуального качества генерирования. Но при этом широко используемая стратегия обучения (MAE + LPIPS + GAN) может вносить визуальные артефакты. Эта неприятная особенность приводит к увеличению количества статей, в которых вводят диффузионный декодер в качестве альтернативы (FlowMo, SSDD).

• Степень сжатия: её увеличение ведёт к ускорению обучения диффузионной модели, ведь сложность вычисления attention’а квадратично зависит от длины последовательности. Следовательно, увеличивая сжатие в два раза по высоте и в два раза по ширине, можно сократить количество операций в 4 раза (при прочих равных). В статье группа исследователей из NVidia показывает, во-первых, что повышение степени сжатия с пропорциональным увеличением количества выходных каналов энкодера снижает качество восстановления (таблица 2 в статье), а во-вторых, что это повышение требует модификаций архитектуры и, в частности, изменения downsample/upsample-блоков.

• Качество генераций диффузионной модели: его можно оценить только по завершении обучения автоэнкодера, уже на этапе обучения диффузионной модели. Качество напрямую связано со свойствами латентного пространства, для оценки которых пока нет достоверных подходов. При этом наращивание способности к восстановлению может негативно сказываться на способности к генерации, как было показано в статье.

Предложенное решение соперничает с лучшими open source-методами. Однако дальнейшее его улучшение требует инноваций, учитывая особенности задачи. А улучшение необходимо, ведь первый, получивший технологию более эффективного сжатия медиа для генеративных сетей, получит, как минимум, значительные преимущества в скорости своих генеративных сетей и явные экономические преимущества.

Коллектив авторов:

Денис Пархоменко, Кирилл Чернышев, Андрей Шуткин, Илья Васильев, Иван Кириллов, Денис Димитров, Дмитрий Михайлов.

Источник: habr.com