Prism: Atención Dispersa Dinámica para el Entrenamiento de Modelos Nativos Conjuntos de Generación de Video y Audio en 2K

Guía de lectura

Problema y motivación

Entrenar de manera nativa modelos de difusión conjuntos de video y audio a resolución 2K permite sintetizar detalles espaciales ricos, dinámicas de movimiento nítidas y sonido sincronizado. Sin embargo, tras la compresión con VAE, un clip en 2K genera desde cientos de miles hasta más de un millón de tokens. La atención completa estándar sufre de dos problemas críticos:

1. Complejidad cuadrática: La huella computacional y de memoria hace que el entrenamiento sea prohibitivamente costoso.
2. Dilución del peso de atención: Masivos tokens de fondo repetitivos dominan la secuencia de tokens, diluyendo el presupuesto de atención asignado al contenido informativo y localizado (como rostros en movimiento, manos y objetos emisores de sonido) y desestabilizando los priors multimodales preentrenados.

Los enfoques previos de atención dispersa aplican aproximaciones sin reentrenamiento, se centran solo en video mudo a resoluciones más bajas (480p/720p) o imponen formas de bloque rígidas y uniformes. Dado que la sincronización audiovisual se concentra estrictamente en regiones compactas generadoras de sonido, los bloques de forma fija mezclan tokens dispares, corrompen las representaciones de bloque obtenidas por promedio (mean-pooling) y no logran adaptarse a los patrones de interacción intermodal.

Ideas clave

• Macro-Zone Partitioning: El volumen espaciotemporal de tokens se particiona en macrozonas de tamaño $Z_T\times Z_H\times Z_W=8\times 8\times 8$ para permitir configuraciones de bloques alineadas con Triton.
• Dual-Guidance Variation Indicator: Cada zona mide dinámicamente la variación direccional de contenido empleando la varianza de canales de video junto con las normas de atención cruzada de audio a video.
• Content-Adaptive 3D Block Shapes: Las zonas eligen entre un conjunto candidato de 16 formas con volúmenes totales $B\in \{64,128,256\}$, particionando finamente a lo largo de los ejes de cambio rápido y de intenso acoplamiento de audio.
• Hybrid Top-$k$ / Top-$p$ Selection: La atención dispersa a nivel de bloque utiliza Top-$k$ para garantizar un contexto base mínimo, expandido adaptativamente mediante Top-$p$ para consultas con distribuciones de atención más dispersas.

Método

Prism opera sobre un Diffusion Transformer (DiT) de doble rama inicializado a partir de MOVA-720p, reemplazando la autoatención completa de video con atención dispersa dinámica mientras mantiene sin cambios la atención cruzada y la autoatención de audio.

1. Formulación del indicador de variación

Para cada macrozona $m$, la proporción de variación direccional de video a lo largo del eje $d\in \{T,H,W\}$ se deriva de las características de valor promediadas ${\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 partir de las normas de atención cruzada de audio a video $a_{t,h,w}^m=\Vert {\hat{\mathbf{h}}}_{t,h,w}^m{\Vert }_2$ almacenadas en caché desde la capa precedente, se calculan la fuerza de acoplamiento de audio $\bar{a}(m)\in [0,1]$ y la varianza direccional normalizada ${\hat{v}}_{a,d}(m)$. El indicador de variación compuesto es: $$g_d(m)=r_{v,d}(m)+\bar{a}(m)\cdot {\hat{v}}_{a,d}(m)+ϵ$$

2. Cuantización de volumen y forma de bloque

El volumen total del bloque $B(m)\in \{64,128,256\}$ se asigna evaluando la densidad de información espaciotemporal normalizada $\rho (m)$ frente a los umbrales ${\tau }_{128}=0.5$ y ${\tau }_{256}=0.25$. Minimizar la pérdida de información intrabloque ${\sum }_d{g}_d(b_d^2-1)$ sujeta a ${\prod }_d{b}_d=B(m)$ produce la longitud de borde continua óptima $b_d^{\ast }\propto g_d^{-1/2}$. En el espacio ${\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 forma de bloque $s^{\ast }=(b_T,b_H,b_W)$ se elige del conjunto de candidatos ${\mathcal{C}}_{B(m)}$ minimizando la distancia euclidiana hacia ${\ell }_d^{\ast }$.

3. Atención híbrida Top-$k$ / Top-$p$

Los bloques clave se promedian mediante mean-pooling en representantes ${\bar{\mathbf{K}}}_j$, y se evalúan frente a los bloques de consulta ${\mathbf{Q}}_i$ para obtener las probabilidades de bloque ${\mathbf{P}}_i$. El conjunto atendido ${\mathcal{S}}_i$ combina umbrales de probabilidad fija y acumulativa: $${\mathcal{S}}_i=\text{Top-}k({\mathbf{P}}_i,k)\cup \text{Top-}p({\mathbf{P}}_i,p)$$ Un kernel Triton con softmax en línea ejecuta la atención dispersa únicamente sobre ${\mathcal{S}}_i$.

Resultados

• Superioridad en benchmarks: En 2K-Bench (300 videos de prueba nativos en 2K), Prism logra $0.61$ en Aesthetic Quality (AQ), $0.982$ en Temporal Flickering (TF), $0.63$ en DeSync y $0.89$ en MotionQ, superando con creces a MOVA (atención completa: $0.42$ AQ, $1.05$ DeSync, $0.53$ MotionQ) y a LTX-2.3 con superresolución ($0.48$ AQ, $0.95$ DeSync, $0.66$ MotionQ).
• Ganancias de eficiencia: La velocidad de entrenamiento es $2.5\times$ más rápida que la atención completa en 2K ($10.6\text{ min}$ frente a $26.5\text{ min}$ por paso). En inferencia con 4 GPUs, Prism se ejecuta en $331\text{ s}$ utilizando $38.0\text{ GB}$ de VRAM por GPU, en comparación con los $1091\text{ s}$ y $72.8\text{ GB}$ de MOVA.
• Escalabilidad de resolución: Prism es el único marco evaluado que muestra mejoras monótonas en todas las métricas de calidad y sincronización intermodal al escalar de 480p a 2K, mientras que la atención completa y las líneas base dispersas de bloque fijo se degradan en 2K debido a la saturación de tokens de fondo.
• Estudio de usuarios: En un estudio perceptual 2AFC con 200 participantes sobre 30 clips en 2K, Prism obtuvo una preferencia superior al $87\%$ sobre todas las líneas base de código abierto en alineación de texto, calidad de audio, detalle visual, realismo de movimiento y sincronización labial.

Limitaciones y preguntas abiertas

• Macrozonas temporales fijas: Las macrozonas tienen una ventana temporal estática de $Z_T=8$. Cortes rápidos de escena o transiciones de zoom dentro de esta ventana introducen disparidades abruptas de características que una sola forma de bloque no puede representar, generando artefactos transitorios de distorsión.
• Escalabilidad de resolución y modelo: Prism se ha evaluado hasta una resolución 2K en una arquitectura base de 16B parámetros; escalar a 4K nativo (lo cual cuadruplica la cantidad de tokens) y explorar arquitecturas DiT mayores a 30B+ siguen siendo retos pendientes.
• Orden de dependencia audiovisual: Debido a que la autoatención precede a la atención cruzada en el bloque DiT, la guía de la atención cruzada de audio depende de normas almacenadas en caché desde la capa anterior, introduciendo un retraso de una capa en la propagación de la señal de guía.