🤖 machine learning

Remembering Distinct Items, Not Tokens: A Learnable Dirichlet-Process Cache Between State-Space Models and Attention

本文介绍了一种基于狄利克雷过程聚类的可学习稀疏缓存，它仅为新颖项目分配内存槽，使状态空间模型能够通过追踪不同项目而非总标记数来实现全注意力召回效率，并优于固定预算的驱逐策略。

原作者： Siddharth Pal, Viktoria Rojkova

发布于 2026-07-14

📖 1 分钟阅读☕ 轻松阅读

原作者： Siddharth Pal, Viktoria Rojkova

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图记住一位朋友讲的一个冗长且乏味的故事，而这位朋友还总是喋喋不休地重复。 “猫坐在垫子上。猫坐在垫子上。猫坐在垫子上。”

大多数计算机大脑处理这种情况的方式要么是两种极端的做法。第一种方式就像一位极其严谨的图书管理员，他会将你朋友说的每一个词都记录在一张新的索引卡上。如果你的朋友说了一个小时，这位图书管理员就会拥有一叠像摩天大楼一样高的卡片。这被称为“全注意力”（full attention）。它能完美地记住一切，但速度很慢，而且卡片的堆叠会变得非常庞大且沉重。

第二种方式就像一个短期记忆游戏，你只能持有固定数量的卡片，比如 32 张。一旦有了新卡片，你就必须扔掉一张旧卡片来腾出空间。这被称为“固定状态模型”（fixed-state model，例如 Mamba 或 S4）。它非常快速且轻量，但如果你的朋友告诉你一个秘密，然后重复了一百遍，你的大脑可能会因为这些噪音而变得非常混乱，以至于在堆叠变得太满时，你把那个秘密给忘了。

中间地带：“新颖性”侦探

这篇论文介绍了一个聪明的第三种选择：一位聪明的侦探，只有当他听到一些“新东西”时才会动笔记录。

如果你的朋友说，“猫坐在垫子上”，侦探会把它记下来。如果他们再次说这句话，侦探只是点点头并说，“我已经知道了”，而不会浪费纸张。他们只为听到的“独特”事物创建一个新的记忆槽位，而不是为这些事物每次重复出现时都创建一个。

作者称之为狄利克雷过程缓存（Dirichlet-Process Cache）。这是一个高级的名字，其规则是：“如果这个新信息与我已有的信息非常不同，就创建一个新槽位；如果它很相似，就只需更新旧的那个。”

“惊喜度”恒温器

论文还提出了第二个更聪明的版本。想象一下，这位侦探拥有一个**“惊喜度”恒温器**。

如果你的朋友开始讲一个充满新角色、情节离奇的新故事，侦探的“惊喜度”就会上升。他们会开启更多的记忆槽位，以捕捉所有新的细节。
一旦故事稳定下来，他们开始重复那些老掉牙的笑话，侦探的“惊喜度”就会下降。他们会关闭多余的槽位并整理好记忆，让记忆重新变得小巧且高效。

这使得记忆能够根据需要增长，并在不需要时收缩，而不会陷入固定的限制或失控的卡片堆叠中。

这篇论文实际证明了什么（以及没能证明什么）

研究人员以一种非常受控的方式测试了这个想法。他们不仅仅是靠猜测；他们运行了模拟和实验来观察其效果。

巨大的胜利： 在测试中，这个“新颖性”侦探可以像记录每一个词的图书管理员一样完美地记住故事。但关键在于：当故事中存在大量重复时（例如，独特的想法比单词数量少四分之一），这位侦探只需要四分之一的内存。
现实世界检验： 他们在四种不同类型的现实数据流上进行了测试：电影推荐、计算机系统日志、医院病人记录和保险理赔。在每种情况下，这位侦探都成功地追踪了独特项目（如独特的电影或独特的医疗代码），同时忽略了成千上万次的重复条目。例如，在 150,000 条保险理赔的数据流中，这位侦探只需要存储大约 3,933 个唯一代码就能记住所有重要的内容，而标准的“固定预算”系统（会丢弃旧信息）则会无法记住那些稀有的、重要的细节。
学习部分： 作者还展示了这种“新颖性”规则不需要由人类进行硬编码。他们训练了一个微小的、简单的门控（仅有两个数字）来自主学习这一规则。令人惊讶的是，一个更大、更复杂的门控反而失败了。这表明，秘诀不在于拥有一个巨大的大脑，而在于拥有正确的“归纳偏置”（inductive bias）——即一种优先考虑新颖性的特定思维方式。

这篇论文排除了什么

