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

Guia de leitura

Problema e motivação

Treinar nativamente modelos de difusão conjuntos de vídeo e áudio em resolução 2K possibilita a síntese de detalhes espaciais ricos, dinâmicas de movimento nítidas e som sincronizado. No entanto, após a compressão por VAE, um clipe de 2K gera de centenas de milhares a mais de um milhão de tokens. A atenção completa padrão sofre com dois problemas fatais:

1. Complexidade quadrática: A pegada computacional e o consumo de memória tornam o treinamento proibitivamente caro.
2. Diluição do peso de atenção (attention weight dilution): Volumes massivos de tokens repetitivos de fundo dominam a sequência de tokens, diluindo o orçamento de atenção alocado para conteúdos informativos e localizados (como rostos em movimento, mãos e objetos emissores de som) e desestabilizando priors multimodais pré-treinados.

Abordagens anteriores de sparse attention aplicam aproximações training-free, focam apenas em vídeos silenciosos em resoluções mais baixas (480p/720p) ou impõem formatos de bloco rígidos e uniformes. Como a sincronização audiovisual é estritamente concentrada em regiões compactas e emissoras de som, blocos de formato fixo misturam tokens heterogêneos, corrompem representações de blocos calculadas por mean-pooling e falham em se adaptar aos padrões de interação intermodal.

Ideias centrais

• Macro-Zone Partitioning (Particionamento em Macrozonas): O volume espaço-temporal de tokens é particionado em macrozonas de tamanho $Z_T\times Z_H\times Z_W=8\times 8\times 8$ para permitir configurações de blocos alinhadas ao Triton.
• Dual-Guidance Variation Indicator (Indicador de Variação por Dupla Orientação): Cada zona mede dinamicamente a variação direcional de conteúdo combinando a variância dos canais de vídeo com as normas da cross-attention de áudio para vídeo.
• Content-Adaptive 3D Block Shapes (Formatos de Blocos 3D Adaptáveis ao Conteúdo): As zonas escolhem a partir de um conjunto candidato de 16 formatos com volumes totais $B\in \{64,128,256\}$, particionando finamente ao longo dos eixos de rápida mudança e forte acoplamento de áudio.
• Hybrid Top-$k$ / Top-$p$ Selection (Seleção Híbrida Top-$k$ / Top-$p$): A atenção esparsa em nível de bloco usa Top-$k$ para assegurar um contexto de referência garantido, expandido adaptativamente via Top-$p$ para queries com distribuições de atenção dispersas.

Método

O Prism opera em um Diffusion Transformer (DiT) de ramo duplo inicializado a partir do MOVA-720p, substituindo a autoatenção completa de vídeo por dynamic sparse attention, mantendo a cross-attention e a autoatenção de áudio inalteradas.

1. Formulação do Indicador de Variação

Para cada macrozona $m$, a razão de variação direcional de vídeo ao longo do eixo $d\in \{T,H,W\}$ é derivada dos recursos de valor médios ${\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 das normas da cross-attention de áudio para vídeo $a_{t,h,w}^m=\Vert {\hat{\mathbf{h}}}_{t,h,w}^m{\Vert }_2$ armazenadas em cache da camada anterior, a força de acoplamento de áudio $\bar{a}(m)\in [0,1]$ e a variância direcional normalizada ${\hat{v}}_{a,d}(m)$ são computadas. O indicador de variação composto é: $$g_d(m)=r_{v,d}(m)+\bar{a}(m)\cdot {\hat{v}}_{a,d}(m)+ϵ$$

2. Volume do Bloco e Quantização do Formato

O volume total do bloco $B(m)\in \{64,128,256\}$ é atribuído avaliando a densidade de informação espaço-temporal normalizada $\rho (m)$ em relação aos limiares ${\tau }_{128}=0.5$ e ${\tau }_{256}=0.25$. Minimizar a perda de informação intrabloco ${\sum }_d{g}_d(b_d^2-1)$ sujeita a ${\prod }_d{b}_d=B(m)$ resulta no comprimento de aresta contínuo ótimo $b_d^{\ast }\propto g_d^{-1/2}$. No espaço ${\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)$$ O formato do bloco $s^{\ast }=(b_T,b_H,b_W)$ é escolhido a partir do conjunto candidato ${\mathcal{C}}_{B(m)}$ minimizando a distância euclidiana em relação a ${\ell }_d^{\ast }$.

3. Atenção Híbrida Top-$k$ / Top-$p$

Os blocos de chaves (key blocks) passam por pooling pela média em representantes ${\bar{\mathbf{K}}}_j$, sendo pontuados em relação aos blocos de consulta (query blocks) ${\mathbf{Q}}_i$ para produzir probabilidades de bloco ${\mathbf{P}}_i$. O conjunto atendido ${\mathcal{S}}_i$ combina limiares de probabilidade fixos e cumulativos: $${\mathcal{S}}_i=\text{Top-}k({\mathbf{P}}_i,k)\cup \text{Top-}p({\mathbf{P}}_i,p)$$ Um kernel Triton de softmax online executa a sparse attention exclusivamente sobre ${\mathcal{S}}_i$.

Resultados

• Superioridade em Benchmarks: No 2K-Bench (300 vídeos de teste nativos em 2K), o Prism alcança $0.61$ em Aesthetic Quality (AQ), $0.982$ em Temporal Flickering (TF), $0.63$ em DeSync e $0.89$ em MotionQ, superando substancialmente o MOVA (Full Attn: $0.42$ AQ, $1.05$ DeSync, $0.53$ MotionQ) e o LTX-2.3 com super-resolução ($0.48$ AQ, $0.95$ DeSync, $0.66$ MotionQ).
• Ganhos de Eficiência: A velocidade de treinamento é $2.5\times$ mais rápida do que a full attention em 2K ($10.6\text{ min}$ vs. $26.5\text{ min}$ por passo). Na inferência com 4 GPUs, o Prism executa em $331\text{ s}$ usando $38.0\text{ GB}$ de VRAM por GPU, em comparação com os $1091\text{ s}$ e $72.8\text{ GB}$ do MOVA.
• Escalabilidade com a Resolução: O Prism é o único framework avaliado que demonstra melhorias monotônicas em todas as métricas de qualidade e sincronização intermodal ao escalar de 480p para 2K, enquanto a full attention e os baselines esparsos de blocos fixos degradam em 2K devido à saturação de tokens de fundo.
• Estudo com Usuários: Em um estudo perceptual 2AFC com 200 participantes em 30 clipes 2K, o Prism obteve $>87\%$ de preferência em relação a todos os baselines de código aberto nos critérios de alinhamento de texto, qualidade de áudio, detalhes visuais, realismo de movimento e sincronia labial.

Limitações e questões em aberto

• Macrozonas Temporais Fixas: As macrozonas têm uma janela temporal estática de $Z_T=8$. Cortes rápidos de cena ou transições de zoom dentro dessa janela introduzem disparidades abruptas de atributos que um único formato de bloco não consegue representar, gerando artefatos transitórios de distorção.
• Escalonamento de Resolução e Modelo: O Prism foi avaliado até a resolução 2K em um backbone de 16B de parâmetros; a extensão para 4K nativo (que quadruplica a contagem de tokens) e a investigação de arquiteturas DiT maiores, com mais de 30B, permanecem inexploradas.
• Ordem de Dependência Audiovisual: Como a autoatenção precede a cross-attention no bloco DiT, a orientação pela cross-attention de áudio depende de normas armazenadas em cache da camada anterior, introduzindo um atraso de uma camada na propagação do sinal de orientação.