Prism : Attention parcimonieuse dynamique pour l'entraînement de modèles natifs de génération conjointe vidéo-audio en 2K

Guide de lecture

Problème et motivation

L'entraînement natif de modèles de diffusion conjoints vidéo-audio en résolution 2K permet de synthétiser des détails spatiaux riches, une dynamique de mouvement nette et un son synchronisé. Cependant, après compression par VAE, un extrait 2K génère de plusieurs centaines de milliers à plus d'un million de tokens. L'attention complète standard souffre de deux problèmes majeurs :

1. Complexité quadratique : L'empreinte mémoire et le coût computationnel rendent l'entraînement excessivement coûteux.
2. Dilution des poids d'attention : Les tokens d'arrière-plan, massifs et répétitifs, dominent la séquence, diluant le budget d'attention alloué aux contenus informatifs et localisés (tels que les visages en mouvement, les mains et les objets émetteurs de son) tout en détériorant les priors multimodaux pré-entraînés.

Les approches antérieures d'attention parcimonieuse appliquent des approximations sans entraînement (training-free), se concentrent uniquement sur la vidéo muette à des résolutions inférieures (480p/720p), ou imposent des formes de blocs rigides et uniformes. La synchronisation audio-visuelle étant strictement concentrée dans des régions émettrices de son compactes, des blocs de forme fixe mélangent des tokens hétérogènes, corrompent les représentations moyennes de ces blocs et échouent à s'adapter aux dynamiques d'interaction cross-modales.

Idées clés

• Partitionnement en macro-zones : Le volume spatio-temporel de tokens est partitionné en macro-zones de taille $Z_T\times Z_H\times Z_W=8\times 8\times 8$ pour autoriser des configurations de blocs alignées sur Triton.
• Indicateur de variation à double guidage : Chaque zone mesure dynamiquement la variation directionnelle du contenu à l'aide de la variance par canal des caractéristiques vidéo, combinée aux normes de cross-attention audio-vers-vidéo.
• Formes de blocs 3D adaptées au contenu : Les zones choisissent parmi un ensemble de 16 formes candidates avec des volumes totaux $B\in \{64,128,256\}$, en partitionnant plus finement le long des axes de changement rapide et de fort couplage audio.
• Sélection hybride Top-$k$ / Top-$p$ : L'attention parcimonieuse au niveau des blocs utilise un Top-$k$ pour garantir un contexte de base minimal, étendu de façon adaptative via un Top-$p$ pour les requêtes ayant des distributions d'attention dispersées.

Méthode

Prism opère sur un Diffusion Transformer (DiT) à deux branches initialisé à partir de MOVA-720p, remplaçant la self-attention vidéo complète par une attention parcimonieuse dynamique, tout en conservant intactes la cross-attention et la self-attention audio.

1. Formulation de l'indicateur de variation

Pour chaque macro-zone $m$, le ratio de variation directionnelle de la vidéo selon l'axe $d\in \{T,H,W\}$ est calculé à partir du pooling moyen des valeurs ${\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)}$$

À partir des normes de cross-attention audio-vers-vidéo $a_{t,h,w}^m=\Vert {\hat{\mathbf{h}}}_{t,h,w}^m{\Vert }_2$ mises en cache depuis la couche précédente, l'intensité de couplage audio $\bar{a}(m)\in [0,1]$ et la variance directionnelle normalisée ${\hat{v}}_{a,d}(m)$ sont calculées. L'indicateur composite de variation est : $$g_d(m)=r_{v,d}(m)+\bar{a}(m)\cdot {\hat{v}}_{a,d}(m)+ϵ$$

2. Quantification du volume et de la forme des blocs

