Несколько месяцев назад мы выпустили токенизаторы для видео и изображений — KVAE-1.0. Сегодня представляем следующее поколение: KVAE-2.0 — две новых модели, которые превосходят Wan 2.2 и HunyuanVideo 1.5 по объективным метрикам и качеству генерации и являются, таким образом, лучшими открытыми видео-токенизаторами. Код и веса выложены в открытый доступ.

Напомним контекст: видео-версия KVAE-1.0 сжимала по трём осям — в 8 раз по ширине, в 8 по высоте, в 4 раза по количеству кадров (формат 4x8x8). Одна из новых моделей поддерживает тот же формат 4x8x8, вторая работает в более агрессивном режиме 4x16x16 — в тексте они так и именуются. Мы рассмотрели актуальные тренды в токенизации, описали архитектурные изменения в KVAE-2.0 и сравнили с альтернативами.

Введение

Последние версии открытых решений нацелены на увеличение сжатия: и Hunyuan-Video, и Wan повысили его с 4x8x8 до 4x16x16, изменив количество каналов с 16 до 32 и 48 соответственно. Перенимая опыт DC-AE, NVidia предложила автоэнкодер для сжатия видео в форматах 4x32x32 и 4x64x64, а Snapchat использует аналогичные приёмы (1, 2) для повышения сжатия при переносе модели генерации на мобильные устройства. Тенденцию на повышение игнорирует LTX-Video-2, оставляя степень сжатия автоэнкодера 8x32x32 с первой версии. Небольшой размер латента по пространственным осям компенсируется увеличением количества каналов, что может замедлить сходимость обучения генеративной модели, однако, как было указано в посте, смещение шага сэмплирования способно компенсировать этот эффект. Увеличение пространственного сжатия с 8 до 16 позволяет уменьшить размер патча в генеративной модели с 2x2 до 1x1, полностью перекладывая ответственность за понижение размерности на токенизатор, что, судя по (1, 2), позволяет увеличить качество генерации.

Смежная задача токенизации изображений, будучи менее требовательной к вычислительным ресурсам, чаще становится объектом исследований. Последние работы в этой области позволяют заявить о существовании группы подходов, внедряющих foundation model либо в обучение автоэнкодера, повышая содержание семантической информации в латенте (RAE и follow-up), либо в обучение диффузионной модели, повышая её скорость сходимости (REPA, REPAe). А некоторые работы предлагают полностью отказаться от токенизаторов, обучая модель генерации сквозным методом в пиксельном пространстве (1, 2). Однако эксперименты в соответствующих статьях чаще всего проводятся на изображениях небольшого разрешения, что выглядит неубедительно и отчасти объясняет отсутствие представительства среди крупных решений на рынке генерации.

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

Описание моделей, изменения

Архитектура новых моделей по большей части совпадает с KVAE-1.0 — это всё так же каузальный автоэнкодер, составленный из свёрточных слоёв, нелинейностей и нормализаций. Однако имеется и ряд важных изменений. Например, в прошлой версии обнаружили узкое место в энкодере: в каждом downsample-слое количество каналов латента увеличивалось менее, чем в два раза, что приводило к частичной потере информации. Стоит упомянуть и замену GroupNorm на RMSNorm: в предыдущей версии статистики (среднее, стандартное отклонение) для GroupNorm считались по группе кадров в 16/17 кадров, что вызывало сложности и в обучении, и в инференсе. Теперь же, в рамках RMSNorm, нормализация проводится покадрово, что позволяет варьировать размер входа при кодировании: при увеличении разрешения во избежание OOM достаточно уменьшить длину обрабатываемого сегмента.

Несмотря на то, что количество параметров между энкодером и декодером и так распределено неравномерно (с большей долей у декодера), в новой модели формата 4x16x16 отношение количества весов энкодера к количеству весов декодера увеличено до 5,3. Для сравнения: в видео-модели KVAE-1.0 4x8x8, также как и в KVAE-2.0 4x8x8, эта величина достигает 1,3. Повышение обеспечивается снижением количества каналов в свёрточных слоях энкодера и увеличением их в декодере. Эта особенность, вкупе с файнтюном декодера с дискриминатором на последних шагах обучения, позволяет сохранить высокочастотные детали на реконструкциях.

Отдельное внимание уделили обучающему датасету. Дополнили видео, которые использовали для обучения KVAE-1.0, а весь расширенный набор отфильтровали с упором на высокую динамичность. В дополнение к этому, для KVAE-2.0 увеличили вероятность сэмплирования изображения в батч с 0,2 до 0,3, что положительно сказалось на качестве их восстановления. При обучении новых токенизаторов активно использовали и исследовали оценки diffusability получающегося латентного пространства. Удачное латентное пространство — ключ к качественным генерациям и быстрой сходимости диффузионной модели.

