📊 statistics

Kernelized Linear Attention: Breaking the Capacity Wall with Symmetric Cones

Este artigo introduz o Kernelized Linear Attention (KATA), um novo framework que aproveita cones simétricos e características PSD de posto um para resolver o compromisso entre capacidade e interferência na atenção linear, alcançando uma recuperação associativa superior e um throughput significativamente maior que o FlashAttention-2, mantendo simultaneamente um desempenho de longo alcance quase perfeito com overhead reduzido de KV-cache.

Autores originais: Ayoub Ghriss, Sourav Chakraborty

Publicado 2026-07-21

📖 4 min de leitura☕ Leitura rápida

CC BY 4.0

Autores originais: Ayoub Ghriss, Sourav Chakraborty

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine que você esteja tentando construir um robô superinteligente que consiga ler um livro e se lembrar de cada detalhe, desde o nome de um personagem secundário até a cor exata de uma porta mencionada três capítulos atrás. No mundo da inteligência artificial, este é o trabalho de um "Transformer", um tipo de modelo que alimenta muitos dos chatbots e ferramentas que usamos hoje. O ingrediente secreto que torna esses robôs tão bons em lembrar é algo chamado "atenção". Pense na atenção como um holofote: quando o robô lê uma nova frase, o holofote brilha nas palavras mais importantes que ele viu anteriormente para ajudá-lo a entender a atual.

No entanto, há um problema. O holofote tradicional é incrivelmente poderoso, mas também incrivelmente pesado. À medida que a história fica mais longa, o holofote tem que escanear cada uma das palavras anteriores para encontrar a correta. Isso é como tentar encontrar uma agulha específica em um palheiro verificando cada pedaço de feno, um por um; leva uma eternidade e requer uma quantidade massiva de espaço de armazenamento (memória) para manter todas essas agulhas à mão. Cientistas têm tentado construir um holofote "linear" que seja mais rápido e leve, um que consiga lembrar das coisas sem precisar reescanear o livro inteiro toda vez. Mas essas versões mais rápidas costumam ter uma memória terrível: elas esquecem os detalhes importantes ou ficam confusas quando muitas coisas parecem semelhantes. Elas são rápidas, mas não são inteligentes o suficiente para lidar com histórias complexas.

É aqui que uma nova ideia chamada Kernelized Linear Attention (KATA) entra em cena. Os pesquisadores por trás deste artigo, Ayoub Ghriss e Sourav Chakraborty, decidiram resolver o problema da memória olhando para ele através da lente da geometria e do empacotamento. Eles perceberam que a razão pela qual os modelos rápidos esquecem as coisas é que eles estão tentando espremer muitas memórias em uma caixa pequena e lotada. Para corrigir isso, eles inventaram uma nova maneira de organizar memórias usando uma forma matemática chamada "cone simétrico".

Pense em uma memória como uma chave única. Nos modelos rápidos antigos, essas chaves eram como formas planas, 2D, que podiam facilmente se sobrepor e se misturar. O KATA, no entanto, usa uma forma 3D especial (especificamente, um "cone semidefinido positivo") para transformar essas chaves planas em algo mais robusto. É como pegar uma folha de papel plana e dobrá-la em um grou de origami complexo. Mesmo que duas folhas de papel pareçam semelhantes quando planas, seus grous dobrados podem ser totalmente diferentes e fáceis de distinguir. Ao usar esse truque de "dobra", o KATA pode empacotar exponencialmente mais memórias únicas no mesmo espaço sem que elas colidam umas com as outras.

O artigo mostra que esse truque geométrico funciona maravilhosamente bem. Eles construíram um novo tipo de mecanismo de atenção que não precisa armazenar uma lista massiva de cada palavra que já viu (o que economiza muita memória). Em vez disso, ele mantém um resumo compacto e organizado. Ao serem testados em tarefas que exigem a lembrança de detalhes específicos de textos longos — como encontrar uma palavra escondida em um mar de distrações — o KATA teve um desempenho quase tão bom quanto os modelos tradicionais pesados e lentos, mas com uma fração da memória. De fato, em alguns testes, ele conseguiu lembrar detalhes de textos 16 vezes mais longos do que aquilo para o qual foi treinado, algo que outros modelos rápidos geralmente falham em fazer.

Os pesquisadores não pararam apenas na teoria; eles construíram o código de computador real para rodar isso em placas gráficas modernas. Eles descobriram que seu novo método é incrivelmente rápido. Em alguns cenários, ele é até 11 vezes mais rápido que o padrão atual de atenção rápida, mantendo a precisão da memória. Eles também descobriram que, embora este novo método seja ótimo para memória pura, às vezes precisa de uma ajuda para entender o fluxo de uma história, sugerindo que os melhores modelos futuros podem combinar essa memória supereficiente com outras ferramentas para lidar tanto com fatos quanto com fluência.

Em resumo, o KATA é como dar ao robô um arquivo superorganizado onde cada arquivo tem uma forma 3D única que evita que ele se perca no meio da confusão. Ele prova que você não precisa escolher entre um robô rápido e um robô inteligente; com a forma geométrica certa, você pode ter ambos.