Le volume total de bloc $B(m)\in \{64,128,256\}$ est attribué en évaluant la densité normalisée d'information spatio-temporelle $\rho (m)$ par rapport aux seuils ${\tau }_{128}=0.5$ et ${\tau }_{256}=0.25$. Minimiser la perte d'information intra-bloc ${\sum }_d{g}_d(b_d^2-1)$ sous la contrainte ${\prod }_d{b}_d=B(m)$ donne la longueur continue optimale des arêtes $b_d^{\ast }\propto g_d^{-1/2}$. Dans l'espace ${\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)$$ La forme de bloc $s^{\ast }=(b_T,b_H,b_W)$ est choisie dans l'ensemble candidat ${\mathcal{C}}_{B(m)}$ en minimisant la distance euclidienne avec ${\ell }_d^{\ast }$.

3. Attention hybride Top-$k$ / Top-$p$

Les blocs de clés font l'objet d'un pooling moyen en représentants ${\bar{\mathbf{K}}}_j$, puis sont évalués face aux blocs de requêtes ${\mathbf{Q}}_i$ pour produire les probabilités de bloc ${\mathbf{P}}_i$. L'ensemble sélectionné ${\mathcal{S}}_i$ combine un seuil fixe et un seuil de probabilité cumulée : $${\mathcal{S}}_i=\text{Top-}k({\mathbf{P}}_i,k)\cup \text{Top-}p({\mathbf{P}}_i,p)$$ Un kernel Triton à softmax en ligne exécute l'attention parcimonieuse uniquement sur ${\mathcal{S}}_i$.

Résultats

• Supériorité sur les benchmarks : Sur le 2K-Bench (300 vidéos de test natives en 2K), Prism obtient $0,61$ en Qualité Esthétique (AQ), $0,982$ en Scintillement Temporel (TF), $0,63$ en DeSync et $0,89$ en MotionQ, surpassant nettement MOVA (Attention complète : $0,42$ AQ, $1,05$ DeSync, $0,53$ MotionQ) et LTX-2.3 avec super-résolution ($0,48$ AQ, $0,95$ DeSync, $0,66$ MotionQ).
• Gains d'efficacité : La vitesse d'entraînement est $2,5\times$ plus rapide que l'attention complète en 2K ($10,6\text{ min}$ contre $26,5\text{ min}$ par étape). Lors de l'inférence sur 4 GPU, Prism s'exécute en $331\text{ s}$ en utilisant $38,0\text{ Go}$ de VRAM par GPU, contre $1091\text{ s}$ et $72,8\text{ Go}$ pour MOVA.
• Mise à l'échelle de la résolution : Prism est le seul cadre évalué montrant une amélioration monotone sur l'ensemble des métriques de qualité et de synchronisation cross-modale lors du passage de 480p à la 2K, alors que l'attention complète et les méthodes de référence à blocs fixes se dégradent en 2K en raison de la saturation causée par les tokens d'arrière-plan.
• Étude utilisateur : Dans une étude perceptuelle 2AFC menée auprès de 200 participants sur 30 vidéos 2K, Prism a obtenu une préférence supérieure à $87\%$ face à l'ensemble des modèles open-source de référence sur l'alignement texte, la qualité audio, les détails visuels, le réalisme du mouvement et la synchronisation labiale.

Limites et questions ouvertes

• Macro-zones temporelles fixes : Les macro-zones ont une fenêtre temporelle statique de $Z_T=8$. Des changements de plan rapides ou des zooms au sein de cette fenêtre créent des disparités de caractéristiques brutales qu'une forme de bloc unique ne peut capturer, provoquant des artéfacts transitoires de distorsion.
• Montée en échelle de la résolution et des modèles : Prism a été évalué jusqu'à la résolution 2K sur un backbone de 16B paramètres ; le passage à la 4K native (qui quadruple le nombre de tokens) et l'exploration d'architectures DiT plus imposantes (30B+) restent à étudier.
• Ordre de dépendance audio-visuel : La self-attention précédant la cross-attention au sein du bloc DiT, le guidage par cross-attention audio repose sur des normes mises en cache depuis la couche précédente, ce qui introduit un décalage d'une couche dans la propagation du signal de guidage.