SALAD: Achieve High-Sparsity Attention via Efficient Linear Attention Tuning for Video Diffusion Transformer

Título: SALAD: Attenção de Alta Esparsidade via Ajuste Eficiente de Atenção Linear para Transformers de Difusão de Vídeo

1. O Problema

Os Transformers de Difusão (DiT) têm demonstrado desempenho excepcional na geração de vídeo. No entanto, eles enfrentam um gargalo crítico de eficiência: a complexidade quadrática $O(N^2)$ do mecanismo de atenção padrão (Full Attention) em relação ao comprimento da sequência ($N$).

• Desafio: Gerar vídeos de alta resolução e longa duração (com dezenas de milhares de tokens) torna-se proibitivamente caro em termos de memória e latência.
• Soluções Existentes e Limitações:
  • Abordagens sem treinamento (Training-free): Métodos de atenção esparsa (como janelas deslizantes) são limitados a esparsidades moderadas (40-60%) para manter a qualidade, resultando em acelerações modestas.
  • Abordagens baseadas em treinamento: Podem atingir esparsidades muito altas (80-95%), mas exigem grandes volumes de dados e custo computacional massivo (ex: VMoBA requer 182 horas de GPU e 36M de amostras).
  • Ajuste com LoRA: A adaptação de baixo rank (LoRA) para recuperar a qualidade após a compressão falha em cenários de ultra-esparsidade. O LoRA sozinho não consegue restaurar adequadamente as interações de longo alcance, levando a artefatos visuais (como duplicação de objetos ou inconsistências temporais) e degradação da qualidade em comparação com a atenção densa.

2. Metodologia: SALAD

O authors propõem o SALAD (High-Sparsity Attention paralleling with Linear Attention for Diffusion Transformer), uma arquitetura que combina atenção esparsa com um ramo de atenção linear leve em paralelo.

Arquitetura Principal:

• Ramos Paralelos: O modelo mantém o ramo de Atenção Esparsa (que domina o processamento para garantir eficiência) e introduz um ramo de Atenção Linear (baseado em funções de kernel, como ReLU, com complexidade $O(N)$) em paralelo.
• Compartilhamento de Parâmetros: Ambos os ramos compartilham as matrizes de projeção ($W_Q, W_K, W_V$) inicializadas a partir do modelo pré-treinado e ajustadas via LoRA. O ramo linear adiciona apenas cerca de 4,99% de parâmetros extras.
• Integração Residual: A saída do ramo linear é somada à saída da atenção esparsa, mas com um controle rigoroso de escala.

Estratégia Chave: Escalonamento Estático-Dinâmico Multi-nível (Multi-level Static-Dynamic Scaling)
Os autores observaram que os ramos de atenção esparsa e linear possuem um desequilíbrio de rank (a atenção linear tende a ter um rank efetivo muito menor). Somá-los diretamente pode suprimir estruturas informativas de alto rank da atenção esparsa. Para resolver isso, o SALAD utiliza:

1. Escalonamento Estático por Camada (LWSS): Uma camada de projeção inicializada com zeros (Zero-initialized) após o ramo linear. Isso permite que o modelo comece o ajuste a partir da atenção esparsa pura e aprenda gradualmente a contribuição do ramo linear.
2. Escalonamento Dinâmico Consciente de Entrada e Passo de Tempo (ITADS): Um módulo leve que gera um fator de escala global ($s \in (0, 1)$) baseado no estado oculto atual e no passo de tempo de denoising ($t$). Isso permite que a contribuição do ramo linear seja adaptativa: maior quando necessário para recuperar informações globais e menor quando a atenção esparsa é suficiente.

Pós-Ajuste (Branch Dropping):
Com base nos fatores de escala dinâmicos, o método permite "descartar" (dropping) ramos de atenção linear que recebem fatores de escala muito baixos durante a inferência, ganhando eficiência adicional sem perda de qualidade.

3. Contribuições Principais

1. Arquitetura SALAD: Uma nova estrutura de atenção que permite modelagem de ultra-esparsidade (até 90%) mantendo a qualidade de geração, através de um ramo linear auxiliar que compensa a perda de interações de longo alcance.
2. Estratégia de Escalonamento Híbrido: A descoberta de que o controle preciso da magnitude do ramo linear (via escalonamento estático e dinâmico) é crucial para evitar a degradação de rank e garantir que o ramo linear atue como um complemento, não como um competidor da atenção esparsa.
3. Eficiência de Treinamento: O método é altamente eficiente, exigindo apenas 2.000 vídeos de amostra, menos de 1.600 passos de treinamento e no máximo 30 horas de GPU (com batch size 8), uma fração mínima do custo de métodos concorrentes.

4. Resultados Experimentais

Os experimentos foram realizados no modelo Wan2.1 (versões 1.3B e 14B) com resoluções de 480p e 720p.

• Desempenho vs. Velocidade:
  • O SALAD atinge 90% de esparsidade (redução de computação) mantendo qualidade comparável ou superior à atenção completa (Full Attention).
  • Aceleração: Ganhos de velocidade de inferência de 1.52x a 2.03x dependendo do modelo e comprimento da sequência.
  • Em comparação com o LoRA puro em atenção esparsa, o SALAD supera significativamente as métricas de consistência de sujeito, fundo e qualidade de imagem.
• Qualidade Visual:
  • Elimina artefatos comuns em modelos esparsos, como duplicação de objetos (ex: dois cães em um quadro) e inconsistências textuais/temporais.
  • Mantém detalhes finos e coerência temporal superior.
• Eficiência de Parâmetros: O SALAD supera variantes LoRA com ranks maiores (até 512) usando menos parâmetros treináveis (165M vs 377M para LoRA r=512), provando que o ganho vem do design da arquitetura e não apenas do aumento de parâmetros.

5. Significância e Impacto

O trabalho SALAD representa um avanço significativo na viabilidade de geração de vídeo de alta resolução e longa duração em hardware limitado.

• Viabilidade Prática: Ao reduzir drasticamente o custo de treinamento (de centenas de horas e milhões de dados para dezenas de horas e milhares de dados), torna a adaptação de modelos de difusão esparsos acessível para mais pesquisadores e empresas.
• Equilíbrio Qualidade-Eficiência: Resolve o dilema tradicional onde alta esparsidade levava a baixa qualidade. O SALAD demonstra que é possível atingir 90% de redução computacional sem sacrificar a fidelidade do vídeo.
• Generalização: A abordagem é aplicável a diferentes famílias de modelos e estratégias de esparsidade (estáticas e dinâmicas), sugerindo um novo padrão para a otimização de Transformers de Difusão.

Em resumo, o SALAD oferece uma solução elegante e eficiente para o problema de escalabilidade em geração de vídeo, combinando a velocidade da atenção esparsa com a capacidade de modelagem global da atenção linear, tudo isso com um custo de ajuste extremamente baixo.