Resumo Técnico: Atenção de Raio de Token para Geração de Vídeo Eficiente

1. Declaração do Problema

Os Video Diffusion Transformers (VDiTs) tornaram-se a arquitetura dominante para a geração de vídeo de alta fidelidade, baseando-se na atenção autoatencional 3D densa para capturar dependências espaciais e temporais complexas. No entanto, a atenção densa incorre em um custo computacional quadrático ($O(N^2)$) em relação ao número de tokens de vídeo $N$, criando um gargalo significativo de inferência para escalonamento e implantação.

Abordagens existentes para mitigar esse custo empregam mecanismos de atenção esparsa, categorizados principalmente em métodos de nível de cabeça (head-level) e de nível de bloco (block-level).

• Métodos de nível de cabeça atribuem padrões espaciais ou temporais estruturados a cabeças de atenção individuais.
• Métodos de nível de bloco selecionam pares de blocos de consulta-chave (query–key) importantes.

A Limitação Central: Ambos os paradigmas tipicamente compartilham um único orçamento de computação (por exemplo, uma densidade de retenção ou raio fixo) para todos os tokens de consulta dentro de uma cabeça ou bloco. Essa suposição de "alocação compartilhada" conflita com a natureza intrínseca da autoatenção, onde a normalização Softmax é realizada independentemente para cada consulta. Consequentemente, consultas com distribuições de atenção concentradas (baixa entropia) desperdiçam computação se forem forçadas a atender a muitas chaves, enquanto consultas com distribuições difusas (alta entropia) podem perder interações críticas se o orçamento compartilhado for muito pequeno. Esse descompasso degrada a qualidade da geração ou falha em atingir a eficiência ideal.

2. Metodologia: Atenção de Raio de Token (TRA)

Os autores propõem a Atenção de Raio de Token (Token Radius Attention - TRA), um framework livre de treinamento que realiza a esparsidade estruturada específica por token sem a necessidade de classificação (ranking) explícita de chaves por consulta. A TRA opera sobre dois insights empíricos derivados da análise dos padrões de atenção dos VDiTs.

Insight I: Esparsidade de Atenção Específica por Token

Os autores observam que a "densidade retida" (a fração de chaves de ranking superior necessária para preservar uma massa de atenção alvo) varia em ordens de magnitude entre as consultas dentro da mesma camada e cabeça. Crucialmente, eles encontem uma forte relação log-linear entre essa densidade retida e a entropia de atenção.

• Mecanismo: Alta entropia indica uma distribuição de atenção difusa que requer um orçamento de tokens maior, enquanto baixa entropia indica uma distribuição concentrada que requer menos tokens.
• Implementação: A TRA utiliza uma aproximação analítica para mapear a entropia da consulta diretamente para um orçamento específico por token ($\hat{B} \approx \tau \exp(H_i) / N$), eliminando a necessidade de ordenação ou classificação dispendiosa por consulta.

Insight II: Padrão de Atenção de Raio de Token

Os autores observam que as interações dominantes para uma consulta formam vizinhanças circulares centradas na consulta no domínio espacial. A probabilidade de atenção normalizada decai aproximadamente de forma exponencial com a distância espacial 2D.

• Mecanismo: Em vez de selecionar chaves top-$k$ arbitrárias, a TRA converte o orçamento de tokens escalar em um raio espacial.
• Decaimento Temporal: Para lidar com sequências de vídeo, o raio espacial é modulado por uma regra de decaimento de distância temporal ($\phi(\delta) = \exp(-\gamma \delta)$), criando uma estrutura de suporte espaço-temporal regular.

O Pipeline da TRA

