SLA2: Sparse-Linear Attention with Learnable Routing and QAT

Resumo Técnico: SLA2

1. Problema e Motivação

O artigo aborda a necessidade de acelerar modelos de difusão (especialmente para geração de vídeo) sem comprometer a qualidade de geração. A atenção padrão (Full Attention) tem complexidade quadrática $O(N^2)$, o que é proibitivo para sequências longas.

O método anterior, Sparse-Linear Attention (SLA), tentou resolver isso combinando uma ramificação de atenção esparsa (para capturar detalhes importantes) e uma ramificação de atenção linear (para capturar dependências de longo alcance de baixo custo). No entanto, o SLA original apresentava duas limitações críticas identificadas pelos autores:

1. Roteamento Heurístico Subótimo: O SLA original dividia a computação entre as ramificações esparsa e linear baseada em uma heurística simples (selecionar os maiores pesos de atenção para a parte esparsa). Essa divisão não é matematicamente ótima e pode não maximizar a esparsidade ou a precisão.
2. Inconsistência na Decomposição: Houve uma análise de erro que revelou uma incompatibilidade entre a formulação do SLA e a motivação original de decompor a matriz de atenção $P$ em uma parte esparsa ($P_1$) e uma parte de baixo posto ($P_2$). O SLA original introduzia um erro de escala (fator $\alpha$) na ramificação esparsa que exigia uma projeção linear adicional para compensar, o que dificultava o aprendizado e não alinhava perfeitamente com a decomposição teórica.

2. Metodologia: SLA2

Os autores propõem o SLA2, uma reformulação da atenção esparsa-linear que corrige as limitações acima através de três componentes principais:

• A. Formulação de Mistura Mais Fiel (Decomposição Direta):
  Em vez de usar uma projeção linear para corrigir erros de escala, o SLA2 introduz um vetor de razão aprendível ($\alpha$) que combina diretamente as saídas das ramificações esparsa e linear.

  • A saída é calculada como: $O = \alpha \odot O_s + (1 - \alpha) \odot O_l$.
  • Isso alinha a saída exatamente com a decomposição teórica $P \approx \alpha P_s + (1-\alpha) P_l$, eliminando a necessidade de correções heurísticas e garantindo que a normalização de linha seja preservada.

• B. Roteador Aprendível (Learnable Router):
  Para resolver o problema do roteamento heurístico, o SLA2 introduz um módulo $R$ que decide dinamicamente quais posições da matriz de atenção devem ser computadas pela ramificação esparsa.

  • Entrada: O roteador utiliza apenas as matrizes de Query ($Q$) e Key ($K$), pois a formação dos pesos de atenção depende delas.
  • Eficiência: Para evitar o custo $O(N^2)$, o roteador aplica pooling (média) em blocos de tokens adjacentes e utiliza projeções lineares aprendíveis ($proj_q, proj_k$) antes de selecionar os top-$k$ pesos.
  • Treinamento: Durante o treinamento, o operador Top-k (não diferenciável) é substituído por um SoftTop-k diferenciável para permitir a retropropagação de gradientes.

• C. Atenção de Baixa Precisão com QAT (Quantization-Aware Training):
  Para obter aceleração adicional, a ramificação esparsa é quantizada para baixo (ex: INT8 ou FP8).

  • Diferente da quantização pós-treinamento (PTQ), o SLA2 utiliza QAT. Durante o treinamento, a passagem frontal (forward) usa a atenção quantizada, enquanto a passagem traseira (backward) permanece em FP16. Isso permite que o modelo se adapte aos erros de quantização, mantendo a precisão durante a inferência.

• D. Estratégia de Treinamento em Duas Etapas:

  1. Inicialização: Treinar o roteador $R$ e o vetor $\alpha$ para minimizar o erro de reconstrução em relação à atenção completa, garantindo uma inicialização estável.
  2. Ajuste Fino (Fine-tuning): Ajustar todo o modelo de difusão de ponta a ponta, otimizando a perda de difusão, permitindo que o modelo se adapte à alta esparsidade.

3. Contribuições Principais

1. Análise e Reformulação: Identificação e correção das limitações teóricas e práticas do SLA original, propondo uma formulação de atenção esparsa-linear que é matematicamente consistente com a decomposição desejada.
2. Roteamento Dinâmico e Aprendível: Substituição de regras fixas por um roteador treinável que otimiza a divisão de trabalho entre as ramificações esparsa e linear.
3. Integração de QAT: Demonstração de que a quantização de baixa precisão pode ser integrada ao SLA2 sem degradação significativa de qualidade, acelerando ainda mais a inferência.
4. Desempenho Superior: O método alcança níveis de esparsidade extremos mantendo (e às vezes superando) a qualidade de modelos com atenção completa.

4. Resultados Experimentais

Os experimentos foram conduzidos nos modelos de vídeo Wan2.1-1.3B e Wan2.1-14B.

• Eficiência:
  • O SLA2 alcançou 97% de esparsidade na atenção.
  • Isso resultou em um aceleração de 18,6x no tempo de execução da atenção em comparação com a atenção completa (Full Attention).
  • Em termos de latência de ponta a ponta, houve uma redução de 2,3x no modelo 1.3B e 4,35x no modelo 14B.
• Qualidade:
  • Mesmo com 97% de esparsidade, o SLA2 superou métodos de base (baselines) operando com apenas 90% de esparsidade em todas as métricas de qualidade de vídeo (IQ, OC, AQ, MS, SC, VR).
  • Curiosamente, em algumas métricas, o SLA2 ajustado superou o modelo com atenção completa (0% de esparsidade), atribuído à alta qualidade do conjunto de dados de ajuste fino.
• Ablação:
  • A remoção do QAT causou queda na qualidade, confirmando a eficácia do treinamento consciente da quantização.
  • O roteador aprendível superou significativamente o roteador Top-k estático do SLA original.

5. Significância

O SLA2 representa um avanço significativo na eficiência de modelos de difusão para vídeo. Ao resolver as inconsistências teóricas do SLA anterior e introduzir mecanismos de aprendizado para roteamento e quantização, o método permite a geração de vídeo de alta qualidade com custos computacionais drasticamente reduzidos. Isso torna viável a execução de modelos grandes de geração de vídeo em hardware mais acessível e acelera o processo de inferência, um gargalo crítico para aplicações em tempo real e produção de conteúdo.