用外部记忆存储长前缀注意力状态,提升生成效率与准确性
Context Memorization for Efficient Long Context Generation

- 将前缀注意力状态预计算存入轻量记忆库,免于重复计算
- 在8K上下文时比全注意力RAG快1.36倍,内存仅需20%
- 无需训练即可适应新前缀,适合动态场景应用
现代大语言模型应用越来越依赖长条件前缀来控制推理行为。尽管前缀增强推理有效,但存在两个结构性限制:一是前缀影响随生成过程减弱,二是前缀注意力计算随长度线性增长。现有方法要么在推理时持续关注前缀并压缩,要么通过梯度训练将前缀内化到模型参数中。前者仍需在推理时访问前缀,后者训练成本高且不支持前缀更新。为此,我们提出注意力状态记忆(attention-state memory),一种无需训练的方法,将前缀与查询词之间的注意力状态外部化为轻量级、基于查找的内存。在ManyICLBench上使用LLaMA-3.1-8B,我们的方法在1K–8K内存预算下优于上下文学习,并在8K时将注意力延迟降低1.36倍;在NBA基准上,仅用20%的内存就超越了全注意力RAG性能。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。