SpotAttention: Plug-In Block-Sparse Routing for Pretrained Long-Context Transformers

Resumo Técnico: SpotAttention

Declaração do Problema

As capacidades de contexto longo tornaram-se um padrão em Grandes Modelos de Linguagem (LLMs) pré-treinados, permitindo tarefas como o processamento de bases de código inteiras ou a fundamentação de respostas em evidências extensas. No entanto, o custo computacional da auto-atenção (self-attention) escala mal com o comprimento da sequência. Durante a fase de preenchimento (prefill), o cálculo cresce quadraticamente com o comprimento da sequência e, durante a fase de decodificação (decode), cada novo token exige a releitura de um cache de Chave-Valor (KV) que cresce linearmente com o tamanho do contexto. Embora o FlashAttention tenha melhorado a eficiência de I/O ao evitar a materialização da matriz de atenção completa, ele não alterou o custo fundamental quadrático do prefill nem o crescimento linear das leituras de memória por etapa.

Métodos de atenção esparsa existentes tentam mitigar isso atendendo apenas a um subconjunto relevante de tokens passados. No entanto, selecionar este subconjunto é frequentemente caro. Métodos como o Quest pontuam blocos de tokens, e o Twilight adiciona a poda top-$p$, mas o próprio seletor torna-se um novo gargalo. Além disso, seletores sem treinamento (training-free) dependem de heurísticas desenhadas à mão que podem não generalizar bem ou podem exigir múltiplas etapas de seleção, incorrendo em overhead cumulativo.

Metodologia: SpotAttention

O SpotAttention é um mecanismo de roteamento de esparsidade em blocos (block-sparse) do tipo plug-in, projetado para ser acoplado a transformers já treinados e congelados. Ele consiste em um "seletor" (ou indexador) leve acoplado a cada camada de atenção total.

1. Arquitetura do Seletor

O seletor é um pontuador Q-K multi-cabeça minúsculo com largura reduzida ($d_{idx}$) e $H_{idx}$ cabeças. Em vez de pontuar tokens individuais, ele opera em blocos de queries e keys (por exemplo, 16 tokens por bloco).

• Mecanismo de Pontuação: Para uma query na posição $q$ e um bloco de key na posição $t$, o seletor computa uma pontuação ReLU ponderada por cabeça:
  $$s_s[q, t] = \sum_{h=1}^{H_{idx}} w_h \text{ReLU}(\mathbf{q}^I_h \cdot \mathbf{k}^I_{h,t})$$
  onde $\mathbf{q}^I_h$ e $\mathbf{k}^I_{h,t}$ são projeções por cabeça, e $w_h$ são pesos de mistura derivados do estado oculto.
• Eficiência: Ao pontuar por blocos, o processo de seleção executa como uma única multiplicação de matriz em tensor-core, evitando o overhead de operações ao nível de token.

2. Treinamento via Destilação KL

O seletor é treinado independentemente contra a distribuição de atenção densa do backbone congelado (o "professor").

• Objetivo: O treinamento minimiza a divergência de Kullback-Leibler (KL) entre a distribuição de atenção do professor ($p_\tau$) e a distribuição estimada pelo seletor ($p_\sigma$).
• Variantes: O artigo explora duas variantes de KL:
  • DenseKL: Computa a perda sobre todas as $T$ keys.
  • SparseKL: Computa a perda apenas sobre o próprio conjunto top-$K$ selecionado pelo seletor, com distribuições renormalizadas.
• Convergência: O seletor converge rapidamente (dentro de ~200 passos) através de vários tamanhos de backbone (4B a 32B) e arquiteturas, aprendendo a corresponder à forma da distribuição do professor.

3. Regra de Seleção Dual Top-$p$

Na inferência, o seletor pontua todos os blocos passados. Para determinar em quais blocos deve atender, o SpotAttention introduz uma regra dual top-$p$ que lê o orçamento diretamente da distribuição calibrada do seletor:

