Don't Read Everything: A Curvature-Conditioned Query for Linear Attention

Resumo Técnico: Não Leia Tudo: Uma Consulta Condicionada à Curvatura para Atenção Linear

1. Definição do Problema

Os mecanismos de atenção linear emergiram como uma solução para o custo computacional quadrático ($O(T^2)$) e a sobrecarga de memória da atenção softmax padrão em Transformers. Ao substituir o produto escalar softmax por um produto escalar kernelizado, a atenção linear reformula a operação de leitura como uma Rede Neural Recorrente (RNN) linear com um estado de valor matricial, alcançando complexidade de tempo linear ($O(T)$) e cache constante ($O(1)$).

No entanto, apesar dos avanços recentes em mecanismos de escrita (ex: gating dependente de dados em GLA, Mamba e RWKV; atualizações de regra-delta em Gated DeltaNet), os modelos de atenção linear subperformam consistentemente a atenção softmax padrão em recuperação de contexto in-context e tarefas de contexto longo. Os autores argumentam que essa lacuna de desempenho decorre não apenas de como a informação é escrita na memória, mas, crucialmente, de como ela é lida.

Na atenção softmax padrão, o normalizador log-sum-exp cria um ambiente competitivo onde a chave mais distinta domina a saída, suprimindo efetivamente correspondências que residem em regiões de alta densidade na memória. A atenção linear descarta esse normalizador para ganhar eficiência. Consequentemente, cada chave passada contribui aditivamente para a saída para sempre. Isso faz com que "alvos úteis" (as "agulhas") sejam diluídos pelo "volume" de vetores armazenados (o "palheiro"), particularmente quando a consulta (query) se sobrepõe a direções onde a memória é densa. As soluções existentes focam no lado da escrita (decidindo o que entra no estado), mas não conseguem reagir a quais direções uma consulta futura lerá contra.

2. Metodologia: Consulta Condicionada à Curvatura (CCQ)

O artigo propõe a Consulta Condicionada à Curvatura (CCQ), um mecanismo de correção leve de tempo de leitura, inspirado pela geometria local do normalizador de partição logarítmica faltante do softmax.

Derivação Teórica

Os autores observam que o termo de competição do softmax, $A_t(q_t) = \log \sum_{j \le t} \exp(k_j^\top q_t)$, é uma função apenas da consulta $q_t$ uma vez que as chaves estão fixas. Eles realizam uma expansão de Taylor de segunda ordem desta função de partição logarítmica em torno do ponto de atenção isotrópica ($q=0$):

$$A_t(q) \approx \text{const} + \mu_t^\top q + \frac{1}{2} q^\top (\bar{C}_t - \mu_t \mu_t^\top) q$$

Aqui, o Hessiano da expansão corresponde à covariância amostral contínua das chaves, $\Sigma_t = \text{Cov}_{\le t}[k]$. Esta matriz de covariância captura as direções de alta densidade na memória.

O Mecanismo

A CCQ utiliza este insight geométrico para construir uma consulta "limpa" antes que ela interaja com o estado da atenção linear. Em vez de ler o estado com a consulta bruta $q_t$, o modelo aplica uma contração linear:

$$q_t^{\text{clean}} = (I - \lambda_t \Sigma_t) q_t = q_t - \lambda_t \Sigma_t q_t$$

• $\Sigma_t$: A covariância contínua das chaves, mantida através do mesmo mecanismo recorrente ou por blocos (chunkwise) do estado principal (requerendo uma matriz adicional $d_k \times d_k$ e um vetor $d_k$ por cabeça).
• $\lambda_t$: Um gate escalar condicionado à consulta ($\lambda_t = \sigma(W_\lambda q_t + b_\lambda)$), controlando a agressividade da contração.
• Efeito: O operador contrai a consulta ao longo das direções de alta variância (alta densidade) da memória. Isso penaliza chaves que se alinham com o volume dos vetores armazenados, efetivamente "puxando" a consulta para longe de direções congestionadas para favorecer alvos isolados.

Integração

A CCQ modifica apenas o passo de leitura. Ela é totalmente composível com qualquer backbone de atenção linear (ex: GLA, Gated DeltaNet) porque não altera a regra de escrita ou a recorrência do estado $S_t = \text{Update}(S_{t-1}; k_t, v_t, \dots)$. A saída final é computada como $o_t = S_t q_t^{\text{clean}}$.