Resumo Técnico: Atenção Linear Kernelizada: Quebrando a Barreira de Capacidade com Cones Simétricos

1. Definição do Problema

Mecanismos de atenção linear oferecem inferência recorrente de tempo constante e representações de estado de tamanho fixo, abordando o custo de treinamento $O(T^2)$ e o uso de memória de KV-cache de $O(T)$ da atenção softmax padrão. No entanto, os modelos de atenção linear existentes sofrem uma degradação acentuada na recuperação associativa (associative recall) conforme o comprimento da sequência aumenta. Essa falha é atribuída à "colisão de memória", onde muitos vínculos de chave-valor devem compartilhar um único estado fixo, causando o colapso do desempenho de recuperação em tarefas de contexto longo.

Os atuais RNNs lineares com portões (ex: DeltaNet, Mamba2, Gated DeltaNet) tentam mitigar isso por meio de portões de esquecimento aprendidos e operações de apagamento endereçadas por conteúdo. Contudo, os limites teóricos de suas geometrias de características permanecem opacos. Especificamente, não há uma compreensão fundamentada sobre:

O que limita a capacidade de recuperação em um tamanho de estado fixo?
Quais mapas de características finitos preservam a isotropia do softmax?
Os mapas de características devem ser não negativos e, se sim, quais geometrias são admissíveis?

2. Metodologia

2.1 Estrutura Teórica: Cones Simétricos e KATA

Os autores formulam a recuperação de atenção como um problema de empacotamento esférico. Eles introduzem a Ativação de Atenção Linear Kernelizada (KATA), um framework derivado de primeiros princípios que certifica pesos de atenção não negativos através de cones homogêneos autoduais.

Premissas Geométricas: O framework assume (1) pesos de atenção não negativos via um cone autodual ( $K = K^*$ ), (2) invariância isométrica do mapa de características e (3) homogeneidade da geometria latente.
Classificação de Koecher–Vinberg: Sob estas premissas, as geometrias de características admissíveis são restritas aos fatores irredutíveis do teorema de Koecher–Vinberg: o ortante positivo, o cone de Lorentz e o cone semidefinido positivo (PSD) simétrico.
Portão de Saída (Output Gating): Os autores provam que o denominador da atenção linear normalizada emerge naturalmente como um portão de saída convexo condicionado ao token e livre de parâmetros. Este portão realiza uma interpolação convexa entre a leitura anterior e o valor atual, eliminando a necessidade de mecanismos de portão aprendidos explícitos na formulação base.

2.2 Teoria de Capacidade

O artigo analisa a "capacidade associativa" (o número de endereços de chaves distinguíveis) para cada geometria de cone:

Ortante e Lorentz: Estas geometrias estão sujeitas ao limite de Welch e à parede de Rankin, respectivamente. Elas oferecem capacidade linear ( $O(d)$ ) ou são estritamente limitadas por limiares de interferência (ex: a interferência de Lorentz não pode cair abaixo de $1/2$ sem forçar chaves mutuamente obtusas).
Cone PSD (Rank-One): Os autores propõem o uso de características semidefinidas positivas (PSD) de rank-um, $\psi(u) = uu^\top$ $ψ (u) = u u^{⊤}$ . Este mapa eleva ao quadrado o produto interno bruto ( $\langle u_i, u_j \rangle^2$ $⟨ u_{i}, u_{j} ⟩^{2}$ ), alterando a geometria da interferência.
- Resultado Teórico: Enquanto os vetores brutos estão sujeitos ao piso de Welch, o emparelhamento ao quadrado permite uma capacidade exponencial na dimensão de projeção $p$ para qualquer tolerância fixa $\epsilon > 0$ . Especificamente, a capacidade escala como $T \geq \frac{1}{2}(1-\epsilon)^{-(p-1)/2}$ .
- Relação Sinal-Ruído: Os autores derivam uma análise de razão sinal-ruído de potência (pSNR) mostrando que o mapa PSD estende o limiar de recuperação idealizado de $\Theta(p)$ (linear) para $\Theta(p^2)$ , efetivamente ampliando o estado sem adicionar parâmetros.

2.3 Implementação Alinhada ao Hardware

Para realizar esses ganhos teóricos, os autores implementam KATA usando kernels Triton fundidos em GPUs NVIDIA.

Paralelismo por Blocos (Chunked Parallelism): A sequência é dividida em blocos. A atenção intra-bloco utiliza Tensor Cores, enquanto a propagação de estado inter-bloco é tratada via um scan associativo (soma de prefixo paralelo).
Redução de Complexidade: Isso reduz a profundidade de recorrência inter-bloco de $O(N_C)$ (sequencial) para $O(\log N_C)$ , onde $N_C$ é o número de blocos.
Dois Pontos de Operação:
1. Quadrático ( $O(T^2)$ ) Forward: Um kernel estilo FlashAttention que evita a materialização do estado completo, alcançando até 1.6 $\times$ o throughput do FlashAttention-2.
2. Linear ( $O(T)$ ) State: Uma forma de estado em blocos exata que atinge $\sim$ 11 $\times$ o throughput do FlashAttention-2 em 131k tokens.
Variantes: O framework inclui KATA-Mg (múltiplos fatores PSD) e KATA- $\Sigma$ g (fatores PSD somados) para equilibrar o tamanho do estado e a eficiência de empacotamento.