1. Blocos Reservados: Um prefixo fixo de blocos "sink" (ex: primeiros 128 tokens) e um sufixo fixo de blocos de "recência" (ex: últimos 256 tokens) são sempre retidos para evitar que sejam dominados pela seleção de núcleo (nucleus selection).
2. Seleção de Núcleo (Nucleus Selection): A amostragem top-$p$ (núcleo) é aplicada à distribuição residual (blocos excluindo sinks e recência). A seleção para quando a massa de probabilidade cumulativa atinge um limiar $p$.
3. Orçamento Dinâmico: Esta abordagem resulta em um orçamento variável por query e por camada, adaptando-se ao conteúdo sem uma etapa de poda separada. Um piso mínimo ($K_{min}$) evita casos degenerativos onde o núcleo é muito pequeno.

4. Quantização

O K-cache do seletor pode ser quantizado para microescala INT4 ou FP4. Como o cache do seletor é distinto do cache KV do backbone, isso reduz a pegada de memória do seletor em aproximadamente 3,5$\times$ sem afetar a precisão ou a acurácia do backbone.

Principais Contribuições

1. Atenção de Esparsidade em Blocos Plug-In: Um seletor leve treinado via destilação KL que se acopla a transformers congelados. Ele executa como um matmul de tensor-core único, evitando o pesado overhead de seletores vistos em trabalhos anteriores.
2. Regra Dual Top-$p$: Um mecanismo de seleção adaptável ao conteúdo que deriva o orçamento por query e por camada diretamente da distribuição aprendida. Ele reserva blocos de sink e recência para garantir estabilidade, igualando a acurácia de $K$ estático com um orçamento médio menor.
3. Comportamento de Treinamento e Inferência: O artigo demonstra que o seletor aprende a assimetria de orçamento por camada do backbone (camadas iniciais atendem a mais contexto do que camadas posteriores), um padrão que o $K$ estático não consegue replicar. O método generaliza de um comprimento de treinamento de 16K para uma avaliação de 128K sem degradação de acurácia.

Resultados

O método foi avaliado em cinco backbones (Qwen3: 4B–32B; Qwen3.5: 4B–9B) através de quatro benchmarks (RULER, BABILong, InfiniteBench, LongBench-v2) em comprimentos de contexto de até 128K.

• Acurácia: O SpotAttention iguala a acurácia do backbone denso em comprimentos de contexto de até 128K (8$\times$ o comprimento de treinamento). Com um orçamento de $K=0,5L$, ele permanece dentro de um erro padrão de bootstrap do baseline denso em todos os modelos.
• Latência e Throughput:
  • Em $L=128K$ no Qwen3-8B, o SpotAttention roda 3,9$\times$ mais rápido que o FlashAttention e 1,8$\times$ mais rápido que o Twilight (o baseline sem treinamento mais forte).
  • O throughput de decodificação é mantido em 70 tokens/s a 128K e 36 tokens/s a 512K, enquanto o FlashAttention e o Twilight sofrem OOM (falta de memória) em 512K e 256K, respectivamente.
  • O próprio seletor é significativamente mais rápido que o seletor do Twilight (1,55 ms vs 4,59 ms em 128K).
• Quantização: Quantizar o K-cache do seletor para INT4 ou FP4 reduz sua pegada de memória em 3,5$\times$ sem perda mensurável de acurácia. Embora isso não melhore a latência de decodificação de fluxo único (dominada pela largura de banda do KV do backbone), reduz significativamente o uso de memória e os custos de prefill.

Significância e Alegações

O artigo posiciona o SpotAttention como uma solução para o "gargalo do seletor" na atenção esparsa. Ao aprender uma distribuição calibrada via destilação KL, o seletor permite um orçamento baseado em núcleo e dinâmico que se adapta às necessidades específicas de cada query e camada.

Os autores afirmam que:

• Eficiência: Alcança o maior throughput entre os métodos comparados em contextos longos (até 512K), superando as limitações de memória que interrompem os baselines densos e outros métodos esparsos.
• Generalização: Um seletor treinado em um contexto de 16K generaliza efetivamente para 128K e além sem necessidade de retreinamento.
• Modularidade: É uma solução do tipo "plug-in" que não requer o retreinamento do backbone, tornando-o aplicável a modelos pré-treinados existentes.
• Limitações: Os autores observam que o maior backbone avaliado é o de 32B, e os resultados são atualmente limitados à família Qwen em benchmarks de inglês. Eles também reconhecem que a latência de prefill esparso não foi testada, embora a arquitetura suporte isso, e que a regra dual top-$p$ depende de o seletor estar bem calibrado (o que se mantém sob o treinamento SparseKL, mas colapsa sob o DenseKL).