A Hippocampus for Linear Attention: An Exact Memory for What the Recurrent State Forgets

📄 arXiv: 2607.02303v1 📥 PDF

作者: Wanyun Cui

分类: cs.AI

发布日期: 2026-07-02

备注: 12 pages


💡 一句话要点

提出HOLA以解决线性注意力模型的记忆遗忘问题

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)

关键词: 线性注意力 海马模型 记忆机制 自然语言处理 信息检索 长序列处理 机器学习

📋 核心要点

  1. 现有的线性注意力模型在处理大量键值对时,早期信息容易被覆盖,导致记忆丢失和召回能力下降。
  2. HOLA通过引入一个有界的KV缓存,结合传统的递归状态,形成了一种新的半参数化记忆结构,以保持重要信息。
  3. 实验结果表明,HOLA在Wikitext和LAMBADA数据集上均表现出显著的性能提升,并在长序列召回任务中展现出更好的鲁棒性。

📝 摘要(中文)

线性注意力和状态空间语言模型将前缀压缩为固定大小的递归状态,导致O(1)内存但损失精确记忆。HOLA(海马线性注意力)在保持压缩记忆的同时,增加了一个有界的精确KV缓存,形成半参数化的测试时记忆。该方法在340M参数下,训练于15B SlimPajama tokens,显著降低Wikitext困惑度并提升LAMBADA性能,同时在32k tokens的召回任务中表现出更强的鲁棒性。

🔬 方法详解

问题定义:论文旨在解决线性注意力模型在处理长序列时的记忆遗忘问题,现有方法在面对多个键值对时,早期信息容易被覆盖,导致召回能力下降。

核心思路:HOLA通过结合传统的递归状态和一个有界的KV缓存,形成了一种新的记忆结构,能够在保持压缩记忆的同时,存储重要的键值对,从而改善信息的召回能力。

技术框架:HOLA的整体架构包括两个主要模块:一个是压缩的递归状态,用于建模线性可压缩结构;另一个是KV缓存,用于存储不应被压缩的关键信息。缓存的写入机制不依赖于学习的驱逐模块,而是基于预测残差进行选择。

关键创新:HOLA的主要创新在于引入了一个有界的KV缓存,使得模型能够在不牺牲记忆精度的情况下,处理更长的输入序列。这一设计与现有方法的本质区别在于,HOLA能够有效地存储和检索重要信息,而不是简单地依赖于压缩的状态表示。

关键设计:HOLA的缓存写入机制基于大于某个阈值的预测残差,确保重要信息被保留。缓存读取采用解耦的RMSNorm-gamma方法,将KV对转化为精确的检索,而非软平均。

🖼️ 关键图片

img_0
img_1

📊 实验亮点

HOLA在340M参数下,训练于15B SlimPajama tokens,成功将Wikitext困惑度从27.32降低至22.92,提升幅度达到16.1%。此外,LAMBADA困惑度也从30.95降至30.26,展现出优于全注意力Transformer++的性能,并在长序列召回任务中表现出更强的鲁棒性。

🎯 应用场景

HOLA的研究成果在自然语言处理、对话系统和信息检索等领域具有广泛的应用潜力。通过改善长序列的记忆和召回能力,该模型能够提升机器翻译、文本生成和问答系统的性能,具有重要的实际价值和未来影响。

📄 摘要(原文)

Linear-attention and state-space language models compress the prefix into a fixed-size recurrent state, yielding O(1) memory at the cost of a lossy exact memory: when many key--value associations compete, earlier facts are overwritten and needle recall degrades. Inspired by Complementary Learning Systems, we give linear attention a hippocampal complement. HOLA (Hippocampal Linear Attention) keeps the usual delta-rule state as a compressive memory and adds a bounded exact KV cache, forming a semiparametric test-time memory: the state models linearly compressible structure, while the cache stores associations that should not be forced through that state. The cache writes without a learned eviction module, keeping tokens with large beta * ||e||, the prediction residual actually committed to the state; a decoupled RMSNorm-gamma cache read then turns these exact KV pairs into sharp retrieval rather than soft averaging. At 340M parameters trained on 15B SlimPajama tokens, HOLA lowers Wikitext perplexity from 27.32 to 22.92 (-16.1%), below a full-attention Transformer++ (26.88), and improves LAMBADA perplexity from 30.95 to 30.26. It also achieves the best linear in-context retrieval and remains much more robust than GDN or a matched HOLA+recency cache on RULER needle-in-a-haystack recall out to 32k tokens (16x its training length).