Source-linked AI summary
Context Memorization for Efficient Long Context Generation
Yasuyuki Okoshi, Hao Mark Chen, Guanxi Lu, Hongxiang Fan, Masato Motomura, Daichi Fujiki
TL;DR
Long prefixes can lose influence during generation while making attention increasingly costly, and existing compression or parameter-internalization approaches retain important trade-offs. The paper proposes attention-state memory, which stores and retrieves precomputed prefix attention states through forward-only construction. Across ICL and RAG benchmarks, it improves or matches full-attention performance while reducing memory or attention cost.
Problem
Long conditioning prefixes can experience prefix decay and impose attention costs that scale linearly with prefix length, while existing alternatives retain inference or training burdens.
Method
Attention-state memory externalizes prefixes into a lightweight lookup memory of precomputed attention outputs, constructed with forward-only inference and retrieved through clustered query representations.
Results
Across ICL and RAG benchmarks, attention-state memory matches or exceeds full-attention performance while reducing prefix attention cost.
Takeaways & Limitations
Attention-state memory provides a favorable accuracy-efficiency trade-off for reusing long prefixes across queries.
Takeaways & Limitations
Larger memory buckets can fail to scale when benchmarks lack enough training data to populate their centroid codebooks.
Abstract
from arXiv · showhide
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.
1 Introduction
Long prefixes can lose influence during generation and impose attention costs that grow with prefix length. Attention-state memory addresses both issues by externalizing prefix attention states into a forward-only lookup memory, with benchmark gains and lower latency.
- Motivation: Long conditioning prefixes suffer prefix decay and attention overhead that scales linearly with prefix length.The overhead applies during prefill and every decode step.
- Prior approaches: Existing alternatives either retain prefix attention through compression or internalize prefixes using gradient-based training.Training-based internalization is slow, memory-intensive, and ill-suited to prefix updates.
- Approach: Attention-state memory externalizes prefixes into a lightweight lookup memory built through forward-only computation.It retrieves precomputed attention states instead of re-attending to the prefix during inference.
- Approach: The method constructs attention states from representative queries, clusters them into centroids, and merges retrieved states with self-attention.The online-softmax identity makes this merge lossless.
- Results: On ManyICLBench, attention-state memory improves accuracy over in-context learning at 1K–8K memory budgets and reduces attention latency by 1.36× at 8K.The evaluation uses LLaMA 3.1-8B.
- Results: On the NBA RAG benchmark, the method surpasses full-attention RAG performance using only 20% of its memory footprint.The reported results cover both ICL and RAG settings.
2 Background
Prior prefix-reduction methods either keep compressed prefixes in attention or encode them into parameters. The paper motivates a third route: losslessly compose precomputed attention states and retrieve them from a compact dictionary.
- Related Work: Prefix-reduction methods fall into prefix internalization and prefix compression.The distinction is whether the prefix remains in inference-time attention.
- Related Work: Prefix internalization encodes prefixes into model parameters through context distillation or hypernetworks.These approaches remove prefix attention but require parameter generation or gradient-based adaptation.
- Related Work: Prefix compression keeps prefixes in attention while shortening tokens, continuous representations, or KV-cache entries.This preserves inference-time prefix attention while reducing its size.
- Motivation: The paper identifies a gap: no prior method simultaneously offers prefix-length-independent decoding latency, training-free construction, and no auxiliary models.This is stated as the authors’ knowledge claim.
- Online-Softmax Identity: The online-softmax identity decomposes attention over concatenated key-value blocks into independently computed attention states and merge weights.The merge recovers attention over the concatenated blocks exactly.
- Online-Softmax Identity: An attention state stores a block’s attention output and normalization term, which suffice to reconstruct that block’s contribution without its original keys and values.The paper denotes these components as (a_b(q), Z_b(q)).
- Implications: This sufficiency and composability motivate externalizing long fixed prefixes into a precomputed dictionary of attention states.Attention-state memory realizes this design in the subsequent method section.
3 Attention-State Memory
Attention-state memory (ASM) externalizes prefix attention into a per-layer, query-indexed dictionary of clustered attention states. It constructs entries through forward passes and clustering, then retrieves and merges them with self-attention during inference without attending to the prefix.
- Offline calibration: Forward-only composability allows long prefixes to be encoded in independent chunks and merged at inference, recovering attention over concatenated blocks.This structure supports offline calibration when peak GPU memory would otherwise scale linearly with prefix length.
- Overview: ASM stores per-layer attention states indexed by representative query vectors and shared across queries through clustering.Each dictionary entry contains a representative query and compressed attention state; inference retrieves the most similar entry.
- Online inference: At inference, each query independently retrieves the closest centroid per layer and merges its stored attention state with self-attention over non-prefix tokens.The merged output then continues through the remainder of the attention block without computing attention to the prefix.
- Lookup key: ASM assumes similar query tokens produce similar attention outputs, while RoPE handling and whitening determine the query representation used for retrieval.The paper considers pre-RoPE versus RoPE-unified representations and optionally whitens lookup keys to improve cosine-similarity retrieval.
- Memory construction: Memory construction collects prefix attention states from forward passes over prefix-trace pairs, then clusters query representations into K aggregated entries per layer.The prefix set may contain examples, instructions, or retrieved documents, and collection records states for response-trace queries.
- Memory construction: Attention-aware aggregation merges cluster members while normalizing Z so centroids behave as averages rather than unbounded merges.The normalization addresses attention-scale mismatch when combining independently encoded contexts.
- Efficiency: Hierarchical centroid indexing reduces lookup from O(K) to O(log K), decoupling retrieval cost from memory size and prefix length.The method also extends to GQA by storing one centroid per KV head because grouped query heads share KV heads.
4 Experiments
Experiments evaluate attention-state memory across ICL and RAG settings, showing competitive accuracy with reduced prefix-attention costs and efficient lookup. Performance depends on memory size, task characteristics, and calibration data.
- Experimental Setup: Experiments cover seven ManyICLBench tasks for ICL and the NBA benchmark from RuleArena for RAG on instruction-tuned LLaMA 3.1-8B.The NBA rulebook contains approximately 20K tokens, while ICL uses five many-shot and two reasoning-oriented tasks.
- In-Context Learning: ASM outperforms or matches ICL across 1K–8K memory entries on five ManyICLBench tasks.Within 4K entries, ASM also scales more efficiently because each entry stores attention output over the entire prefix.
- In-Context Learning: KVZip underperforms both ASM and ICL across most settings, while ASM preserves critical information by storing attention outputs aggregated over the prefix.The passage specifically identifies label tokens in ICL prompts as information that KV compression may collapse.
- In-Context Learning: ASM’s gains are modest on math and science reasoning tasks because ICL itself shows only marginal improvement as sequences lengthen.Across 1K–16K bucket sizes, ASM tracks the ICL baseline closely while retaining comparable prefix information.
- Retrieval-Augmented Generation: ASM surpasses full-rulebook ICL on NBA at K=4K using 20% of the memory footprint without supplying reference rules during inference.Accuracy is non-monotonic in entry count, so the optimal codebook size should be tuned as a task-specific hyperparameter.
- Efficiency Analysis: Chunked construction closely matches the 16K banking77 baseline while processing no more than 4K tokens in one forward pass.With equal total tokens consumed, chunking also exceeds the single-4K baseline, indicating complementary information across chunks.
- Efficiency Analysis: Hierarchical retrieval with top-m ≥16 matches flat top-1 accuracy across most settings of nL1, so the evaluation adopts top-m=16.This is the smallest tested value achieving the match and therefore minimizes retrieval cost.
- Efficiency Analysis: Hierarchical lookup grows much more slowly than full attention as the memory budget increases, suppressing inference-time growth with memory size.It has relatively higher overhead at small budgets, while optimized ASM combines linear and hierarchical lookup.
5 Conclusion
The paper introduces attention-state memory to decouple reusable prefix knowledge from repeated per-query attention. Across ICL and RAG, it reports a favorable accuracy–efficiency trade-off and positions the approach as a step toward compact reusable representations.
- Conclusion: Attention-state memory externalizes prefix knowledge into a training-free, lookup-based memory of precomputed attention outputs.The design targets prefix reuse without repeated full attention over the entire context.
- Conclusion: On ICL, ASM surpasses full attention across 1K–8K memory entries on average while reducing attention latency by 1.36× at 8K.This is the paper’s reported ICL accuracy–efficiency result.
- Conclusion: On RAG, ASM outperforms full attention with only 20% of the memory footprint.The conclusion presents this as the corresponding RAG result.
- Conclusion: The work frames attention-state memory as one step toward externalizing LLM knowledge into compact, reusable representations beyond text and model parameters.This broader direction is stated as the paper’s perspective rather than as an evaluated result.
A Detailed Hyperparameters
The appendix details task-dependent construction limits, latency measurement settings, and computational resources used to evaluate attention-state memory.
- Hyperparameters: Construction iterations are determined by each task’s available training split, which bounds the distinct examples observed during centroid updates.Large-pool tasks use 5,000–6,000 iterations, while limited-data tasks are capped lower.
- Latency Evaluation: Latency is measured during decoding on one NVIDIA RTX Ada 4500 GPU with batch size 1, 512-token questions, and 100-token generation.Prefix lengths or memory entries vary from 1K to 16K, with FlashAttention and a custom Triton lookup implementation.
- Compute Resources: Non-latency experiments run on one NVIDIA H100 GPU, while NBA memory construction and evaluation each take up to approximately 1.5 hours.Other tasks complete substantially faster than these NBA upper bounds.
B Generalizability of Attention-State Memory to Different Models
Attention-state memory generalizes across model scales and families in both ICL and RAG settings without architecture-specific tuning. Its advantage is strongest at smaller banking77 memory budgets, while NBA performance consistently exceeds the reported baselines.
- The evaluation isolates model-size and model-family effects using LLaMA-3.2-3B and Qwen3-8B on banking77 and NBA.Banking77 represents ICL, while NBA represents RAG; reasoning tasks are excluded because ICL provides no measurable benefit there.
- ASM matches or exceeds same-budget ICL accuracy across memory sizes on banking77 for both Qwen3-8B and LLaMA-3.2-3B-Instruct.The gap is most pronounced at smaller budgets and narrows as the memory budget grows.
- ASM consistently outperforms both full-rulebook ICL and zero-shot settings on the NBA benchmark for both evaluated backbones.This follows the trend observed with LLaMA-3.1-8B in the main experiments.
- Optimal memory size differs between the two backbones, attributed to differences in attention-head geometry and grouping factor c.
- Overall, ASM transfers across model scale and family without architecture-specific tuning.
C Limitation
Attention-state memory relies on local query structure so cached centroids can approximate attention faithfully. This assumption may fail under substantial prefix drift, requiring lookup strategies for non-stationary query distributions.
- The method assumes queries exhibit local structure, allowing a small set of cached centroids to faithfully approximate attention.
- Substantial prefix drift can make query distributions diffuse in long multi-turn conversations spanning varied exchanges.
- Under non-stationary query distributions, query-dependent retrieval may need replacement by alternative lookup strategies.The paper identifies extending externalization to diverse workloads as future work.