Обновлённый токенизатор видео K-VAE 2.0 от Сбера

Обновлённый токенизатор видео K-VAE 2.0 от Сбера

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

Архитектура и форматы сжатия

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

Современные тенденции в токенизации видео направлены на увеличение степени сжатия. Так, Hunyuan-Video и Wan перешли от формата 4x8x8 к 4x16x16, увеличив количество каналов с 16 до 32 и 48 соответственно. NVidia, вдохновляясь подходом DC-AE, предложила автоэнкодеры с ещё более высоким сжатием — 4x32x32 и 4x64x64. Snapchat использует схожие методы для адаптации моделей под мобильные устройства.

LTX-Video-2 отклоняется от этой тенденции, сохраняя степень сжатия 8x32x32. При этом небольшой размер латента по пространственным осям компенсируется большим количеством каналов, что может замедлить обучение. Однако, как показано в ряде исследований, этот эффект можно компенсировать изменением стратегии сэмплирования.

Увеличение пространственного сжатия с 8 до 16 позволяет снизить размер патча в генеративной модели с 2x2 до 1x1. Это полностью перекладывает ответственность за понижение размерности на токенизатор. Как показывают исследования, такой подход способствует повышению качества генерации.

Тренды в токенизации изображений

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

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

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

Описание моделей и ключевые изменения

Архитектура KVAE-2.0 во многом сохраняет структуру KVAE-1.0: это каузальный автоэнкодер на основе свёрточных слоёв, нелинейностей и нормализаций. Однако внесены важные улучшения.

В предыдущей версии в энкодере наблюдалось узкое место: на каждом этапе downsample количество каналов увеличивалось менее чем вдвое, что приводило к потере информации. Также использовалась GroupNorm, где статистики рассчитывались по группе из 16–17 кадров. Это создавало сложности при обучении и инференсе.

Теперь применена RMSNorm с покадровой нормализацией. Это позволяет гибко менять длину входного сегмента: при увеличении разрешения можно уменьшить длину видео для избежания переполнения памяти (OOM).

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

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

На протяжении обучения активно оценивалась diffusability латентного пространства. Удачное латентное пространство — ключ к качественной генерации и быстрой сходимости диффузионной модели.

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

Для тестирования использовались открытые датасеты MCL-JCV (1280x720) и BVI-DVC. В формате 4x8x8 сравнивались Wan-2.1 и HunyuanVideo-1.0. Инференс выполнялся через библиотеку diffusers. Для HunyuanVideo применялся тайлинг (параметры по умолчанию).

В формате 4x16x16 сравнивались Wan-2.2 и HunyuanVideo-1.5. Инференс Wan-2.2 проводился через оригинальный репозиторий. HunyuanVideo-1.5 также запускался с тайлингом.

Прямое сравнение качества генерации

Качество токенизатора определяется не только метриками восстановления, но и итоговым качеством генерации. Мы провели прямое сравнение (side-by-side, SBS) с участием нескольких пользователей. Каждому показывали пары изображений, сгенерированных по одному промпту разными моделями. Оценка проводилась по трём критериям: соответствие промпту, визуальное и семантическое качество.

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

Результаты SBS с KVAE-2.0 4x16x16 показывают, что диффузионная модель с этим токенизатором чаще превосходит модель с Hunyuan-1.5, особенно по семантическому качеству.

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

Обучение каждого нового токенизатора требует последующего полного обучения генеративной модели — это крайне ресурсоёмко. Возникает вопрос: можно ли предсказать итоговое качество генерации без долгого обучения, опираясь только на свойства латентного пространства?

Для анализа мы использовали модель Брэдли-Терри (БТ) — вероятностный фреймворк, оценивающий «силу» моделей на основе матрицы побед и поражений. Метрики восстановления, такие как PSNR и rFID, показали почти нулевую корреляцию с БТ-оценкой визуального качества.

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

Поскольку латентное пространство автоэнкодера структурно аналогично сетке патч-токенов ViT, этот подход применим и к нашим моделям.

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

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

Заключение

Современные продукты стремятся к увеличению степени сжатия токенизаторов. Однако нельзя игнорировать и исследования, предлагающие расширить обучение автоэнкодеров за счёт foundation models (например, DiNOv2) или вообще отказаться от токенизаторов, работая в пиксельном пространстве.

KVAE-2.0 — шаг к синтезу этих подходов. Модель 4x8x8 — улучшенная версия KVAE-1.0. Модель 4x16x16 переходит к более сжатому пространству, что требует архитектурных изменений, таких как перераспределение весов между энкодером и декодером.

Будущие модели будут не только сильнее сжимать данные, но и лучше передавать семантическую информацию в латентном пространстве. Это неизбежно приведёт к более качественной генерации видео.

Читать оригинал