Prism: динамическое разреженное внимание для нативного обучения моделей совместной генерации видео и аудио в 2K

Разбор статьи

Проблема и мотивация

Нативное обучение диффузионных моделей совместной генерации видео и аудио в разрешении 2K позволяет синтезировать богатые пространственные детали, четкую динамику движения и синхронизированный звук. Однако после VAE-сжатия 2K-клип порождает от сотен тысяч до более чем одного миллиона токенов. Стандартное полное внимание (full attention) сталкивается с двумя критическими проблемами:

1. Квадратичная сложность: вычислительные затраты и потребление памяти делают обучение непомерно дорогим.
2. Размытие весов внимания (Attention Weight Dilution): огромное количество повторяющихся фоновых токенов доминирует в последовательности, растрачивая бюджет внимания, предназначенный для информативного, локализованного контента (такого как движущиеся лица, руки и объекты, издающие звуки), и разрушая предобученные мультимодальные априорные распределения (priors).

Существующие подходы к разреженному вниманию либо применяют приближения без дообучения (training-free), либо ориентированы исключительно на немое видео в более низком разрешении (480p/720p), либо используют жесткие, однородные формы блоков. Поскольку аудиовизуальная синхронизация строго сконцентрирована в компактных зонах, производящих звук, фиксированные формы блоков смешивают разнородные токены воедино, искажают представления блоков после mean-pooling и не способны адаптироваться к паттернам кросс-модального взаимодействия.

Ключевые идеи

• Разбиение на макрозоны (Macro-Zone Partitioning): пространственно-временной объем токенов разбивается на макрозоны размером $Z_T\times Z_H\times Z_W=8\times 8\times 8$ для обеспечения конфигураций блоков, выровненных под Triton.
• Индикатор вариации с двойным управлением (Dual-Guidance Variation Indicator): каждая зона динамически измеряет направленную вариацию контента, используя поканальную дисперсию видео наряду с нормами кросс-внимания аудио-к-видео.
• Адаптивные к контенту формы 3D-блоков: зоны выбирают конфигурацию из набора 16 кандидатов с общими объемами $B\in \{64,128,256\}$, выполняя более мелкое разбиение вдоль осей с быстрыми изменениями и интенсивной аудиовизуальной связью.
• Гибридный выбор Top-$k$ / Top-$p$: разреженное внимание на уровне блоков использует Top-$k$ для гарантированного базового контекста, адаптивно расширяемого с помощью Top-$p$ для запросов (queries) с рассеянным распределением внимания.

Метод

Prism работает на базе двухветвевого Diffusion Transformer (DiT), инициализированного из MOVA-720p, заменяя полное самовнимание (self-attention) видео на динамическое разреженное внимание, оставляя кросс-внимание и самовнимание аудио без изменений.

1. Формулирование индикатора вариации

Для каждой макрозоны $m$ коэффициент направленной вариации видео вдоль оси $d\in \{T,H,W\}$ вычисляется на основе усредненных (mean-pooled) признаков value ${\bar{\mathbf{V}}}_{m,i}^{(d)}$: $${\sigma }_{v,d}^2(m)=\frac{1}{D}\sum _{c=1}^D\frac{1}{Z_d}\sum _{i=1}^{Z_d}{\left({\bar{V}}_{m,i,c}^{(d)}-{\mu }_{m,c}^{(d)}\right)}^2,r_{v,d}(m)=\frac{{\sigma }_{v,d}^2(m)}{{\sum }_{d^{\prime }}{\sigma }_{v,d^{\prime }}^2(m)}$$

На основе норм кросс-внимания аудио-к-видео $a_{t,h,w}^m=\Vert {\hat{\mathbf{h}}}_{t,h,w}^m{\Vert }_2$, кэшированных с предыдущего слоя, вычисляются сила аудиосвязи $\bar{a}(m)\in [0,1]$ и нормализованная направленная дисперсия ${\hat{v}}_{a,d}(m)$. Составной индикатор вариации равен: $$g_d(m)=r_{v,d}(m)+\bar{a}(m)\cdot {\hat{v}}_{a,d}(m)+ϵ$$

2. Квантование объема и формы блоков

