Context Memorization for Efficient Long Context Generation

2026-05-18

AI Summary

This work addresses the challenges of influence decay in long-context prefixes during generation and the linearly growing computational cost of attention with sequence length. The authors propose a training-free attention state memory mechanism that externalizes the prefix into a lightweight, dynamically updatable store of precomputed attention states. By leveraging query-prefix attention precomputation and memory-efficient lookup, the method enables efficient reuse of prefix information during inference. This approach is the first to compress long-prefix information into a flexible external memory without requiring any training. On the ManyICLBench benchmark, it surpasses in-context learning accuracy under memory budgets of 1K–8K tokens—achieving a 1.36× reduction in inference latency at 8K—and outperforms full-attention RAG on the NBA benchmark using only 20% of the memory.

Abstract

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.

Problem

• long context generation
• prefix conditioning
• attention efficiency
• context memorization
• inference latency

Innovation

• attention-state memory
• long context generation
• training-free
• prefix externalization
• efficient inference