Resumo Técnico: Echo: Recordação Associativa sem KV-Cache com Operadores Espectrais de Koopman

1. Declaração do Problema

O raciocínio em contextos longos e a chamada de ferramentas por agentes geram rastros que abrangem dezenas de milhares de tokens. Embora arquiteturas baseadas em Transformers dependam de caches de Chave-Valor (KV) para manter a recuperação endereçada por conteúdo, esses caches crescem linearmente com o comprimento da sequência, criando um gargalo de memória em hardware comercial. Por outro lado, Modelos de Espaço de Estado (SSMs) como o Mamba oferecem inferência em tempo linear e memória constante por passo, comprimindo a sequência em um estado recorrente de tamanho fixo. No entanto, essa compressão é perdedora: os SSMs sofrem de um "penhasco de memória", onde a precisão da recuperação colapsa exponencialmente assim que a lacuna entre um fato armazenado e sua consulta excede o horizonte efetivo do estado recorrente.

Remédios atuais envolvem arquiteturas híbridas que intercalam camadas de SSM com blocos de atenção causal. Embora isso restaure as capacidades de recuperação, reintroduz o crescimento linear de memória do cache KV. O artigo postula que a recuperação endereçada por conteúdo pode ser redefinida como regressão de crista (kernel ridge regression), onde estatísticas suficientes são acumuladas como somas em execução em memória fixa, eliminando assim o cache KV completamente sem sacrificar o desempenho de recuperação.

2. Metodologia: Echo e Atenção Espectral de Koopman (SKA)

Os autores introduzem o Echo, uma arquitetura sem KV-cache construída em torno da Atenção Espectral de Koopman (SKA), um substituto direto para camadas de atenção padrão. O Echo mantém SSMs (especificamente Mamba-2) para processamento sequencial local, mas substitui mecanismos de recuperação global pela SKA.

2.1 Mecanismo Central: Atenção Espectral de Koopman

A SKA estima um modelo dinâmico linear compacto a partir de estatísticas suficientes em fluxo para recuperar valores em posições de consulta sem materializar uma matriz de atenção $T \times T$ ou manter um cache KV.

• Acumulação de Estatísticas Suficientes: A SKA acumula três estatísticas aditivas em memória $O(r^2 + rP)$ (onde $r$ é um rank de projeção pequeno e $P$ é a dimensão da cabeça):

  1. Uma matriz Gram regularizada $G = \sum \hat{z}_t \hat{z}_t^\top + \epsilon I$.
  2. Uma covariância intertemporal $M = \sum \hat{z}_t \hat{z}_{t-1}^\top$.
  3. Uma covariância valor-chave $C_v = \sum v_t \hat{z}_t^\top$.
     Essas estatísticas são estatísticas suficientes exatas para o cálculo de recuperação, o que significa que sequências mais longas melhoram o condicionamento das estimativas em vez de degradá-las.

• Estimação do Operador e Recuperação:

  • O sistema ajusta um operador de Koopman branqueado $A_w = L^{-1} M L^{-\top}$ (onde $L$ é o fator de Cholesky da matriz Gram regularizada) via regressão de crista.
  • A leitura de valor é computada como $B_v = C_v \tilde{G}^{-1}$, que é o minimizador único do objetivo de mínimos quadrados regularizado $\min_B \sum \|v_t - B\hat{z}_t\|^2 + \epsilon \|B\|_F^2$.
  • Diferentemente da atenção linear (que corresponde a um único passo de gradiente a partir de zero) ou da atenção softmax (que aproxima implicitamente essa solução via otimização não convexa), a SKA computa a solução de crista de forma fechada diretamente via fatoração de Cholesky.

• Filtragem Espectral:

  • O operador branqueado $A_w$ é normalizado espectralmente para limitar seu raio espectral.
  • Um filtro iterado por potência ($A_w^K$) é aplicado à consulta branqueada. Este filtro amplifica modos eigen persistentes (onde $|\lambda| \approx 1$, representando ligações estáveis) e suprime modos transitórios (onde $|\lambda| \ll 1$, representando ruído ou distratores). Essa seletividade espectral é uma capacidade não presente em objetivos de regressão padrão sozinhos.

2.2 Componentes da Arquitetura