3. Principais Contribuições

1. Identificação do Gargalo do Passo de Leitura: O artigo identifica o passo de leitura como um alavancagem distinta e subexplorada para fechar a lacuna de desempenho entre a atenção softmax e a linear, complementando as existentes soluções de escrita.
2. Mecanismo CCQ: A proposta de uma contração linear consciente da curvatura derivada de uma única expansão de Taylor do log-particionamento do softmax. É aditiva a qualquer regra de escrita de atenção linear.
3. Validação Empírica: Experimentos extensos mostrando que a CCQ melhora a perplexidade, a precisão zero-shot em tarefas downstream, a recuperação de agulha-no-palheiro (S-NIAH) e a extrapolação de comprimento e o desempenho no LongBench quando anexada aos backbones GLA e Gated DeltaNet.

4. Resultados Experimentais

Os autores avaliaram a CCQ em modelos de 500M e 1.3B de parâmetros treinados no subset FineWeb-Edu.

• Modelagem de Linguagem e Precisão Downstream:
  • No modelo de 1.3B, a CCQ-GLA alcançou a melhor perplexidade de WikiText (16.92) e precisão LAMBADA (47.82%) entre os modelos recorrentes, superando seu backbone por margens significativas.
  • A CCQ-Gated DeltaNet alcançou a melhor precisão média em 7 tarefas (56.24%) entre os modelos recorrentes.
• Recuperação Sintética (S-NIAH):
  • A CCQ resgatou significativamente os backbones em tarefas de extrapolação de comprimento. Por exemplo, no modelo de 1.3B, a CCQ-GLA melhorou a precisão S-NIAH-1 (pass-key) em contexto de 4K em +49.0% sobre o backbone GLA.
  • Enquanto os Transformers padrão colapsaram para 0% de precisão em contexto de 8K (devido ao limite de treinamento de 4K), as variantes recorrentes da CCQ mantiveram uma precisão substancial.
• Extrapolação de Comprimento:
  • Em seis corpora de contexto longo (GovReport, QMSum, etc.), as variantes CCQ consistiram em atingir perplexidades menores que seus backbones de 4K a 20K de contexto.
  • Em 20K, as variantes CCQ superaram seus respectivos backbones em 1 a 2.5 pontos de perplexidade.
• Compreensão de Contexto Longo (LongBench):
  • A CCQ-GLA alcançou a melhor pontuação média (15.77) entre os modelos recorrentes no LongBench de 14 tarefas, superando o Mamba2 em +2.64 pontos.
  • Ambas as variantes CCQ mostraram melhorias marcantes em tarefas que exigem geração de saída longa (ex: GovReport, MultiNews) e raciocínio few-shot.

5. Significância e Alegações

O artigo alega que a lacuna de desempenho na atenção linear não é meramente uma questão de capacidade, mas sim geométrica: a falta de competição no passo de leitura permite que a densidade da memória afogue alvos específicos.

• Insight Geométrico: Os autores demonstram que a competição "faltante" no da atenção linear pode ser aproximada pela curvatura local (Hessiano) do log-particionamento do softmax, que coincide com a covariância das chaves.
• Correção Consciente da Memória: Diferente das projeções de consulta baseadas apenas no input, a CCQ é "consciente da memória" porque usa as estatísticas contínuas das chaves armazenadas para condicionar a consulta.
• Modularidade: O método é apresentado como uma substituição direta para o passo de leitura, compatível com qualquer backbone de atenção linear sem modificar a dinâmica de escrita.
• Escopo Modesto: Os autores explicitamente notam limitações, afirmando que os resultados são atualmente limitados às escalas de 500M/1.3B, dois backbones específicos (GLA, Gated DeltaNet) e benchmarks específicos. Eles não alegam que a CCQ resolve todos os problemas da atenção linear ou que se generaliza para escalas de 7B+ ou modelos de espaço de estados como o Mamba sem investigação adicional.

Em resumo, o artigo argumenta que, ao "não ler tudo" (isto é, contrair a consulta ao longo das direções densas da memória), a atenção linear pode recuperar grande parte da capacidade de recuperação e de contexto longo perdida ao descartar o normalizador softmax, com um custo computacional mínimo.