Resumo Técnico: PISA - Atenção Esparsa por Partes para Transformers de Difusão Eficientes

1. Declaração do Problema

Transformers de Difusão (DiTs) tornaram-se o padrão para geração de imagem e vídeo de alta fidelidade. No entanto, sua escalabilidade é limitada pela complexidade quadrática ($O(L^2)$) do mecanismo de autoatenção à medida que os comprimentos de sequência de entrada ($L$) aumentam para resoluções mais altas ou vídeos mais longos.

As soluções existentes enfrentam limitações significativas:

• Atenção Esparsa ("Manter-ou-Descartar"): Métodos que descartam blocos de chave-valor (KV) não críticos aceleram a computação, mas sofrem degradação de desempenho em altas taxas de esparsidade devido à perda de contexto. Eles também apresentam dificuldades com sequências curtas, onde o overhead de seleção supera os benefícios.
• Atenção Linear/Híbrida: Abordagens que alteram a distribuição da atenção (por exemplo, via mapeamentos de características ou combinações aditivas de ramos esparsos e lineares) frequentemente perturbam a normalização intrínseca de modelos pré-treinados. Isso exige retreinamento ou ajuste fino (fine-tuning) dispendiosos, impedindo o reuso direto de pesos pré-treinados.
• Truncamento Rígido (Hard Truncation): Descartar blocos diretamente leva a desvios acentuados da distribuição de atenção original, causando perda de qualidade no conteúdo gerado.

O desafio central é projetar um mecanismo que acelere a inferência com complexidade subquadrática enquanto mantém a precisão da atenção total e permaneça compatível com pesos pré-treinados sem necessidade de retreinamento.

2. Metodologia: Atenção Esparsa por Partes (PISA)

O PISA introduz um novo paradigma "exato-ou-aproximado", indo além da tradicional estratégia de "manter-ou-descartar". Ele trata a atenção como um processo por partes que cobre todo o espectro da atenção com complexidade subquadrática.

Insights Principais

Os autores analisaram DiTs pré-treinados e descobriram duas propriedades críticas:

1. Estabilidade Distribuicional: Os escores de atenção pré-softmax de blocos não críticos exibem uma distribuição simétrica em forma de sino centrada em zero ou valores negativos. Isso os torna altamente suscetíveis a uma aproximação precisa via expansão de Taylor, em vez de exigir computação exata ou descarte.
2. Consistência de Normalização: Para manter a compatibilidade sem treinamento, as aproximações devem ser integradas internamente no numerador e denominador do softmax, em vez de simplesmente adicionar saídas de diferentes ramos (o que viola as regras de soma ponderada).

Algoritmo Central

O PISA particiona os blocos KV em dois conjuntos:

1. Blocos Críticos ($S_i$): Computados de forma exata para preservar informações de alta saliência.
2. Blocos Não Críticos ($U_i$): Aproximados eficientemente usando expansão de Taylor por bloco.

A saída $o_t$ é derivada pela normalização da agregação ponderada dos valores:
$$o_t = \frac{N_t}{D_t}$$

Onde o denominador ($D_t$) e o numerador ($N_t$) são compostos por:

• Termos Esparsos Exatos: Computados para os blocos selecionados $j \in S_i$.
• Aproximação de Ordem Zero por Bloco: Para blocos não selecionados, a contribuição é aproximada usando o centroide do bloco $\bar{k}_j$ e a soma dos valores.
• Correção Global de Primeira Ordem: Em vez de computar um termo de primeira ordem distinto para cada bloco não selecionado (o que seria limitado pela memória), o PISA emprega uma aproximação de ordem híbrida. Ele utiliza uma estatística global $\bar{H}$ (a média das matrizes de primeira ordem por bloco entre todos os blocos) para aproximar o gradiente de primeira ordem para todo o conjunto de cauda.

A formulação para o numerador $N_t$ torna-se:
$$N_t = \sum_{j \in S_i} \sum_{n} e^{q_t k_{j,n}^\top} v_{j,n} + \sum_{j \in U_i} \alpha_{t,j} \left(\sum_{n} v_{j,n}\right) + q_t \bar{H} \sum_{j \in U_i} \alpha_{t,j}$$
Onde $\alpha_{t,j} = \exp(q_t \bar{k}_j^\top)$.

Roteamento Consciente de Covariância

Para minimizar o erro de aproximação, o PISA utiliza uma Estratégia Top-k Consciente de Covariância. O escore de seleção de um bloco incorpora não apenas o escore de atenção, mas também a norma da matriz de covariância do bloco ($M_j = \|H_j - \bar{H}\|_2$). Blocos com alta relevância semântica ou alta heterogeneidade são selecionados para computação exata, enquanto blocos mais suaves são delegados à expansão de Taylor.