3. Principais Contribuições

Classificação de Cones para Atenção Linear: O artigo caracteriza as geometrias de características não negativas admissíveis (Ortante, Lorentz, PSD) sob as restrições de pesos não negativos e invariância isométrica, fundamentando a atenção linear na classificação de Koecher–Vinberg.
Teoria de Capacidade de Recuperação Associativa: Ao tratar a recuperação como um empacotamento esférico discreto, os autores provam a parede de Rankin de Lorentz e os limites de empacotamento dependentes da dimensão para raios PSD de rank-um. Eles demonstram que características PSD de rank-um oferecem uma vantagem de capacidade exponencial sobre características lineares abaixo do piso de Welch.
Kernels Otimizados para Hardware: Os autores liberam kernels Triton fundidos para a recorrência KATA. Estes kernels utilizam um scan associativo para reduzir a profundidade inter-bloco, alcançando ganhos significativos de throughput sobre baselines de atenção linear sequencial e FlashAttention-2.
Validação Empírica: Experimentos em tarefas sintéticas (MQAR, sobreposição de chave repetida, indução) e modelos de linguagem de 340 milhões de parâmetros mostram que as variantes KATA (particularmente as do cone PSD) superam o DeltaNet com portão e outros baselines lineares em recuperação associativa, mantendo uma recuperação quase perfeita em extrapolação de comprimento fora do domínio de 16 $\times$ .

4. Resultados

Benchmarks Sintéticos:
- MQAR (Multi-Query Associative Recall): KATA-M1 retém 0.985 de precisão em 16 $\times$ o comprimento de contexto de treinamento, aproximando-se do desempenho do softmax utilizando aproximadamente um quarto das entradas de KV-cache.
- Sobreposição de Chave Repetida (Repeated-Key Overwrite): Variantes com a regra delta (DeltaKATA) alcançam precisão quase perfeita de sobreposição, corrigindo vínculos obsoletos em endereços repetidos.
- Needle-in-a-Haystack (NIAH): Em tarefas de alta entropia (UUIDs), KATA-M2 e DeltaKATA-M1 superam significativamente o Gated DeltaNet, que degrada para uma recuperação próxima de zero (0.004) em UUIDs.
Pré-treinamento de Modelos de Linguagem:
- Modelos pré-treinados de 340M de parâmetros mostram que as variantes PSD mantêm precisão zero-shot competitiva em benchmarks padrão (ARC, PIQA, etc.), enquanto exibem um "trade-off de fluência" (maior perplexidade que o softmax, mas melhor recuperação que outros modelos lineares).
- Indução: KATA preserva capacidades de indução quase perfeitas, uma tarefa onde a atenção linear frequentemente falha.
Desempenho:
- O kernel de estado linear KATA-M2 alcança $\sim$ 11 $\times$ o throughput de forward do FlashAttention-2 em 131k tokens.
- O forward quadrático KATA-M1 atinge $\sim$ 1.6 $\times$ o throughput do FlashAttention-2.
- Implementações de scan associativo apresentam, em média, $\sim$ 2.4 $\times$ o throughput de blocos de atenção linear sequencial correspondentes.

5. Significância e Alegações

O artigo afirma resolver a "barreira de capacidade" da atenção linear ao identificar o substrato geométrico correto (cones simétricos, especificamente o cone PSD) para a não negatividade da atenção.

Insight Teórico: Estabelece que o "preço" do estado constante na atenção linear não é inerente à própria linearidade, mas à escolha da geometria de características. Ao mudar para características PSD de rank-um, a capacidade pode ser exponencialmente aumentada sem aumentar o número de parâmetros, apenas a dimensão do estado.
Impacto Prático: O trabalho demonstra que uma recuperação associativa de alta capacidade é alcançável com inferência de tempo linear e estados de tamanho fixo, desde que o mapa de características esteja alinhado ao hardware (via kernels fundidos) e a geometria seja teoricamente sólida (via cones simétricos).
Limitações: Os autores observam que, embora a capacidade geométrica seja alta, a aprendibilidade é um fator; os modelos podem ter dificuldade em aprender o empacotamento esférico de alta dimensão (ex: KATA-M1 tem uma capacidade teórica maior, mas às vezes performa pior que KATA-M2 em tarefas contextuais devido à dinâmica de treinamento). Além disso, a implementação atual foca em cones de números reais, deixando extensões para números complexos como trabalho futuro.

Em resumo, o KATA fornece uma ponte fundamentada entre os limites teóricos da memória associativa e as restrições práticas de inferência eficiente em hardware, oferecendo um caminho para Transformers de tempo linear que não sacrificam a recuperação de longo alcance.