Busca de Memória Substitui Atenção Linear em Prefixos Longos

Pesquisadores do Institute of Science Tokyo e Imperial College London publicaram um método que substitui atenção linear sobre prefixos longos por busca de memória de tempo constante. No LLaMA 3.1-8B, a abordagem reduz latência de atenção 1.36x com orçamento de memória de 8K e supera RAG de atenção completa no benchmark NBA com 20% do footprint de memória, sem atualizações de gradiente no modelo.

Na inferência, atenção sobre um prefixo (prompt de sistema, documentos recuperados, exemplos in-context) escala linearmente com comprimento do prefixo durante tanto prefill quanto cada passo de decodificação. Prefix caching (abordagem do Anthropic no Claude Code) amortiza custo de prefill mas deixa overhead de passo de decodificação e memória KV-cache intactos. Compressão de atenção ainda lê o prefixo. Destilação de contexto e ajuste fino de hypernetwork requerem treinamento baseado em gradiente, que é caro e quebra quando o prefixo muda. Este método evita ambos.

O método precomputa atenção de prefixo em modo forward-only. Ele executa queries representativas através do modelo, coleta suas saídas de atenção e agrupa-as em centroides. Na inferência, uma query chegando recupera o centroide mais próximo e mescla-o com self-attention usando uma operação online-softmax—uma reconstrução sem perda que pula tokens de prefixo. Custo de busca escala logaritmicamente com número de centroide. Uma vez construído, custo por-decodificação permanece constante enquanto o prefixo cresce.

No ManyICLBench com LLaMA 3.1-8B, o método melhora acurácia sobre aprendizado in-context padrão em orçamentos de 1K a 8K centroides, reduzindo latência 1.36x em 8K. No NBA (uma tarefa RAG), supera baselines de atenção completa com 20% da memória KV. O artigo também valida em RuleArena para rule-following sobre prompts de sistema longos. Especificidades de hardware e throughput não são divulgados.

Construir a memória requer um corpus de query representativa. A qualidade dos centroides depende de quão bem queries de construção cobrem distribuição de query em tempo de execução. Para RAG de propósito geral com queries imprevisíveis, selecionar esse corpus é não-trivial. Segundo, a memória se liga a um prefixo específico. Qualquer atualização de prefixo requer rebuild—uma passagem forward-only mas um passo operacional real. O artigo não reporta custo de construção em tempo de wall-clock ou GPU-hours, que arquitetos precisam para orçamentar cadência de rebuild.

O artigo nota que influência de prefixo decai conforme geração prossegue, mesmo sob atenção completa. Mas a memória é construída de uma distribuição de query estática. Em sequências generativas longas onde estado de modelo desliza, o centroide recuperado pode divergir do que atenção completa produziria. Tarefas de avaliação executam geração de comprimento moderado. Se o ganho de latência 1.36x se mantém além de 32K passos de decodificação permanece untested.

Código está disponível em github.com/yasu0001/AttentionMemory.

Conclusão do arquiteto: para prefixos estáveis e de alta reutilização (prompts de sistema fixos, corpora de documentos atualizadas diariamente), precomputar centroides uma vez e servir buscas em tempo de decodificação é um ganho de latência plug-and-play sem retreinamento. Risco: um corpus de query mal construído em tempo de build degrada qualidade de recuperação em runtime.