@omarsar0:新论文值得一读。(收藏)基本思路是将压缩循环状态与一个小型精确记忆配…
摘要
HOLA(海马体线性注意力)受海马记忆启发,通过一个有界的精确KV缓存增强线性注意力,在不牺牲效率的情况下改善长程召回和困惑度。在340M参数规模下,它在Wikitext上优于全注意力Transformer,并在高达32k token的范围内实现稳健的针召回。
查看缓存全文
缓存时间: 2026/07/04 22:56
NEW paper worth reading.
(bookmark it)
基本思路是将一个压缩性循环状态与一个小型精确记忆配对,这有助于恢复长程回忆,同时不牺牲线性注意力的效率。
更多内容如下:
线性注意力和状态空间模型将整个前缀压缩成一个固定大小的状态。这带来了 O(1) 内存的代价,但当大量键值关联竞争时,早期事实会被覆盖,针检索性能下降。
HOLA 为线性注意力提供了一个海马体补充。它保留了通常的 delta 规则状态作为压缩记忆,并添加了一个有界精确 KV 缓存,形成了一种半参数的测试时记忆。
状态模型处理线性可压缩的结构,而缓存存储那些不应强制通过状态压缩的关联。缓存写入无需学习到的驱逐模块,仅保留那些预测残差实际被提交到状态的 token。
在 15B SlimPajama token 上训练的 340M 参数模型中,HOLA 将 Wikitext 困惑度从 27.32 降至 22.92,低于全注意力 Transformer++ 的 26.88,并且在 RULER 针检索测试中,在 32k token 的上下文长度下(训练长度的 16 倍)仍然保持鲁棒。
论文:https://arxiv.org/abs/2607.02303
在我们的学院中学习构建有效的 AI 智能体:https://academy.dair.ai
线性注意力的海马体:循环状态遗忘内容的精确记忆
来源:https://arxiv.org/html/2607.02303
摘要
线性注意力和状态空间语言模型将前缀压缩成一个固定大小的循环状态,以 O(1) 内存为代价,但牺牲了精确记忆:当大量键值关联竞争时,早期事实会被覆盖,针检索性能下降。受互补学习系统启发,我们为线性注意力添加了一个海马体补充。HOLA(海马体线性注意力)保留了通常的 delta 规则状态作为压缩记忆,并添加了一个有界精确 KV 缓存,形成了一种半参数的测试时记忆:状态建模线性可压缩的结构,而缓存存储那些不应强制通过该状态的关联。缓存写入无需学习到的驱逐模块,仅保留那些具有较大 β·∥e∥ 的 token,即实际被提交到状态的预测残差;然后,一个解耦的 RMSNorm-γ 缓存读取将这些精确 KV 对转化为锐利检索而非软平均。在 15B SlimPajama token 上训练 340M 参数模型后,HOLA 将 Wikitext 困惑度从 27.32 降至 22.92(-16.1%),低于全注意力 Transformer++(26.88),并将 LAMBADA 困惑度从 30.95 降至 30.26。它还实现了最佳的线性上下文内检索,并且在 RULER 针检索测试中,在上下文长度达到 32k token(训练长度的 16 倍)时,比 GDN 或匹配的 HOLA+recency 缓存更鲁棒。
DNGLAGSAGDNT++KDAHOLA20202222242426262828303029.029.028.828.828.228.227.327.326.226.226.926.922.922.9ppl(a) Wikitext 困惑度↓线性基线全注意力 (T++)HOLA (我们的)2k4k8k16k32k00.20.20.40.40.60.60.80.811上下文长度 (训练于 2k)召回率(b) S-NIAH-1 召回率 vs. 上下文Transformer++GDN (无缓存)HOLA+recencyHOLA
图 1:HOLA 降低了困惑度并提高了长度鲁棒的针检索。(a) 在 340M 参数下,HOLA 将 Wikitext 困惑度从 27.3 降至 22.9(-16.1%),低于全注意力 Transformer++(26.9)。(b) 在 RULER S-NIAH-1 上,随着上下文增长到 32k token,HOLA 仍然比 GDN 和 HOLA+recency 更强。
1 引言
人脑并不依赖单一记忆。以海马体为中心的系统可以一次性记录特定、新颖的事件,并在之后精确回忆(Tse 等,2007),而以新皮层为中心的系统则缓慢地从许多经验中提炼出可泛化的结构(Frankland & Bontempi,2005);互补学习系统(CLS)理论(McClelland 等,1995;Kumaran 等,2016)认为,两者必须分开——一个为缓慢、压缩性泛化而构建的系统如果被强制快速写入单一事实,会遭受灾难性干扰,因此不能同时服务于快速、精确、一次性的记忆。高效的语言模型恰好缺少这种互补性。
线性注意力正是这样一个“新皮层”:出色的压缩器,但精确记忆不佳。DeltaNet、Gated DeltaNet (GDN) 以及相关的线性注意力/状态空间模型(Katharopoulos 等,2020;Gu & Dao,2023;Yang 等,2024a;c)用固定大小的循环状态 S 替换了软注意力不断增长的键值缓存,以 O(1) 内存汇总所有历史,并以次二次时间处理序列。状态在线更新——delta 规则写入每个新值与状态当前预测之间的残差——使其成为高效的运行摘要。但固定大小的状态是有损记忆:它在新的写入覆盖共享键方向的旧内容之前,只能容纳有限数量的不同键→值关联。这恰恰出现在最需要的地方:多项目关联回忆(Arora 等,2023)、密码/干草堆中的针检索(Mohtashami & Jaggi,2023),以及逐字复制特定的远距离 token(Jelassi 等,2024)。在这些任务中,全软注意力——以 O(T) 内存和 O(T²) 计算精确保留每个 token——仍然是黄金标准。问题变成:我们能保持线性注意力的成本,同时恢复其丢失的精确回忆吗?CLS 指明了答案——给这个新皮层一个海马体:一个独立的、预算有限的、精确的情景记忆,优先存储惊喜。我们将这种互补形式形式化为半参数测试时记忆回归:循环状态是线性可压缩结构的参数估计器,而有界精确缓存是对状态不应强制吸收的 KV 对的非参数修正。这个框架决定了两个设计选择:存储什么(令人惊讶的项)和如何读取(精确检索,而非软平均)。
存储什么:模型已经计算出的惊喜信号。 预算有限的情景记忆只有在保留正确的 token——那些令人惊讶的、压缩状态无法吸收的 token——时才有用。最明显的选择是保留最近的 token,即滑动窗口,大多数线性注意力混合模型都这样做(Xiao 等,2024;Zhang 等,2023;Du 等,2025;Wang 等,2025;Fang 等,2025);但近因无法表示那些遥远却仍需精确回忆的信息——一旦它滑出窗口,就像被状态覆盖一样丢失。我们的关键观察是,delta 规则模型已经测量了每个 token 对其状态的惊讶程度:如果状态对 token 值预测较差,且模型强烈提交修正,该 token 恰好是非参数缓存应保留的那种支持点。我们使用这个内在写入幅度作为缓存的驱逐分数,保留那些改变状态最多的 token,而非仅仅最近的 token。对照实验确认了预期的分离:在 340M 参数下,重要性驱逐在困惑度和长上下文检索上优于近因缓存,而常识性任务保持在单种子噪声范围内(第 4 节)。
如何读取:检索,而非平均。 一个精确副本如果像另一个线性摘要一样被读取,则毫无用处。直接重用 DeltaNet/GDN 使用的 L2 归一化查询和键会导致缓存 logits 过小,因此缓存 softmax 几乎均匀,精确记忆退化为有损平均。一个简单的解决方法是在缓存路径上仅使用 Qwen3 风格的 RMSNorm-γ(Qwen Team,2025),与状态更新解耦。这恢复了锐利、接近 argmax 的检索,同时保留了 delta 规则稳定状态更新所需的单位范数键。
综合这些想法,我们得到了海马体线性注意力(HOLA),它恢复了状态丢失的精确回忆(图 1)。在 15B SlimPajama token 上训练的 340M 参数模型中,该模型达到了所有比较模型中最低的困惑度——低于全注意力 Transformer++(Wikitext 22.9 vs. 26.9)——并且,与 GDN 不同,在硬 RULER 召回测试中在 32k token(训练 2k 长度的 16 倍)下保持长度鲁棒,而循环状态饱和且 GDN 衰减。
贡献
- 受 CLS 启发的线性注意力海马体精确记忆。 我们认为固定循环状态是一个新皮层式的压缩器,但精确记忆不佳,并添加了一个有界、海马体式的 KV 存储,用于精确的一次性关联。
- 半参数测试时记忆框架。 我们将记忆读出形式化为对前缀键-值观测的测试时回归:纯 GDN 是参数估计器 q⊤S_t,全注意力是无界非参数核估计器,而 HOLA 是有界半参数情况。
- 具体的写入/读取算法。 缓存通过 delta 规则已计算的固有惊喜信号——提交给状态的残差更新——进行写入,并通过解耦的 RMSNorm-γ 缓存路径进行读取,使精确 KV 对可检索而非软平均。
- 困惑度和检索方面的大幅实证提升。 在 340M 参数下,HOLA 相对于相同骨干网络的 GDN 锚点(27.32 → 22.92)显著降低了 Wikitext 困惑度,低于全注意力 Transformer++(26.88),实现了最佳的线性上下文内检索,并且在长上下文针检索上比 GDN 或近因缓存更鲁棒。
2 相关工作
我们的工作连接了三条线索:线性时间循环注意力、高效语言模型的混合精确记忆机制,以及半参数记忆。
线性时间循环注意力。
线性注意力(Katharopoulos 等,2020)、状态空间模型(Gu & Dao,2023;Dao & Gu,2024)和 DeltaNet 系列(Schlag 等,2021;Yang 等,2024c)用固定大小的循环状态替换了软注意力的增长 KV 缓存。Gated DeltaNet (GDN)(Yang 等,2024a)通过数据依赖的衰减增强了 DeltaNet,而 GLA(Yang 等,2024b)、GSA(Zhang 等,2024)和 Kimi Delta Attention (KDA)(Kimi Team 等,2025)是典型的现代线性注意力基线。这种固定状态高效但有损:作为关联记忆,它在新的写入干扰旧内容之前具有有限容量(Hopfield,1982;Ramsauer 等,2021)。Based 和 Zoology(Arora 等,2024;2023)明确了这个回忆-吞吐量权衡,表明高效循环模型在多项目关联回忆上落后于软注意力。HOLA 保留了循环骨干网络,但为状态不应强制压缩的 KV 对添加了有界精确记忆。
高效语言模型的混合精确记忆。
许多高效 LM 通过将循环/SSM 骨干网络与软注意力配对来恢复一些精确回忆。层间混合模型如 Jamba(Lieber 等,2024)、Samba(Ren 等,2024)和 Griffin(De 等,2024)混合了循环层与局部注意力层;层内混合模型如 NHA(Du 等,2025)、RAttention(Wang 等,2025)和 AHN(Fang 等,2025)将循环记忆与最近 token 窗口或学习到的压缩存储结合。这些设计在架构上不同,但其无损组件主要基于近因:一旦旧 token 离开窗口,必须由压缩状态表示。HOLA 则使精确记忆对前缀进行选择,保留令人惊讶的 KV 对,即使它们远离当前查询。
最接近的工作是 LTE(He & Garner,2025),它也用有界可驱逐 KV 缓存增强了 GDN。主要区别在于缓存的选择和集成方式:LTE 使用额外的 CNN 模块学习驱逐分数,并交替使用 GDN 与稀疏注意力层,而 HOLA 使用 delta 规则自身的写入幅度作为无参数惊喜分数,并将缓存附加在每个循环层内部。因此,LTE 表明可驱逐缓存可以帮助 delta 规则模型,而 HOLA 则认为循环更新本身已经暴露了精确记忆应保留内容的正确信号。
半参数和检索增强记忆。
半参数语言模型将参数化序列建模与显式非参数记忆配对,如 kNN-LM(Khandelwal 等,2020)、RETRO(Borgeaud 等,2022)、Memorizing Transformers(Wu 等,2022)和 SPALM(Yogatama 等,2021)。这些系统从外部或长期数据存储中检索。我们的设置是序列内且测试时:参数部分是循环状态 S_t,非参数部分是有界精确上下文内 KV 对集,读出在两者之间插值。这种半参数观点也将软注意力和 GDN 置于同一谱系上:全注意力是所有前缀 token 上的无界非参数估计器,而 GDN 是纯粹的参数状态估计器。
3 方法
我们从 delta 规则更新方程推导 HOLA:一个方程同时告诉我们为什么缓存是必要的(第 3.2 节)以及在缓存中存储什么(第 3.3 节);然后我们解决如何读取(第 3.4 节)和完整层(第 3.5 节)。图 2 给出了概述。
上下文 token x₁,…,xₜ 状态记忆 S_t:所有历史,压缩(有损)≈ 新皮层 精确 KV 记忆 A_t:top-w 令人惊讶的 KV 对 ≈ 海马体 + o_t:所有 token,由 β·∥e∥ 选择 q_t⊤S_t
相似文章
利用指数衰减记忆增强注意力提升查询感知的KV稀疏性
本文探讨了RAT+中的指数衰减记忆模块如何提升长上下文语言模型的查询感知稀疏推理方法,在针尖干草垛任务中,跨多种稀疏预算展示了一致的准确率提升。
@rohanpaul_ai: 来自上海大学论文的一个宏大想法——赋予AI模型人类式记忆。想象一下阅读一本巨大的悬疑小说……
本文介绍了HOLA,一种为快速AI模型(如线性注意力和状态空间模型)提供额外记忆缓存以存储令人惊讶的事实的方法,在不牺牲速度的情况下提高了它们在长上下文任务中的回忆能力。
@shikhargupta02: 我一直在学习潜在注意力(来自 DeepSeek)。它不是为每个 token 存储完整的 K 和 V 向量,而是……
作者分享了他用 DeepSeek 的潜在注意力训练一个小模型的见解,观察到潜在空间的使用因层而异,以及一个可以将 KV 缓存减少 4 倍且损失不变的测试时技巧。
@jiqizhixin: 如果AI的记忆不必随着每多一句话而膨胀呢?牛津大学、Technion、AITHYRA 等…
介绍了KV-Compression Aware Training (KV-CAT) 方法,该方法鼓励Transformer在训练过程中学习可压缩的键值缓存,在不牺牲性能的情况下提高长上下文任务的记忆效率。
基于门控关联检索的通用三重潜在压缩
本文介绍了通用三重潜在循环模型,该模型将令牌对交互压缩为潜在状态,并提出一种改进精确召回的门控关联检索变体。该混合模型在字节级WikiText-2和分词语言基准上优于Transformer,实现了高达41.9%的关联召回率(对比25%)。