1. Entropia-para-Orçamento: Durante uma fase inicial de "aquecimento" (warm-up) densa, a TRA computa a entropia de atenção para cada consulta. Esta entropia é mapeada para um orçamento de tokens específico.
2. Orçamento-para-Raio: O orçamento é convertido em um raio base específico por consulta. Este raio é então ajustado para cada frame com base na distância temporal para formar uma máscara de disco 2D.
3. Execução Eficiente:
   • Layout Tile-Major: Os tokens de vídeo são rearranjados em uma ordem tile-major para garantir que tokens espacialmente vizinhos (que caem dentro do raio) sejam contíguos na sequência.
   • Kernels CUDA Fundidos: Um kernel customizado funde o cálculo de distância, comparação de raio, poda de blocos e votação de tokens para converter as máscaras de raio irregulares de nível de token em máscaras de esparsidade de bloco regulares compatíveis com o FlashInfer.
   • Reuso entre Passos (Cross-Step Reuse): A entropia é computada apenas durante a fase de aquecimento e reutilizada para os passos esparsos subsequentes, aproveitando a estabilidade observada nos padrões de entropia através dos passos de denoising.

3. Principais Contribuições

1. Descoberta de Esparsidade Específica por Token: O artigo revela que a densidade retida varia dramaticamente entre as consultas, mas segue uma relação log-linear previsível com a entropia de atenção, e que as interações dominantes seguem um padrão de raio centrado na consulta.
2. Atenção de Raio de Token (TRA): Um framework livre de treinamento que transforma demandas de computação específicas por token em suportes espaço-temporais regulares através de um pipeline de entropia-para-orçamento-para-raio, evitando a classificação de chaves explícita.
3. Co-design de Sistema: Os autores realizam o co-design de um kernel de máscara de raio e um kernel de entropia fundido para minimizar o overhead, permitindo a execução eficiente em backends de esparsidade de bloco existentes.
4. Validação Abrangente: A TRA é validada em sete configurações de Wan2.1, Wan2.2 e HunyuanVideo (abrangendo de 1.3B a 14B parâmetros) tanto para tarefas de Texto-para-Vídeo (T2V) quanto de Imagem-para-Vídeo (I2V).

4. Resultados Experimentais

A TRA foi avaliada contra a atenção densa e três baselines de esparsidade livres de treinamento (SVG, SVG2 e Radial) no conjunto de benchmarks VBench.

• Eficiência: A TRA retém apenas 9%–19% das interações de atenção. Ela alcança um aumento de velocidade de 1.56× a 2.05× sobre os modelos densos em todas as configurações testadas. Por exemplo, no HunyuanVideo-13B, ela alcança um aumento de 2.05× para T2V e 2.02× para I2V.
• Qualidade: A TRA mantém uma qualidade de geração competitiva, frequentemente alcançando os melhores scores de "Overall" do VBench entre os métodos esparsos. No Wan2.1-14B T2V, ela quase iguala o score da atenção densa enquanto fornece um aumento de velocidade de 1.75×.
• Fidelidade: Resultados qualitativos mostram que a TRA preserva melhor a identidade do sujeito, a integridade estrutural e a consistência temporal em comparação com outros métodos esparsos, com menos artefatos como cintilação (flickering) ou distorção.
• Estudos de Ablação: Remover o mapeamento de orçamento guiado por entropia ou a máscara de raio (substituindo-a por distância 1D) degrada significativamente o desempenho, confirmando a necessidade tanto da adaptatividade específica por token quanto do suporte espacial estruturado.

5. Significância e Alegações

O artigo afirma que a Atenção de Raio de Token aborda uma ineficiência fundamental nos paradigmas atuais de atenção esparsa: o descompasso entre os orçamentos de alocação compartilhados e as demandas de atenção específicas por token. Ao aproveitar a relação intrínseca entre entropia e esparsidade de atenção, a TRA alcança um equilíbrio favorável entre qualidade e eficiência sem exigir treinamento adicional ou mecanismos complexos de classificação online.

Os autores posicionam a TRA como uma técnica complementar a outros métodos de aceleração (como caching ou quantização) que foca especificamente na redução do custo quadrático de interação consulta–chave dentro das camadas de atenção executadas. O trabalho demonstra que a adaptatividade ao nível de token pode ser realizada eficientemente através de máscaras estruturadas baseadas em raio, ofereando um caminho escalável para implantar modelos de geração de vídeo de alta fidelidade.