EpMAN: Episodic Memory AttentioN for Generalizing to Longer Contexts¶

会议: ACL 2025
arXiv: 2502.14280
代码: -
领域: others
关键词: long context, episodic memory, attention, context length extension, retrieval-augmented generation

一句话总结¶

提出 EpMAN 方法，通过情景记忆模块估计上下文块的相对相关性，用该相关性重新加权解码器的自注意力（differentiating attention），配合噪声训练和注意力范围扩展策略，在 16k-256k 上下文长度范围内实现了比长上下文 LLM 和 RAG 更强且更鲁棒的表现。

研究背景与动机¶

LLM 在长上下文输入上的泛化仍是重大挑战
现有方案的不足：
继续预训练长序列：计算开销巨大，自注意力 \(O(n^2)\) 复杂度
位置编码外推（PI、RoPE scaling）：需额外微调，效果有限
稀疏注意力 / 滑窗注意力：可能丢失关键信息
RAG：检索模型与 LLM 参数记忆冲突，检索噪声导致幻觉或忽略上下文
长上下文的三大难题：
- 近因偏差（Recency Bias）：LLM 偏向关注上下文末尾
- 干扰信息影响：无关文档会降低模型准确性
- 注意力稀释（Attention Dilution）：softmax 归一化导致注意力被稀释
本文灵感：受 Kahneman 双系统理论启发——自注意力是快速直觉的"系统1"，EpMAN 模拟缓慢计算的"系统2"

方法详解¶

整体框架：双层注意力¶

情景记忆层（Episodic Memory）：将长上下文分块存入记忆模块，估计各块相对于 query 的相关性
自注意力层：用情景注意力 \(a_{mem}\) 对标准自注意力进行重新加权

情景记忆操作¶

Write 操作：将上下文分为固定大小的 chunk（256 tokens/chunk），使用预训练检索器（Dragon）编码后存入记忆
Read 操作：用 query 编码与 chunk 编码的余弦相似度计算情景注意力 \(a_{mem}\)
情景记忆同时存储 chunk 的 KV cache（存于 CPU 内存以应对大规模上下文）

差异化注意力（Differentiating Attention）¶

将标准注意力与情景注意力相乘：

\[a_{epman} = \text{softmax}\left(\frac{qK^T}{\sqrt{d_z}}\right)(V * a_{mem})\]

\(a_{mem}\) 是 chunk 级别的权重，broadcast 到 token 级别
效果：放大相关 chunk 的注意力，抑制无关 chunk

训练数据¶

预训练数据： - 用 Mixtral-8x22B 生成合成段落 - 加入 Wikipedia 干扰段落增加上下文长度 - 使用 next-token prediction + episodic attention loss

QA 合成数据： - Topic-sampled：教师模型生成段落和问答对 - Wikipedia-based：从 Wikipedia 段落生成问答 - Hard negative mining：挖掘与相关 chunk 主题相似但无关的 Wikipedia 段落作为强干扰

训练配置：episode size = 16（1 relevant + 15 distractors），chunk = 256 tokens，有效训练上下文 = 4K tokens

噪声训练去噪（Noisy Training）¶

问题：固定使用 \(a_{mem}\) 权重训练，模型会过拟合于"最高权重=最相关"的模式；OOD 时检索器可能将相关 chunk 排在较低位置
解决方案：top-K chunks 接收 0.9-1.0 之间的随机权重 + 随机排列顺序
提供去噪目标，使解码器即使在 retriever 不准确时也能找到相关信息

损失函数¶

\[L = \mathbb{E}_\mathcal{D} [\alpha \ln p(l|q, C) + \ln(a|q, C, a_{mem})]\]

第一项：情景注意力损失（cross-entropy，\(\alpha = 0.1\)）
第二项：next-token prediction 损失

BroadAttn：推理时注意力范围扩展¶

NarrowAttn：仅关注 top-K chunks
BroadAttn：扩展到 top-K 各 chunk 的相邻 chunk
解决信息截断问题（例如"Albert Einstein was born in Germany"在一个 chunk，"He taught himself algebra"在相邻 chunk）
保持原始上下文中的 chunk 顺序

实验¶

实验设置¶

解码器：Mistral-7B-Instruct-v0.2（LoRA 微调）
检索器：Dragon（多轮上下文编码器）
评测任务：Needle-in-a-Haystack（Paul Graham / PG19）、FactRecall-en、MultifieldQA、LoogleSD
使用 LV-Eval 框架的最难设置（同时包含 CFI 混淆信息 + KPR 关键词替换）

主实验结果¶

Needle-in-a-Haystack（16k-128k）： - EpMAN 在 Paul Graham 和 PG19 数据源上均达到近乎完美的召回率（99-100%） - 对比：Mistral-7B 在 128k 时仅 25.4%，Phi-3-128k 在 16k 时仅 26.4% - Dragon + Mistral 改善但远不如 EpMAN

FactRecall-en（16k-256k，LV-Eval CFI+KPR）：

方法	16k	32k	64k	128k	256k	均值
Mistral-7B	65.3	72.5	41.0	22.5	11.5	42.6
Phi-3-128k	82.0	80.5	81.0	63.0	34.5	68.2
Dragon+Mistral	74.2	71.8	66.0	77.2	69.0	71.7
EpMAN (noisy+BroadAttn)	81.8	75.2	76.0	75.2	80.2	77.7

EpMAN 在 256k 时保持 80.2%，而 Phi-3-128k 暴跌到 34.5%

MultifieldQA（LLM-as-Judge）： - EpMAN (noisy+BroadAttn) 均值 74.3，超越 Dragon+Mistral 的 69.7 和 Phi-3 的 42.6

LoogleSD： - EpMAN (uniform+BroadAttn) 均值 78.6，超越 Dragon+Mistral 的 77.4

消融与关键发现¶

Noisy training vs. Uniform training：噪声训练在 FactRecall 和 MultifieldQA 上显著优于均匀训练（77.7 vs. 75.1）
BroadAttn vs. NarrowAttn vs. Exact：BroadAttn 始终最优，因为解决了信息截断问题
Exact 推理表现差：因为在 CFI+KPR 设置下 Dragon 检索器不总能正确排序相关 chunk
对于 LoogleSD（Wikipedia 来源数据），检索器和训练数据都源自 Wikipedia，uniform training 即可

亮点与洞察¶

双系统理论的巧妙应用：自注意力 = 系统1（快而不准），情景记忆注意力 = 系统2（慢而准确）
噪声训练提供去噪目标：对 OOD 泛化至关重要，使模型能容忍检索器的排序错误
BroadAttn 推理策略用极低成本解决了 chunk 边界信息截断问题
在 4K token 上训练即可泛化到 256K——极高的训练效率
在带混淆信息和关键词替换的高难度设置下仍表现出色

局限性¶

需要存储完整 KV cache（目前在 CPU），对超大文档内存开销大
较大的 top-K 值增加训练内存需求
uniform/noisy training 和 exact/narrow/broad attention 的最优组合取决于任务性质
仅在 Mistral-7B 上验证，未测试更大模型
检索器（Dragon）的质量对整体性能有显著影响

评分 ⭐⭐⭐⭐¶

巧妙地将情景记忆与自注意力结合，训练高效（4K 泛化到 256K），在最难设置下表现优异。方法直觉清晰、实验充分，但 scalability（KV cache 内存）和检索器依赖是实际部署的瓶颈。