Memorização de Contexto para Geração Eficiente de Contextos Longos

Resumo

Aplicações modernas de modelos de linguagem de grande escala (LLMs) dependem cada vez mais de prefixos de condicionamento longos para controlar o comportamento do modelo durante a inferência. Embora a inferência aumentada por prefixo seja eficaz, ela acarreta duas limitações estruturais: i) a influência do prefixo diminui à medida que a geração prossegue, e ii) a computação de atenção sobre o prefixo escala linearmente com seu comprimento. As abordagens existentes ou mantêm o prefixo na atenção enquanto o comprimem, ou o internalizam nos parâmetros do modelo por meio de treinamento baseado em gradiente. A primeira ainda atende ao prefixo durante a inferência, enquanto a segunda é intensiva em treinamento e inadequada para atualizações do prefixo. Para abordar essas questões, propomos a memória de estados de atenção, uma abordagem livre de treinamento que externaliza o prefixo em uma memória leve baseada em consulta de estados de atenção pré-computados entre tokens de prefixo e de consulta. No ManyICLBench com LLaMA-3.1-8B, nosso método melhora a precisão em relação ao aprendizado em contexto com orçamentos de memória de 1K a 8K, enquanto reduz a latência da atenção em 1,36x em 8K, e supera o desempenho do RAG com atenção completa no benchmark NBA usando apenas 20% do seu consumo de memória.

English

Modern large language model (LLM) applications increasingly rely on long conditioning prefixes to control model behavior at inference time. While prefix-augmented inference is effective, it incurs two structural limitations: i) the prefix's influence fades as generation proceeds, and ii) attention computation over the prefix scales linearly with its length. Existing approaches either keep the prefix in attention while compressing it, or internalize it into model parameters through gradient-based training. The former still attends to the prefix at inference, while the latter is training-intensive and ill-suited to prefix updates. To address these issues, we propose attention-state memory, a training-free approach that externalizes the prefix into a lightweight, lookup-based memory of precomputed attention states between prefix and query tokens. On ManyICLBench with LLaMA-3.1-8B, our method improves accuracy over in-context learning at 1K-8K memory budgets while reducing attention latency by 1.36x at 8K, and surpasses full-attention RAG performance on NBA benchmark using only 20% of its memory footprint.