Context Memorization for Efficient Long Context Generation

Resumo Técnico: Memória de Estado de Atenção para Geração Eficiente de Contexto Longo

Declaração do Problema
Aplicações modernas de Modelos de Linguagem de Grande Escala (LLM) dependem cada vez mais de prefixos de condicionamento longos (por exemplo, exemplos em contexto, documentos recuperados ou instruções de agentes) para guiar a inferência. No entanto, a inferência aumentada por prefixos enfrenta duas limitações estruturais primárias:

1. Decaimento do Prefixo: À medida que a geração prossegue, a atenção do modelo distribui-se sobre os tokens recém-gerados, fazendo com que a influência do prefixo diminua, particularmente em cenários de contexto longo.
2. Ineficiência de Inferência: O cálculo de atenção sobre o prefixo escala linearmente com seu comprimento ($O(L)$), impondo latência significativa e sobrecarga de memória durante tanto o pré-preenchimento (prefill) quanto cada etapa de decodificação. Embora o cache de prefixo amortize os custos de pré-preenchimento, ele ainda incorre em consumo substancial de memória.

Soluções existentes caem em duas categorias, ambas com desvantagens:

• Internalização do Prefixo: Métodos que ajustam os parâmetros do modelo ou utilizam hiper-redes para internalizar o comportamento do prefixo. Estes requerem treinamento baseado em gradiente, tornando-os intensivos em recursos, lentos para atualização e inadequados para mudanças dinâmicas no prefixo.
• Compressão do Prefixo: Métodos que mantêm o prefixo no mecanismo de atenção, mas o comprimem (por exemplo, poda de tokens ou evicção de cache KV). Estes ainda exigem atenção ao prefixo comprimido durante a inferência, o que significa que a latência escala com o comprimento comprimido e a influência do prefixo permanece sujeita ao decaimento.

Metodologia: Memória de Estado de Atenção (ASM)
Os autores propõem a Memória de Estado de Atenção, uma abordagem sem treinamento que externaliza o prefixo em uma memória leve baseada em busca de estados de atenção pré-calculados. Este método aproveita a identidade do online-softmax para decompor a atenção sobre um bloco de chaves concatenado em atenção sobre sub-blocos.

A ideia central é que, para uma dada consulta, o estado de atenção $(a, Z)$ — composto pela saída de atenção e pelo termo de normalização — é suficiente para reconstruir a contribuição do prefixo para a atenção, sem re-atenção às chaves e valores originais.

O framework opera em duas fases:

1. Fase de Calibração Offline (Construção da Memória):

   • Coleta: O modelo realiza passagens forward sobre um conjunto de traços representativos de consulta-resposta concatenados com o prefixo. Para cada token de consulta na resposta, o sistema registra o vetor de consulta e o estado de atenção correspondente $(a, Z)$ sobre o prefixo.
   • Agrupamento: Os vetores de consulta coletados são agrupados (por exemplo, via K-means) em $K$ centróides.
   • Agregação: Para cada cluster, uma única entrada representativa é criada. O vetor de consulta do centróide é a média das consultas do cluster. O estado de atenção $(\bar{a}, \bar{Z})$ é calculado via agregação consciente de atenção, mesclando os estados de todos os membros do cluster usando a regra de atualização do online-softmax. Isso garante que o centróide represente uma média sem perdas da atenção do prefixo do cluster.
   • Eficiência: Para lidar com prefixos muito longos (por exemplo, 16K+ tokens) sem exceder a memória da GPU, o prefixo pode ser particionado em blocos, codificado independentemente e mesclado dentro da memória usando a propriedade de composabilidade da identidade do online-softmax.

2. Fase de Inferência Online:

   • Recuperação: Para uma consulta entrante, o sistema identifica o centróide de memória mais próximo usando similaridade de cosseno (opcionalmente com transformações de branqueamento e tratamento de RoPE).
   • Mesclagem: O estado de atenção pré-calculado recuperado $(\bar{a}, \bar{Z})$ é mesclado com a saída de auto-atenção da consulta (calculada sobre tokens não-prefixo) usando o operador de mesclagem definido pela identidade do online-softmax. Isso recupera a saída de atenção exata que teria sido produzida se o prefixo estivesse presente, sem atender aos tokens do prefixo.
   • Complexidade: O custo de busca escala logaritmicamente ($O(\log K)$) com o tamanho da memória via indexação hierárquica, independente do comprimento original do prefixo.

Principais Contribuições

1. Externalização sem Treinamento: O artigo propõe um método para externalizar prefixos longos em um dicionário de memória compacto sem exigir treinamento baseado em gradiente ou modelos auxiliares.
2. Extensão do Online-Softmax: Os autores estendem a identidade do online-softmax do cálculo eficiente de atenção para o domínio da reutilização de prefixos entre consultas, permitindo a reconstrução sem perdas da atenção do prefixo via recuperação.
3. Desempenho e Eficiência: Experimentos demonstram que a ASM iguala ou supera o desempenho da atenção completa enquanto reduz significativamente os custos computacionais.

Resultados Experimentais
O método foi avaliado no ManyICLBench (Aprendizado em Contexto) e RuleArena/NBA (Geração Aumentada por Recuperação) usando LLaMA-3.1-8B.

• Aprendizado em Contexto (ICL): No ManyICLBench, a ASM melhorou a precisão sobre o aprendizado em contexto padrão em orçamentos de memória de 1K a 8K entradas. Com um orçamento de memória de 8K, reduziu a latência de atenção em 1,36× em comparação com a atenção completa. A ASM superou as bases de compressão de cache KV (KVzip), que frequentemente colapsavam informações críticas como tokens de rótulo.
• Geração Aumentada por Recuperação (RAG): No benchmark NBA (envolvendo ~20K tokens de regras), a ASM superou o desempenho do RAG com atenção completa usando apenas 20% da pegada de memória. O desempenho ótimo foi observado em um tamanho intermediário de código (4K entradas), sugerindo que o número de entradas é um hiperparâmetro específico da tarefa, em vez de algo que deva simplesmente ser maximizado.
• Latência: Embora a busca hierárquica incorra em maior sobrecarga em orçamentos de memória muito pequenos, o crescimento da latência é sublinear. A ASM torna-se mais rápida que a atenção completa em aproximadamente 4K entradas de memória e alcança um acréscimo de velocidade de 1,8× em 16K entradas.

Significado e Alegações
O artigo afirma que a Memória de Estado de Atenção aborda os custos estruturais da geração de contexto longo ao desacoplar o conhecimento do prefixo do cálculo de atenção por consulta. Ao externalizar o prefixo em uma memória pré-calculada baseada em busca, o método elimina a escala linear da latência de inferência com o comprimento do prefixo e remove o problema de decaimento do prefixo associado aos mecanismos de atenção padrão.

Os autores posicionam este trabalho como um passo em direção a uma direção de pesquisa mais ampla: externalizar o conhecimento de LLMs em representações compactas e reutilizáveis que existem além de prompts de texto e parâmetros do modelo. Eles enfatizam que sua abordagem oferece uma troca favorável entre precisão e eficiência, fornecendo uma alternativa sem treinamento à internalização de parâmetros e uma alternativa mais eficiente à compressão de prefixo.