Token Sparse Attention: Efficient Long-Context Inference with Interleaved Token Selection

1. Declaração do Problema

Os Modelos de Linguagem de Grande Escala (LLMs) enfrentam um gargalo fundamental na inferência de contexto longo devido à complexidade quadrática ($O(L^2)$) do mecanismo de atenção durante a etapa de pré-preenchimento (prefill). Embora otimizações de hardware como o FlashAttention reduzam a sobrecarga de E/S de memória, elas não alteram a complexidade algorítmica.

Os métodos de aceleração existentes sofrem com duas limitações principais:

• Esparsidade Estruturada: Métodos como Minference ou FlexPrefill aplicam esparsidade em nível de bloco. Isso frequentemente retém tokens irrelevantes se estiverem agrupados com tokens importantes no mesmo bloco, levando a uma esparsidade subótima.
• Evicção de Tokens: Métodos como FastKV ou GemFilter descartam permanentemente tokens em camadas iniciais. Isso assume que a importância do token é estática, ignorando a natureza dinâmica da relevância do token entre diferentes camadas e cabeças de atenção. Uma vez evictos, um token não pode ser reconsiderado, levando à degradação da precisão.

2. Metodologia: Token Sparse Attention

Os autores propõem o Token Sparse Attention, um mecanismo de esparsificação leve, dinâmico e reversível em nível de token. Diferentemente de métodos baseados em evicção, ele não remove tokens permanentemente, mas comprime o espaço de computação e depois restaura a sequência.

Arquitetura Central

O método opera em duas etapas por cabeça de atenção:

1. Compressão (Seleção QKV):
   • Para cada cabeça de atenção $h$, um subconjunto de índices de token $S_h$ é selecionado com base em pontuações de importância.
   • Os tensores de Consulta ($Q$), Chave ($K$) e Valor ($V$) são reunidos para formar tensores comprimidos $\hat{Q}, \hat{K}, \hat{V}$ com um comprimento de sequência reduzido $L' \ll L$.
   • Característica Chave: A seleção é realizada independentemente por cabeça, permitindo que diferentes cabeças foquem em subconjuntos diferentes de tokens, abordando a heterogeneidade entre cabeças.
   • Os tensores comprimidos permanecem densos e contíguos, permitindo o uso de kernels otimizados como o FlashAttention.
2. Descompressão (Restauração de Saída):
   • A saída de atenção $\hat{O}$ (forma $L' \times d$) é dispersa de volta para um tensor inicializado com zeros do comprimento de sequência original $L \times d$ usando os índices específicos da cabeça.
   • Característica Chave: As posições não selecionadas são zeradas, mas a conexão residual da camada anterior preserva a informação dos tokens pulados. Isso permite que o modelo "revisite" e reavalie a importância do token em camadas subsequentes.

Cobertura Dinâmica de Tokens

Para determinar quais tokens manter e quantos, o artigo introduz a Cobertura Dinâmica de Tokens:

• Pontuação: Um mapa de atenção proxy leve é computado usando consultas recentes contra todas as chaves. As pontuações de importância dos tokens são derivadas somando os pesos de atenção verticalmente.
• Alocação de Orçamento: Em vez de manter os top-$k$ tokens, o método ordena os tokens por importância crescente e poda a "cauda" (tokens menos importantes) até que a massa cumulativa dos tokens podados exceda um limiar de cobertura $\tau$. Isso trata a poda como regularização estrutural contra ruído de atenção.
• Seleção por Cabeça: Após determinar o orçamento global, cada cabeça seleciona independentemente seus top-$k$ tokens com base em suas pontuações específicas.

Seleção de Camada Esparsa

Os autores observam que aplicar esparsidade a todas as camadas degrada o desempenho. Eles introduzem o Desvio de Representação Inter-Camada ($R_\ell$) para medir a estabilidade das representações de tokens entre camadas.

• Camadas com baixo desvio (representações estáveis) são selecionadas para esparsificação.
• Camadas com alto desvio (onde as representações mudam significativamente) são mantidas densas para preservar o comportamento do modelo.
• Essa seleção é uma etapa de pré-processamento única por modelo.

3. Contribuições Principais

1. Esparsificação Reversível em Nível de Token: Um design inovador de "Comprimir e Descomprimir" que reduz o custo computacional sem evictar tokens permanentemente, permitindo a reavaliação dinâmica da importância do token entre camadas.
2. Heterogeneidade por Cabeça: O método suporta seleção independente de tokens para cada cabeça de atenção, capturando o fato de que diferentes cabeças priorizam diferentes características semânticas.
3. Compatibilidade Transparente: A abordagem é totalmente compatível com kernels densos existentes (FlashAttention) e esparsos (Minference, FlexPrefill), atuando como um acelerador complementar em vez de um substituto.
4. Orçamentação Dinâmica: Uma abordagem orientada por dados para determinar níveis de esparsidade com base nas distribuições de importância dos tokens e na estabilidade por camada, em vez de proporções fixas.

4. Resultados Experimentais

Os experimentos foram conduzidos em LLaMA-3.1-8B-Instruct e Mistral-Nemo-12B-Instruct usando benchmarks como RULER, InfiniteBench e Needle-in-a-Haystack.

• Compromisso Precisão-Latência:
  • O Token Sparse Attention melhora consistentemente o compromisso velocidade-precisão.
  • Em contexto de 128K, alcança até $\times 3,23$ de aceleração de atenção comparado ao FlashAttention com menos de 1% de degradação de precisão.
  • Quando combinado com FlexPrefill, melhora a aceleração de $\times 2,44$ para $\times 2,76$ mantendo a precisão média idêntica (87,27%).
• Comparação com Métodos de Evicção:
  • Supera significativamente métodos de evicção de tokens (PyramidInfer, FastKV, GemFilter). Por exemplo, em 128K, o Token Sparse Attenção alcança 86,84% de precisão versus 78,49% para o PyramidInfer em acelerações similares.
• Escalabilidade:
  • Os ganhos de aceleração aumentam com o comprimento do contexto (por exemplo, ganho mínimo em 4K, substancial em 128K) porque a esparsidade de atenção aumenta com o comprimento da sequência.
• Sobrecarga:
  • A sobrecarga para pontuação e gerenciamento de índices é mínima, representando <11% da latência total de atenção mesmo em altos níveis de esparsidade.

5. Significado

Este artigo aborda uma lacuna crítica na inferência de LLMs de contexto longo, afastando-se da evicção irreversível de tokens em direção a uma esparsificação dinâmica e intercalada.

• Insight Teórico: Demonstra que a importância do token é altamente dinâmica entre camadas e cabeças, necessitando de mecanismos de seleção flexíveis em vez de poda estática.
• Impacto Prático: Fornece uma otimização "plug-and-play" que pode ser aplicada sobre kernels de atenção existentes de última geração (FlashAttention, Minference, etc.) para alcançar acelerações massivas em cenários de contexto longo sem retreinar modelos ou sacrificar precisão significativa.
• Direção Futura: Estabelece um novo ponto de design para inferência escalável, sugerindo que a compressão reversível é uma estratégia superior para lidar com o gargalo $O(L^2)$ comparada à remoção permanente de tokens.