Общий объем блока $B(m)\in \{64,128,256\}$ назначается путем оценки нормализованной пространственно-временной плотности информации $\rho (m)$ относительно порогов ${\tau }_{128}=0.5$ и ${\tau }_{256}=0.25$. Минимизация потери информации внутри блока ${\sum }_d{g}_d(b_d^2-1)$ при условии ${\prod }_d{b}_d=B(m)$ дает оптимальную непрерывную длину ребра $b_d^{\ast }\propto g_d^{-1/2}$. В пространстве ${\log }_2$ ($S={\log }_{2}B(m)$): $${\ell }_d^{\ast }=-\frac{1}{2}{\log }_2{g}_d+\frac{1}{3}\left(S+\frac{1}{2}{\log }_2{g}_T+\frac{1}{2}{\log }_2{g}_H+\frac{1}{2}{\log }_2{g}_W\right)$$ Форма блока $s^{\ast }=(b_T,b_H,b_W)$ выбирается из набора кандидатов ${\mathcal{C}}_{B(m)}$ путем минимизации евклидова расстояния до ${\ell }_d^{\ast }$.

3. Гибридное внимание Top-$k$ / Top-$p$

Блоки ключей (key) усредняются через mean-pooling в репрезентативные векторы ${\bar{\mathbf{K}}}_j$ и сопоставляются с блоками запросов (query) ${\mathbf{Q}}_i$ для получения вероятностей блоков ${\mathbf{P}}_i$. Множество блоков для внимания ${\mathcal{S}}_i$ объединяет фиксированный и кумулятивный пороги вероятности: $${\mathcal{S}}_i=\text{Top-}k({\mathbf{P}}_i,k)\cup \text{Top-}p({\mathbf{P}}_i,p)$$ Пользовательское ядро Triton с online softmax выполняет разреженное внимание исключительно по множеству ${\mathcal{S}}_i$.

Результаты

• Превосходство на бенчмарках: на 2K-Bench (300 тестовых нативных 2K-видео) Prism достигает эстетического качества (AQ) $0.61$, временного мерцания (TF) $0.982$, DeSync $0.63$ и MotionQ $0.89$, существенно опережая MOVA (Full Attn: $0.42$ AQ, $1.05$ DeSync, $0.53$ MotionQ) и LTX-2.3 с супер-разрешением ($0.48$ AQ, $0.95$ DeSync, $0.66$ MotionQ).
• Прирост эффективности: скорость обучения в $2.5\times$ выше по сравнению с полным вниманием при 2K ($10.6\text{ мин}$ против $26.5\text{ мин}$ на шаг). При инференсе на 4 GPU Prism выполняется за $331\text{ с}$, потребляя $38.0\text{ ГБ}$ VRAM на GPU, по сравнению с $1091\text{ с}$ и $72.8\text{ ГБ}$ у MOVA.
• Масштабирование разрешения: Prism — единственный протестированный фреймворк, демонстрирующий монотонное улучшение по всем метрикам качества и кросс-модальной синхронизации при переходе от 480p к 2K, тогда как полное внимание и бейзлайны с фиксированными блоками деградируют на 2K из-за насыщения фоновыми токенами.
• Пользовательское исследование: в исследовании восприятия 2AFC с участием 200 респондентов на 30 клипах 2K Prism получил $>87\%$ предпочтений по сравнению со всеми открытыми аналогами по соответствию тексту, качеству звука, визуальной детализации, реалистичности движения и синхронизации губ.

Ограничения и открытые вопросы

• Фиксированные временные макрозоны: макрозоны имеют статическое временное окно $Z_T=8$. Быстрая смена кадров или резкие зум-переходы внутри этого окна вызывают резкие расхождения в признаках, которые невозможно адекватно выразить одной формой блока, что приводит к кратковременным артефактам дисторсии.
• Масштабирование разрешения и моделей: Prism оценивался в разрешении до 2K на бэкбоне с 16B параметров; масштабирование до нативного 4K (что учетверяет число токенов) и исследование более крупных архитектур DiT (30B+) остаются открытыми вопросами.
• Порядок аудиовизуальной зависимости: поскольку в блоке DiT самовнимание предшествует кросс-вниманию, управление через аудио-кросс-внимание опирается на нормы, кэшированные с предыдущего слоя, что вносит задержку передачи управляющего сигнала в один слой.