Prism: Dynamic Sparse Attention for Native 2K Joint Video-Audio Generation Model Training

Reading guide

Problem and motivation

Natively training joint video-audio diffusion models at 2K resolution enables the synthesis of rich spatial details, sharp motion dynamics, and synchronized sound. However, after VAE compression, a 2K clip yields hundreds of thousands to over one million tokens. Standard full attention suffers from two fatal issues:

1. Quadratic complexity: The computational and memory footprints make training prohibitively expensive.
2. Attention weight dilution: Massive repetitive background tokens dominate the token sequence, diluting the attention budget assigned to informative, localized content (such as moving faces, hands, and sound-producing objects) and disrupting pretrained multimodal priors.

Prior sparse attention approaches either apply training-free approximations, focus only on silent video at lower resolutions (480p/720p), or enforce rigid, uniform block shapes. Because audio-visual synchronization is strictly concentrated in compact sound-producing regions, fixed block shapes mix disparate tokens together, corrupt mean-pooled block representations, and fail to adapt to cross-modal interaction patterns.

Key ideas

• Macro-Zone Partitioning: The spatiotemporal token volume is partitioned into macro-zones of size $Z_T\times Z_H\times Z_W=8\times 8\times 8$ to permit Triton-aligned block configurations.
• Dual-Guidance Variation Indicator: Each zone dynamically measures directional content variation using video channel-wise variance alongside audio-to-video cross-attention norms.
• Content-Adaptive 3D Block Shapes: Zones choose from a 16-shape candidate set with total volumes $B\in \{64,128,256\}$, partitioning finely along axes of rapid change and intense audio coupling.
• Hybrid Top-$k$ / Top-$p$ Selection: Block-level sparse attention uses Top-$k$ to ensure a guaranteed baseline context, adaptively expanded via Top-$p$ for queries with dispersed attention distributions.

Method

Prism operates on a dual-branch Diffusion Transformer (DiT) initialized from MOVA-720p, replacing full video self-attention with dynamic sparse attention while keeping cross-attention and audio self-attention unmodified.

1. Variation Indicator Formulation

For each macro-zone $m$, the directional video variation ratio along axis $d\in \{T,H,W\}$ is derived from mean-pooled value features ${\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)}$$

From the audio-to-video cross-attention norms $a_{t,h,w}^m=\Vert {\hat{\mathbf{h}}}_{t,h,w}^m{\Vert }_2$ cached from the preceding layer, the audio coupling strength $\bar{a}(m)\in [0,1]$ and normalized directional variance ${\hat{v}}_{a,d}(m)$ are computed. The composite variation indicator is: $$g_d(m)=r_{v,d}(m)+\bar{a}(m)\cdot {\hat{v}}_{a,d}(m)+ϵ$$

2. Block Volume and Shape Quantization

Total block volume $B(m)\in \{64,128,256\}$ is assigned by evaluating normalized spatiotemporal information density $\rho (m)$ against thresholds ${\tau }_{128}=0.5$ and ${\tau }_{256}=0.25$. Minimizing intra-block information loss ${\sum }_d{g}_d(b_d^2-1)$ subject to ${\prod }_d{b}_d=B(m)$ yields the optimal continuous edge length $b_d^{\ast }\propto g_d^{-1/2}$. In ${\log }_2$ space ($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)$$ The block shape $s^{\ast }=(b_T,b_H,b_W)$ is chosen from candidate set ${\mathcal{C}}_{B(m)}$ by minimizing Euclidean distance to ${\ell }_d^{\ast }$.

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

Key blocks are mean-pooled into representatives ${\bar{\mathbf{K}}}_j$, scored against query blocks ${\mathbf{Q}}_i$ to yield block probabilities ${\mathbf{P}}_i$. The attended set ${\mathcal{S}}_i$ combines fixed and cumulative probability thresholds: $${\mathcal{S}}_i=\text{Top-}k({\mathbf{P}}_i,k)\cup \text{Top-}p({\mathbf{P}}_i,p)$$ An online softmax Triton kernel executes sparse attention solely across ${\mathcal{S}}_i$.

Results

• Benchmark Superiority: On the 2K-Bench (300 native 2K test videos), Prism achieves $0.61$ Aesthetic Quality (AQ), $0.982$ Temporal Flickering (TF), $0.63$ DeSync, and $0.89$ MotionQ, substantially beating MOVA (Full Attn: $0.42$ AQ, $1.05$ DeSync, $0.53$ MotionQ) and LTX-2.3 with super-resolution ($0.48$ AQ, $0.95$ DeSync, $0.66$ MotionQ).
• Efficiency Gains: Training speed is $2.5\times$ faster than full attention at 2K ($10.6\text{ min}$ vs. $26.5\text{ min}$ per step). In 4-GPU inference, Prism runs in $331\text{ s}$ using $38.0\text{ GB}$ VRAM per GPU, compared to MOVA's $1091\text{ s}$ and $72.8\text{ GB}$.
• Resolution Scaling: Prism is the only tested framework showing monotonic improvements across all quality and cross-modal sync metrics when scaling from 480p to 2K, whereas full attention and fixed-block sparse baselines degrade at 2K due to background token saturation.
• User Study: In a 200-participant 2AFC perceptual study across 30 2K clips, Prism achieved $>87\%$ preference over all open-source baselines across text alignment, audio quality, visual detail, motion realism, and lip synchronization.

Limitations and open questions

• Fixed Temporal Macro-Zones: Macro-zones have a static temporal window of $Z_T=8$. Rapid shot cuts or zoom transitions within this window introduce abrupt feature disparities that a single block shape cannot represent, causing transient distortion artifacts.
• Resolution and Model Scaling: Prism has been evaluated up to 2K resolution on a 16B parameter backbone; scaling to native 4K (which quadruples token counts) and exploring larger 30B+ DiT architectures remain unexplored.
• Audio-Visual Dependency Order: Because self-attention precedes cross-attention in the DiT block, audio cross-attention guidance relies on norms cached from the previous layer, introducing a single-layer delay in guidance signal propagation.