论文非常明确地说明了它不是什么。

它不是解决所有问题的万灵药。作者明确指出，这是一个在受控数据上进行的“机制研究”。他们尚未在处理庞大书籍的完整、现实世界的语言模型（如聊天机器人）上测试过它。那是未来“伴随研究”的工作。
如果故事从未改变，它并不比“固定预算”更好。如果输入是稳定且可预测的，一个简单的固定大小的记忆同样有效。只有当故事变得混乱并改变主意时，“惊喜度”版本才会大放异彩。
它不是在所有情况下都取代“图书管理员”（全注意力）。如果你只需要读一次故事且不再读第二次，图书管理员可能就足够了。但如果你需要读完一个长篇故事并在稍后回答相关问题，侦探会高效得多。

底线

这篇论文表明，通过将记忆视为一种“新颖性检测器”而非“单词计数器”，我们可以构建出能够记住长篇故事中重要且独特的部分，而不会被重复内容所淹没的 AI。这是一个比图书管理员更便宜、比短期记忆游戏更聪明的中间地带。

然而，作者谨慎地表示，这是一个有前景的步骤，而非最终成品。他们已经证明了该机制在特定任务和真实数据流上的有效性，但将其应用于巨大的、现实世界语言模型中的终极测试，仍将在未来到来。目前，他们已经证明，有时通过“减少记忆”（即忽略重复）才是更好地记住更多内容的方法。

技术摘要：记忆离散项而非 Token

1. 问题与差距

序列模型面临着记忆效率与关联召回能力之间的根本权衡。

固定状态模型（如 S4, Mamba）： 将无限的历史压缩为有限的状态。虽然计算成本低（线性时间，恒定推理规模），但受限于“状态维度”。它们无法进行超出状态规模的基于内容的召回，从而导致召回性能存在硬性上限。
全注意力机制（Full Attention）： 通过为每个 Token 维护一个键-值（key-value）条目来突破这一上限。这实现了完美的召回，但会产生二次方计算成本以及随序列长度线性增长的缓存。

本文识别了一个常被忽视的“中间地带”：基于离散项的稀疏缓存。现实世界的流数据包含显著的冗余（重复的 Token）。作者认为，对于关联召回而言，关键信息是由离散项而非每一个 Token 携带的。一个仅在输入具有真正新颖性时才分配新槽位，并将冗余输入合并到现有槽位的记忆系统，其存储的信息量将与离散项的数量（ $K$ ）成正比，而非总 Token 数（ $L$ ），从而以极小的内存代价实现媲全注意力的召回效果。

2. 方法论

作者提出了一种 狄利克雷过程（Dirichlet-Process, DP）缓存，作为深度循环骨干网络的记忆算子。其核心机制是一个源自 DP-means 聚类（狄利克雷过程混合模型的小方差、最大后验似然极限）的在线分配规则。

核心机制：基于新颖性的分配

对于输入的 Token $(k_t, v_t)$ ，系统会计算其相对于已占用槽位的新颖性：
$\text{nov}_t = 1 - \max_{i \in \text{slots}} \text{sim}(k_t, \kappa_i)$
其中 $\text{sim}$ 为余弦相似度。

分配规则： 如果 $\text{nov}_t > \tau$ （阈值），则创建一个新槽位。否则，将该 Token 合并到最近的现有槽位中。
结果： 缓存大小收敛于离散项的数量（ $K$ ），而非序列长度（ $L$ ）。
读取： 查询通过对（较小的）槽位集合进行注意力计算来完成，而非针对整个 Token 流。

两种方法

静态狄利克雷过程缓存： 使用固定的浓度参数（阈值 $\tau$ ）。缓存大小追踪流中的离散项。
惊喜度自适应缓存（Surprise-Adaptive Cache）： 浓度是动态的。由近期分配率驱动的“惊喜信号”（温度 $T_t$ $T_{t}$ ）会调整槽位预算（ $M_t$ $M_{t}$ ）。
- 在出现大量新颖内容的过程中，预算会扩张。
- 当流变得可预测（重复项占主导）时，预算会收缩，剔除使用频率最低的槽位。
- 这使得缓存能够追踪非平稳的需求，同时保持有界。

可学习性

论文研究了分配规则是否可以进行端到端学习。他们使用一个可微的新颖性阈值门控（ $g_t = \sigma(a(\text{nov}_t - b))$ ）取代了硬阈值。

发现： 仅通过任务损失训练的一个极简的两参数门控，即可精确恢复分配规则。
启示： 成功之处在于归纳偏置（使用新颖性特征），而非门控本身的容量。缺乏这种特定偏置的过度参数化门控无法学习该规则。