• Espinha Dorsal SSM: A maioria das camadas usa Mamba-2 para processamento local.
• MLP Espectral de Koopman: As camadas feedforward padrão SwiGLU são substituídas por um MLP espectral que aplica uma rotação em bloco-diagonal (constrita ao disco unitário) à entrada. Este design preserva o rank do gradiente e reduz a contagem de parâmetros em um terço em comparação com o SwiGLU, mantendo a preservação de norma.
• Inferência: O Echo mantém memória $O(1)$ por camada independentemente do comprimento do contexto. O estado consiste nos acumuladores ($G, M, C_v$) e uma norma máxima escalar, totalizando aproximadamente 77 KB para um modelo de 50M de parâmetros.

3. Contribuições Principais

1. Eliminação do Penhasco de Memória: O artigo demonstra que a falha de recuperação de SSMs puros é uma consequência estrutural da recorrência multiplicativa, não uma limitação de capacidade. O Echo elimina esse penhasco usando um solucionador de crista de forma fechada sobre estatísticas em fluxo.
2. Recordação Associativa sem KV-Cache: O Echo alcança recuperação perfeita em benchmarks onde SSMs puros falham e iguala modelos híbridos SSM+Atenção, mas com memória de inferência constante.
3. Conexão Teórica com Regressão: O trabalho formaliza o vínculo entre recuperação endereçada por conteúdo e regressão de crista, mostrando que a SKA resolve esse problema de otimização explicitamente via fatoração de Cholesky, contornando o cenário de otimização não convexa da atenção.
4. Preservação do Fluxo de Gradiente: Os autores fornecem uma análise teórica mostrando que a SKA e o MLP Espectral de Koopman introduzem zero compressão de rank adicional na passagem de trás (diferentemente da atenção softmax, que comprime gradientes via o rank efetivo dependente de entropia da matriz de atenção). Isso sugere um mecanismo para melhoria na eficiência de dados.

4. Resultados Experimentais

Os autores avaliaram o Echo em três regimes: tarefas de transferência sintéticas, Recordação Associativa de Múltiplas Consultas (MQAR) e modelagem de linguagem.

• Transferência Sintética (Escala Sub-Milionária): Em cinco tarefas de recuperação retidas (incluindo recuperação de prompt de sistema, rastro de ferramenta e agulha no palheiro), um modelo Echo de 982K parâmetros (SSM+SKA) superou tanto um SSM puro (1,11M de parâmetros) quanto um híbrido SSM+Atenção (998K parâmetros). O Echo alcançou médias de precisão de 81,1%, comparado a 54,4% para SSM puro e 76,2% para o híbrido. Crucialmente, o Echo manteve alta precisão (65% em 4.096 tokens) em tarefas de generalização de comprimento onde as linhas de base colapsaram para níveis próximos ao acaso.
• MQAR (Escala 50M): No benchmark MQAR com lacunas de distratores de até 4.096 tokens e 32 pares KV, um Mamba-2 puro marcou ~3% (acaso). O modelo Echo (Mamba-2+SKA) alcançou 100% de precisão de recuperação em todas as configurações, confirmando que o penhasco de memória é estrutural e pode ser eliminado nesta escala.
• Modelagem de Linguagem (Escala 180M): Um modelo Echo-180M treinado em 10B tokens igualou ou superou o desempenho de outros modelos de 100M–400M parâmetros (incluindo GPT-2 345M e variantes Mamba-3) em cinco de seis benchmarks zero-shot (HellaSwag, PIQA, ARC, WinoGrande, LAMBADA), apesar de usar 10x menos tokens de treinamento do que algumas linhas de base. Isso sugere que a arquitetura é tanto eficiente em memória quanto eficiente em dados.

5. Significado e Alegações

O artigo alega que o Echo demonstra que SSMs não precisam trocar suas vantagens de eficiência (tempo linear, memória constante) para recuperar as capacidades de recordação da atenção. Ao redefinir a recuperação como um problema de regressão espectral de forma fechada, o Echo fecha a lacuna de desempenho com modelos baseados em atenção em tarefas de recuperação, mantendo um envelope de memória fixo.

Os autores enfatizam que o ganho de recuperação é impulsionado pelo operador espectral e pela solução de forma fechada, e não por estratégias de mascaramento de prefixo. Eles argumentam que essa abordagem aponta para uma nova classe de arquiteturas onde a recorrência lida com o contexto local e módulos leves de operadores lidam com buscas esparsas de longo alcance, tudo dentro de um orçamento de memória constante. O trabalho sugere que o "penhasco de memória" é um problema estrutural solucionável, e não uma limitação inerente de modelos de espaço de estado.