Сравнение

Восстановление

Для теста использовали открытые датасеты MCL-JCV (видео в разрешении 1280x720) и BVI-DVC. В качестве альтернатив для формата 4x8x8 рассматривали Wan-2.1 и HunyuanVideo-1.0. Инференс проводили с помощью реализаций из библиотеки diffusers. Для модели HunyuanVideo, ввиду наличия у неё блока full attention, применяли тайлинг (параметры по-умолчанию). Ниже представлены результаты сравнения по метрикам PSNR, SSIM и LPIPS (с признаками от AlexNet).

Для формата 4x16x16 выбрали последние версии этих же семейств: Wan-2.2 и HunyuanVideo-1.5. Для инференса Wan-2.2 использовали оригинальный репозиторий (ссылка на Wan-2.2). Модель HunyaunVideo-1.5, также как и HunyuanVideo-1.0, запускали с тайлингом.

Прямое сравнение

Цель токенизатора состоит в формировании латентного пространства для генеративной модели, поэтому его превосходство можно установить только с помощью оценки качества генераций. Для этого мы напрямую сравнили моделей (side-by-side, SBS) с участием нескольких пользователей. Каждому показывали пары изображений, созданных по одному и тому же запросу. Люди оценивали каждую пару по трём характеристикам: следование промпту, визуальное и семантическое качество. Достаточно много размеченных пар позволяют установить отношение лучше-хуже между парой моделей. Честность сравнения обеспечивает фиксированный обучающий датасет для генеративной модели, её архитектура, а также стратегия обучения (параметры оптимизатора, количество шагов, размер батча и другие гиперпараметры). Ниже представлены результаты двух SBS с KVAE-2.0 4x16x16:

Видео ниже иллюстрируют последнее side-by-side сравнение: диффузионная модель с KVAE-2.0 чаще превосходит диффузионную модель с токенизатором Hunyuan-1.5, особенно по семантической составляющей.

Hunyuan 1.5:

K-VAE 2.0:

Hunyuan 1.5:

K-VAE 2.0:

Hunyuan 1.5:

K-VAE 2.0:

Оценка генеративного качества токенизатора

Каждый новый эксперимент с токенизатором требует не только его обучения, но и полного обучения генеративной модели. В такой постановке итеративная разработка токенизатора становится крайне затратной. Эта проблема ставит вопрос: можно ли предсказать итоговое качество, минуя долгое обучение генеративной модели, опираясь только на свойства латентного пространства? Для ответа на него мы поискали метрики, сильно коррелирующие с визуальным качеством.

Для агрегации попарных сравнений использовали модель Брэдли-Терри (БТ) — вероятностный фреймворк, который оценивает «силу» каждой модели из матрицы побед и поражений. Метрики качества реконструкции (PSNR, rFID) показывают околонулевую корреляцию с БТ-оценкой визуального качества, в то же время адаптированный CDS (correlation decay slope) из статьи iREPA обеспечивает коэффициент корреляции Пирсона 0,82. В оригинальной статье его использовали для оценки visual foundation models (VFM) в рамках REPA-обучения, и один из выводов авторов такой: качество генерации определяет именно пространственная структура патч-токенов, а не точность классификации, например. Но пространственная карта латентов автоэнкодера аналогична сетке патч-токенов ViT, каждая позиция представляется вектором признаков, что позволяет применить этот подход и в рассматриваемом случае.

PSNR, rFID и CDS против БТ-оценки визуального качества. На графиках представлены flux и flux-2, а также 16-канальный KVAE-1.0 2D (model_3), 32-канальный KVAE без регуляризации (model_5) и 32-канальный KVAE с VF-регуляризацией (model_4), остальные модели представляют внутренние эксперименты с KVAE 2D:

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

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

Заключение

В настоящий момент общее направление продуктовых решений состоит в увеличении степени сжатия токенизаторов. Однако невозможно игнорировать и вызовы с фронтира исследований, призывающих либо расширить обучение автоэнкодеров дистилляцией foundation models (как DiNOv2), либо вообще проводить моделирование полностью в пиксельном пространстве. Дальнейшее развитие должно учитывать обе эти позиции, искать между ними баланс. KVAE-2.0 является шагом на этом пути: модель 4x8x8 есть улучшенная версия KVAE-1.0, а 4x16x16 переходит к более сжатому пространству, где возникает необходимость в модификации модели (как, например, перераспределение весов между энкодером и декодером). Так что следующие модели будут не только обладать большей степенью сжатия, но и лучше моделировать семантическую информацию внутри латента, что непременно повлечёт более качественные генерации.

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