3. 主要贡献与主张

本文声称其贡献限于受控机制研究范围内：

识别差距： 将基于 DP 的、基于新颖性分配的稀疏记忆定位在固定状态递归模型与全注意力机制之间。
离散项效率： 缓存按离散项的数量进行存储，在达到全注意力召回水平的同时，仅使用极小部分的内存（例如，在 4 倍冗余度下实现 4 倍内存缩减）。
优越的前沿表现： 在召回率与规模的权衡前沿上，优于固定预算剔除缓存（如 H2O, StreamingLLM, Recency），特别是在重尾分布的任务中，固定预算会浪费容量处理频繁项而忽略了稀有的长尾项。
混合架构： 当与状态空间骨干网络结合时，该模型能同时实现长程整合（通过 SSM）和内容寻址召回（通过缓存），并达到测试中最低的内存成本。
可学习的分配： 分配策略可以通过一个极简的学习门控进行端到端恢复，证明了给定正确的归纳偏置时，该规则是任务损失下的涌现属性。
现实世界有效性： “离散项”属性在四个无关领域（推荐、系统日志、临床事件、保险理赔）均成立，无需针对特定领域进行调优。

4. 实验结果

合成关联召回

设置： 包含 $K$ 个离散项且每个项重复 $r$ 次的流（$L = rK$）。
性能： DP 缓存实现了 1.00 召回率（匹配全注意力），同时仅存储了 $K$ 个槽位。
对比：
- 固定状态 SSM：约 0.83 召回率（受限于状态维度）。
- 固定预算剔除法（H2O/Recency）：约 0.70–0.92 召回率（未能以相等预算覆盖所有离散项）。
- DP 缓存： 在离散项预算下实现完美召回。

现实世界重尾任务

MovieLens（类型召回）： DP 缓存使用约 1,370 个槽位（每个电影一个）实现了 0.99 召回率。固定预算缓存（H2O, StreamingLLM, SnapKV）在相同预算下仅实现 0.23–0.61 召回率，因为它们优先考虑热门电影而丢弃了稀有电影。
系统日志（Loghub）： 分配规则充当了在线解析器，将 2,000 行日志分组为 15 个模板（真实值为 14），准确率为 0.89，实现了 143 倍的压缩。
临床与保险数据： 观察到类似结果，缓存存储的是离散实体（药物、诊断代码）而非事件，实现了约 0.99 的召回率。

可学习性与成本

学习门控： 一个两参数新颖性门控达到了 1.00 召回率，与基于规则的方法一致。而一个过度参数化的 MLP 门控则失败了（0.37 召回率），证实了归纳偏置而非容量才是核心要素。
成本权衡：
- 构建成本： $O(N \cdot K)$ （由于需要寻找最近槽位，比注意力的 $O(N)$ 追加操作更重）。
- 读取成本： $O(K)$ （相对于序列长度是常数）对比注意力的 $O(N)$ 。
- 结论： 对于需要重复读取的长上下文推理，当内存是主要约束时，该缓存更具优势。

自适应浓度（非平稳需求）

在工作集大小波动的流中，自适应缓存在相同的平均成本下优于固定预算缓存。
它在“困难”阶段（高新颖性）扩张容量，并在“容易”阶段（高重复性）进行整合；而固定预算要么浪费资源（按最坏情况配置），要么导致召回失败（按平均情况配置）。

5. 重要性与范围

论文明确将其贡献定位为将现有组件进行综合（在线 DP-means、聚类注意力、增长型缓存）为一种新型的深度循环模型记忆算子。

范围： 证据源于受控机制研究，规模适中（合成探测和真实流统计）。基线是重新实现的机制，而非完整的生产级系统（例如，并非完整的 Mamba 或 Memorizing Transformers）。
局限性： 研究中的键（Keys）是通过构造实现清晰分离的。论文承认，在真实的、存在重叠的数据上，阈值和合并步骤需要精细的工程设计。
未来工作： 作者指出，在“真实骨干网络、真实语料库验证”（例如，带有此缓存的 Mamba 骨干网络在 enwik8 或 WikiText-103 上）是必要的下一步，并正在伴随研究中进行。

结论： 本文确立了通过狄利克雷过程分配规则实现的、以离散项数量为规模的稀疏缓存，为有限状态空间模型与全注意力机制之间提供了一个可学习且高效的中间方案，并在涉及冗余和重尾分布的召回任务中，表现优于固定预算的剔除策略。