LoSA: Locality Aware Sparse Attention for Block-Wise Diffusion Language Models

Resumo Técnico: LoSA para Modelos de Linguagem de Difusão em Blocos

1. O Problema: Inchaço de KV (KV Inflation) em DLMs

Os Modelos de Linguagem de Difusão (DLMs) baseados em blocos oferecem uma alternativa promissora aos modelos autoregressivos, permitindo a geração de múltiplos tokens em qualquer ordem dentro de um bloco. No entanto, eles enfrentam um gargalo crítico em cenários de contexto longo: a atenção baseada em memória (memory-bound attention).

• Falha da Atenção Esparsa Ingênua: Métodos de atenção esparsa comuns em LLMs autoregressivos falham quando aplicados diretamente a DLMs. Isso ocorre devido ao problema de Inchaço de KV (KV Inflation).
• Mecanismo do Problema: Em um bloco de $B$ tokens, cada token de consulta (query) seleciona um subconjunto diferente de posições de chave/valor (KV) no prefixo. Como a atenção é limitada pela largura de banda de memória, o custo não é determinado pelo número de tokens selecionados por cada consulta, mas sim pela união de todas as posições KV selecionadas por todo o bloco.
• Consequência: A união das posições KV acessadas torna-se muito grande (potencialmente até $B$ vezes maior que o orçamento por consulta), anulando os ganhos de velocidade esperados pela esparsidade e mantendo a latência alta.

2. Observação Fundamental: Localidade de Mudanças de Representação

Os autores identificaram uma estrutura única nos DLMs de difusão que não existe na decodificação autoregressiva: a Localidade de Mudanças de Representação entre passos de denoising consecutivos.

• Tokens Ativos vs. Estáveis: Entre dois passos de denoising ($t-1$ e $t$), apenas uma pequena fração de tokens (os tokens ativos) sofre mudanças significativas em seus estados ocultos (hidden states), geralmente aqueles que estão sendo atualizados de [MASK] para tokens reais.
• Tokens Estáveis: A maioria dos tokens (tokens estáveis) mantém seus estados ocultos quase constantes. Consequentemente, seus padrões de atenção e saídas de atenção permanecem aproximadamente inalterados.
• Visualização: Gráficos de MSE (Erro Quadrático Médio) entre passos mostram que a mudança de representação é altamente localizada, com a maioria dos tokens tendo variação mínima.

3. Metodologia: LoSA (Atenção Esparsa Consciente da Localidade)

Com base na observação acima, os autores propõem o LoSA, um método que reduz drasticamente o tráfego de memória ao reutilizar resultados de atenção cacheados para tokens estáveis.

Fluxo de Trabalho do LoSA:

1. Poda de Localidade (Locality Pruning):
   • Calcula-se uma pontuação de localidade (MSE) para cada token no bloco, comparando os vetores de consulta ($Q$) do passo atual com o anterior.
   • Identifica-se o conjunto $A$ de tokens ativos (aqueles com maior mudança) e o conjunto de tokens estáveis.
2. Reutilização para Tokens Estáveis:
   • Para tokens estáveis, o LoSA não recalcula a atenção. Ele reutiliza diretamente as estatísticas de atenção do prefixo (saída $o_p$ e normalizador logarítmico $L_p$) do passo anterior.
   • Isso preserva a informação de atenção densa completa para a maioria dos tokens, evitando a perda de precisão comum em métodos esparsos.
3. Atenção Esparsa para Tokens Ativos:
   • Apenas para o conjunto reduzido de tokens ativos ($A$), aplica-se um seletor de atenção esparsa (como o algoritmo QUEST) para escolher as posições de KV relevantes.
   • Como o conjunto de consultas ativas é muito menor que o tamanho do bloco ($|A| \ll B$), a união das posições KV selecionadas é drasticamente reduzida.
4. Fusão e Atualização:
   • Calcula-se a atenção dentro do bloco (suffix) densamente (já que $B$ é pequeno).
   • Usa-se a técnica online-softmax para fundir as contribuições do prefixo (reutilizado ou recalculado) e do bloco atual.
   • As estatísticas atualizadas são armazenadas no cache para o próximo passo.

4. Contribuições Principais

• Identificação do Problema KV Inflation: Demonstrar que a aplicação ingênua de atenção esparsa em DLMs falha devido à inflação da união de índices KV, tornando a esparsidade ineficiente.
• Descoberta de Localidade: Evidenciar que as mudanças de representação em DLMs são altamente localizadas, permitindo a distinção entre tokens ativos e estáveis.
• Proposta do LoSA: Um algoritmo que combina reutilização de cache (para estabilidade) e atenção esparsa dinâmica (para mudanças), reduzindo o número de índices KV carregados em ~1,5x na prática.
• Eficiência e Precisão: O método não apenas acelera a inferência, mas também melhora a precisão em comparação com a atenção esparsa padrão, pois preserva a informação densa para a maioria dos tokens.

5. Resultados Experimentais

O LoSA foi avaliado em vários modelos (Trado-8B, Trado-4B, SDAR-8B) e benchmarks (LongBench, raciocínio comum).

• Precisão (LongBench):
  • O LoSA supera consistentemente o baseline QUEST (atenção esparsa adaptada).
  • Em níveis de esparsidade agressivos (orçamento de 128 tokens), o LoSA alcançou até +10,43 pontos de precisão média em relação ao QUEST no modelo Trado-8B.
  • Mantém precisão próxima à do modelo denso (dense baseline) mesmo com alta esparsidade.
• Densidade de KV:
  • O LoSA opera com 1,54x menor densidade de cache KV em comparação com o QUEST, significando que carrega menos dados da memória.
• Aceleração (Latência):
  • Em GPUs NVIDIA RTX A6000, o LoSA alcançou um speedup de 4,14x em relação à atenção densa, superando o QUEST (3,26x).
  • Em GPUs mais novas (RTX 5090), o speedup foi de 3,67x, demonstrando que a melhoria escala com hardware mais rápido, mantendo o gargalo na largura de banda de memória.
• Generalização: Os resultados foram consistentes em diferentes tamanhos de modelo (4B e 8B) e tamanhos de bloco (16 e 32).

6. Significado e Impacto

O trabalho é significativo porque resolve uma barreira fundamental para a escalabilidade de Modelos de Linguagem de Difusão em contextos longos.

• Viabilidade Prática: Ao mitigar o problema de KV Inflation, o LoSA torna a inferência de DLMs competitiva em termos de latência e custo de memória em relação aos modelos autoregressivos tradicionais.
• Eficiência de Memória: A estratégia de reutilizar resultados de cache para tokens estáveis é uma inovação que reduz o tráfego de memória sem sacrificar a qualidade do modelo, algo que métodos puramente esparsos não conseguem fazer.
• Futuro: Abre caminho para a implementação de DLMs em sistemas de serviço de larga escala, onde a eficiência de memória é crítica.

Em resumo, o LoSA transforma a natureza iterativa da difusão de um obstáculo em uma vantagem, explorando a estabilidade temporal dos tokens para otimizar radicalmente o acesso ao cache de atenção.