AB-Attention acelera a geração de vídeo em até 2,23× e pula 92% da atenção

A equipe da InfRec, em parceria com pesquisadores da CUHK Shenzhen e das universidades de Shandong, Tóquio e Fudan, além de Berkeley, HKUST e NUS, lançou o AB-Attention. Trata-se de um modo de acelerar a geração de vídeo em transformers de difusão (Wan2.1 e HunyuanVideo) sem retreinar os pesos.

A operação mais cara nesses modelos é a atenção: cada fragmento do vídeo é comparado com todos os outros, no espaço e no tempo, então o custo cresce de forma quadrática com o tamanho da sequência. O AB-Attention pula 89–92% desses cálculos. Mesmo assim, o vídeo gerado fica mais próximo do resultado da atenção completa “de verdade” do que o dos concorrentes.

O PSNR (métrica de similaridade com o vídeo de referência, em decibéis: quanto maior, mais próximo) é melhor em 1,6–4,4 dB, e a geração completa em uma única placa NVIDIA H20 é 1,84–2,23 vez mais rápida.

Onde os esquemas esparsos anteriores falham

A atenção esparsa calcula apenas parte dos pares “query — chave”, e não todos. Os primeiros métodos definiam um padrão de antemão, como janelas 3D locais. Eles não sabem o que é importante para um prompt, uma cabeça de atenção ou uma etapa de remoção de ruído específicos, e por isso acabam mantendo muita coisa desnecessária. Os métodos dinâmicos escolhem blocos de chaves e valores com base nas ativações atuais, mas, segundo os autores, têm três pontos fracos:

• As estimativas de importância são grosseiras ou caras. Tirar a média do bloco em um único vetor mistura tokens de comportamentos diferentes, e a clusterização global (SVG2) exige k-means sobre toda a sequência e permutação de tokens. • A seleção top-k corta a cauda. O bloco descartado some tanto do numerador da atenção quanto da normalização do softmax. Estimativas melhores não resolvem isso: mesmo um top-k ideal, com 10,8% das chaves, preserva só cerca de 69% da massa de atenção. • A unidade de seleção é o tile de computação. O bloco inteiro é mantido ou descartado, mesmo quando só algumas de suas chaves são necessárias.

Os três recursos do AB-Attention

O primeiro recurso, o IBCS (Intra-Block Clustering Selector), estima a importância dos blocos de chaves. As queries são divididas em blocos de 128 tokens. Dentro de cada bloco, elas são agrupadas em 8 clusters, cada um com um centroide e um peso igual à fração de queries no cluster.

O bloco de chaves recebe como nota a atenção ponderada dos centroides às suas chaves médias. Os clusters servem apenas para a estimativa; o cálculo em si segue a ordem original dos tokens. Assim, um bloco observado por apenas um cluster não se perde, como acontece quando se tira a média do bloco inteiro.

A diferença de cobertura da seleção entre a média simples e o top-k ideal cai pela metade.

O segundo recurso, o SIP (Sampling with Inclusion-Probability Weighting), troca a seleção rígida por amostragem. Um bloco entra na amostra com probabilidade proporcional à sua nota. Os blocos mais importantes são sempre escolhidos, e os demais têm chance diferente de zero. Cada bloco selecionado é multiplicado pelo inverso da probabilidade de inclusão: um bloco com 25% de probabilidade conta como quatro. É o estimador clássico de Horvitz-Thompson, que fornece estimativas não enviesadas tanto do numerador quanto da normalização do softmax. O erro de atenção fica 41% menor que o do top-k ideal com o mesmo número de blocos. Não é preciso uma passada extra para correção: o cálculo continua sendo a atenção softmax comum, com deslocamento de logits.

O terceiro recurso é a seleção fina em hardware de granularidade grossa. Os blocos de chaves e valores têm apenas 4 tokens, mas o kernel reúne os blocos escolhidos, por índices, em tiles completos dos tensor cores. Os tokens não são permutados, e nenhuma cópia desnecessária é criada. Reduzir o bloco de 128 para 4 tokens eleva a fração da massa de atenção retida pela seleção ideal de 59% para 69%, sem aumentar o número de tiles com orçamento fixo. O plano de amostragem é recalculado a cada quatro etapas de remoção de ruído, enquanto a atenção em si é calculada em todas.

Como foi a avaliação e quais foram os resultados

Foram comparados três modelos em 720p: duas versões do Wan2.1 (81 quadros) e o HunyuanVideo (129 quadros), em 600 prompts do benchmark Penguin Video. Como referência, usou-se o vídeo gerado com atenção completa, com o mesmo prompt e a mesma seed de geração. Além do PSNR, foram calculados SSIM e LPIPS, além de cinco avaliações do próprio vídeo: estética, consistência do fundo, qualidade da imagem, fluidez do movimento e constância do objeto. Os concorrentes são RadialAttention (máscara estática com decaimento de energia), SVG2 e DFSAttn, todos com o código oficial e o backend FlashInfer com FlashAttention-3. O primeiro bloco do transformer e as primeiras 10 das 50 etapas rodam com atenção completa em todos os métodos.

• Esparsidade: 89,16% no Wan2.1 e 92,35% no HunyuanVideo. • Os concorrentes, nesse cenário, calculam de 1,8 a 3,6 vezes mais atenção. • O AB-Attention é o melhor em PSNR, SSIM e LPIPS nos três modelos. • Aceleração de ponta a ponta: 1,84–2,23× em uma única H20, em BF16.

O que isso muda na prática

O método não exige retreino, então pode ser aplicado a pesos já prontos. Para as equipes que rodam Wan2.1 ou HunyuanVideo, isso significa quase metade do tempo por vídeo, com qualidade próxima à da referência. O ganho vai crescer em vídeos longos e de alta resolução, em que a atenção pesa mais na latência. Há também limitações: as medições foram feitas na H20 e em três modelos, e o kernel e o plano de amostragem foram projetados para os tensor cores da NVIDIA. O código está aberto no GitHub, no repositório InfRec-Cardinal/AB-Attention, então é possível conferir os resultados no seu próprio hardware.