Context Memorization for Efficient Long Context Generation

📄 arXiv: 2605.18226v1 📥 PDF

作者: Yasuyuki Okoshi, Hao Mark Chen, Guanxi Lu, Hongxiang Fan, Masato Motomura, Daichi Fujiki

分类: cs.CL, cs.AI

发布日期: 2026-05-18

💡 一句话要点

提出Attention-State Memory，解决长文本生成中前缀信息衰减和注意力计算效率问题。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长文本生成 注意力机制 上下文学习 记忆网络 大型语言模型

📋 核心要点

  1. 现有长文本生成方法存在前缀信息利用率低和注意力计算效率低下的问题。
  2. 论文提出Attention-State Memory，通过存储预计算的注意力状态来高效利用前缀信息。
  3. 实验表明，该方法在提高准确性的同时，显著降低了注意力计算延迟，并减少了内存占用。

📝 摘要（中文）

现代大型语言模型（LLM）应用越来越依赖于长上下文前缀来控制推理时的模型行为。虽然前缀增强推理有效，但它存在两个结构性限制：i) 前缀的影响随着生成过程的进行而减弱；ii) 对前缀的注意力计算随其长度线性扩展。现有方法要么在压缩前缀的同时保持对其的注意力，要么通过基于梯度的训练将其内化到模型参数中。前者仍然需要在推理时关注前缀，而后者是训练密集型的，并且不适合前缀更新。为了解决这些问题，我们提出了一种注意力状态记忆（attention-state memory）方法，这是一种免训练的方法，它将前缀外化为一个轻量级的、基于查找的记忆，其中存储了前缀和查询token之间预先计算的注意力状态。在使用LLaMA-3.1-8B的ManyICLBench上，我们的方法在1K-8K的记忆预算下提高了上下文学习的准确性，并在8K时将注意力延迟降低了1.36倍，并且在使用仅20%的内存占用情况下，在NBA基准测试中超过了全注意力RAG的性能。

🔬 方法详解

问题定义：长文本生成任务中，如何有效利用长上下文前缀信息是一个关键问题。现有方法，如直接注意力机制，计算复杂度高，且前缀信息容易在生成过程中衰减。基于梯度训练的方法虽然能将前缀信息内化到模型参数中，但训练成本高昂，且难以适应前缀的动态更新。

核心思路：论文的核心思路是将前缀信息转化为预计算的注意力状态，并存储在一个外部记忆模块中。在生成过程中，模型可以通过查询该记忆模块来高效地获取前缀信息，避免了对整个前缀进行重复的注意力计算。这种方法无需训练，可以灵活地适应前缀的更新。

技术框架：整体框架包含三个主要部分：1) 前缀编码：将输入的前缀文本编码成一系列的token表示。2) 注意力状态计算与存储：计算前缀token和查询token之间的注意力状态，并将这些状态存储在外部记忆模块中。3) 生成阶段：在生成阶段，模型通过查询外部记忆模块来获取前缀信息，并将其融入到当前的生成过程中。

关键创新：最重要的创新点在于将注意力计算从在线计算转化为离线预计算和存储。通过预先计算并存储前缀token和查询token之间的注意力状态，模型可以在生成阶段快速检索相关信息，从而显著降低计算复杂度，提高生成效率。与现有方法的本质区别在于，该方法无需训练，且能有效解决长文本生成中前缀信息衰减的问题。

关键设计：关键设计包括：1) 注意力状态的表示：注意力状态可以表示为注意力权重或注意力头的输出。2) 记忆模块的组织方式：记忆模块可以采用键值对的形式，其中键是查询token的表示，值是对应的注意力状态。3) 查询机制：可以使用近似最近邻搜索等技术来高效地查询记忆模块。4) 超参数：记忆模块的大小、查询算法的参数等。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，提出的Attention-State Memory方法在ManyICLBench数据集上，使用LLaMA-3.1-8B模型，在1K-8K的记忆预算下，提高了上下文学习的准确性，并在8K时将注意力延迟降低了1.36倍。在NBA基准测试中，该方法仅使用20%的内存占用，就超过了全注意力RAG的性能。这些结果表明，该方法在提高效率和降低内存占用方面具有显著优势。

🎯 应用场景

该研究成果可广泛应用于需要长上下文信息的自然语言生成任务中，例如：文档摘要、机器翻译、对话生成、代码生成等。通过高效利用长上下文信息，可以显著提高生成质量和效率，并降低计算成本。该方法尤其适用于需要频繁更新上下文信息的场景。

📄 摘要（原文）

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.