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

📄 arXiv: 2607.02303 📥 PDF

作者: Wanyun Cui

分类: cs.AI

发布日期: 2026-07-05

💡 一句话要点

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

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

关键词: 线性注意力 海马体机制 记忆检索 自然语言处理 深度学习

📋 核心要点

  1. 现有的线性注意力模型在处理多个键值对时，容易导致早期信息的丢失，影响记忆的准确性。
  2. HOLA模型通过引入海马体机制，结合压缩记忆和精确KV缓存，提升了信息的存储和检索能力。
  3. 实验结果表明，HOLA在多个基准测试中显著降低了困惑度，并在长序列检索任务中表现出更高的鲁棒性。

📝 摘要（中文）

线性注意力和状态空间语言模型将前缀压缩为固定大小的递归状态，导致O(1)内存，但在多个键值关联竞争时，早期信息被覆盖，回忆能力下降。受互补学习系统的启发，本文提出HOLA（海马线性注意力），在保持常规的压缩记忆的同时，增加了一个有界的精确KV缓存，形成半参数化的测试时记忆。该缓存在没有学习驱逐模块的情况下写入，保留与状态相关的预测残差，从而实现精确的KV对检索。实验结果显示，HOLA在340M参数下训练于15B SlimPajama数据集，Wikitext困惑度从27.32降低至22.92，LAMBADA困惑度从30.95降低至30.26，且在32k令牌的检索任务中表现出色。

🔬 方法详解

问题定义：本文旨在解决线性注意力模型在处理大量键值对时，早期信息被覆盖导致的记忆丢失问题。现有方法在保持O(1)内存的同时，无法有效保留重要信息。

核心思路：HOLA模型通过引入海马体机制，结合常规的压缩记忆和一个有界的KV缓存，形成半参数化的记忆结构，以便在不损失信息的情况下进行有效的检索。

技术框架：HOLA的整体架构包括两个主要模块：压缩记忆模块和KV缓存模块。压缩记忆模块负责状态的线性压缩，而KV缓存模块则存储不应被压缩的关联信息。

关键创新：HOLA的主要创新在于引入了KV缓存机制，使得模型能够在不依赖学习驱逐模块的情况下，保留重要的预测残差，从而实现更精确的记忆检索。

关键设计：模型的参数设置为340M，训练数据为15B SlimPajama令牌。KV缓存的写入机制基于预测残差的大小，采用解耦的RMSNorm-gamma进行缓存读取，以实现精确的KV对检索。

🖼️ 关键图片

img_0
img_1

📊 实验亮点

HOLA模型在Wikitext数据集上的困惑度从27.32降低至22.92，提升幅度达到16.1%。在LAMBADA数据集上，困惑度从30.95降低至30.26。此外，HOLA在32k令牌的检索任务中表现出色，显示出比GDN和匹配的HOLA+近期缓存更高的鲁棒性。

🎯 应用场景

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).