Implementação de Hardware

O PISA é implementado via um kernel fundido (fused kernel) que intercala a computação exata e a varredura aproximada dentro do processo de softmax online (estilo FlashAttention). O pipeline envolve:

1. Pré-computação: Cálculo dos centroides dos blocos, somas de valores e da estatística global $\bar{H}$.
2. Fase 1 (Exata): Carregamento dos blocos selecionados na SRAM para atenção exata.
3. Fase 2 (Aproximação): Varredura de vetores agregados de blocos para blocos não selecionados usando expansão de ordem zero.
4. Fase 3 (Correção): Injeção do termo de correção de primeira ordem global.

Este design garante que o método seja eficiente em termos de memória e evite os gargalos associados ao carregamento de matrizes distintas para cada bloco.

3. Principais Contribuições

1. Atenção Esparsa por Partes Inovadora: Um mecanismo livre de treinamento que permite a cobertura de todo o espectro de atenção com complexidade subquadrática, unificando a computação esparsa exata e a aproximação por bloco.
2. Esquema de Aproximação de Ordem Híbrida: Uma estratégia que combina a expansão de ordem zero por bloco com uma correção de primeira ordem global. Isso aumenta significativamente a precisão com um custo computacional negligenciável em comparação com a atenção esparsa padrão.
3. Roteamento Consciente de Covariância: Uma derivação de análise de erro que minimiza a divergência de aproximação ao selecionar blocos baseando-se tanto nos escores de atenção quanto nas normas de covariância do bloco.
4. Desempenho de Estado da Arte: Demonstrou eficiência e qualidade superiores em diversas tarefas (geração de vídeo e imagem) sem a necessidade de retreinamento.

4. Resultados Experimentais

Os autores avaliaram o PISA em modelos de geração de vídeo (Wan2.1-1.3B/14B, Hunyuan-Video) e modelos de geração de imagem (SD 3.5, FLUX.1).

• Geração de Vídeo:

  • Wan2.1-14B: Alcançou um aceleração de 1.91× sobre a atenção densa (1468s $\to$ 818s) com 87,5% de esparsidade, mantendo ou até melhorando levemente as métricas de qualidade (VBench, PSNR, SSIM) em comparação com a atenção densa.
  • Hunyuan-Video: Alcançou uma aceleração de 2.57× (1651s $\to$ 641s) com preservação consistente de qualidade.
  • Comparação: Superou métodos esparsos existentes como SpargeAttn e SVG2 tanto em velocidade quanto em qualidade, particularmente na preservação de detalhes finos e consistência temporal.

• Geração de Imagem:

  • FLUX.1-dev: Alcançou uma aceleração de 1.2× (8.32s $\to$ 6.87s) com 85% de esparsidade. Notavelmente, o PISA obteve pontuações mais altas em métricas de preferência humana (ImageReward, Hpsv2) e métricas de similaridade (PSNR, SSIM) em comparação tanto com a atenção densa quanto com o SpargeAttn.
  • SD 3.5: Demonstrou acelerações significativas (ex: 1.4× no SD 3.5-M) mantendo pontuações FID comparáveis ou superiores à atenção densa.

• Eficiência de Kernel:

  • O PISA superou consistentemente o FlashAttention-2/3 e o SpargeAttn em latência através de vários comprimentos de sequência (4K a 32K) e densidades.
  • Ao contrário do SpargeAttn, que degrada em sequências curtas, o PISA manteve vantagens de velocidade mesmo em 4K tokens.

5. Significância e Alegações

O artigo afirma que o PISA representa uma mudança de paradigma do "manter-ou-descartar" para a atenção esparsa "exato-ou-aproximado". Sua principal significância reside em:

• Ponte para o Gap Velocidade-Qualidade: Ele efetivamente une o gap entre velocidade de inferência e qualidade de geração, alcançando uma aproximação quase sem perdas com overhead negligenciável.
• Compatibilidade Sem Treinamento: Ao integrar as aproximações diretamente na estrutura de normalização do softmax, o PISA permite o reuso direto de pesos pré-treinados, eliminando a necessidade de retreinamento ou ajuste fino dispendioso exigido por métodos de atenção linear ou híbrida.
• Escalabilidade: Fornece uma solução unificada tanto para geração de imagens de alta resolução quanto para geração de vídeos de longa duração, abordando o gargalo da complexidade quadrática sem sacrificar a fidelidade do conteúdo gerado.

Os autores enfatizam que seu método é o primeiro a integrar nativamente a expansão de Taylor por bloco e a atenção esparsa em um framework unificado que preserva a distribuição intrínseca dos modelos pré-treinados.