记住独特项目而非标记：一种可学习的狄利克雷过程缓存，介于状态空间模型与注意力机制之间

# 状态空间模型与注意力机制之间的可学习狄利克雷过程缓存  
来源：https://arxiv.org/html/2607.09889  

## 记住不同项目而非标记：状态空间模型与注意力机制之间的可学习狄利克雷过程缓存（预印本）  

###### 摘要  
固定状态序列模型（如结构化状态空间模型）将无限历史压缩为有界状态，这使其成本低廉，但将关联检索能力限制在大致状态维度以内；注意力机制通过为*每个*标记保留一个键值条目来突破限制，但计算复杂度呈二次方，且缓存随序列长度增长。我们研究中间地带：一种稀疏缓存，仅在输入为*新*时分配一个槽位，因此其大小跟踪*不同*项目的数量而非标记数量。分配规则是 DP-means 聚类规则，即狄利克雷过程混合的小方差极限，但并非用作潜在变量推断，而是作为深度循环骨干网络的键值内存算子——据我们所知，这种综合尚未见报道，尽管每个组成部分（在线 DP-means 聚类、聚类注意力、增长缓存）都独立存在。我们以两种形式开发此缓存：一种是具有固定浓度的静态狄利克雷过程缓存，另一种是惊奇自适应变体，其浓度通过运行中的惊奇信号进行调整，在新内容爆发时扩大容量，在流稳定后回收容量。在一个带有冗余性的受控关联检索基准测试中，我们在机制层面展示了：该缓存匹配全注意力检索性能，同时仅存储不同项目（四倍冗余时减少四倍）；在检索-大小前沿上优于固定预算驱逐缓存；在状态空间骨干网络上，它以测试模型中最低的内存同时回答了一个检索查询和一个长程聚合查询。我们进一步证明，分配是*端到端可学习的*：一个仅通过任务损失训练的两参数新颖性阈值门完美恢复该规则，而过度参数化的门则失败，因此起作用的因素是归纳偏差而非容量。我们明确范围：核心证据是一系列在适度规模下进行的受控机制研究——关联检索、与已部署驱逐缓存的检索-大小前沿比较、可学习门消融实验、状态空间混合、成本缩放、以及自适应浓度测试——在多个机器（Apple M1 和 M4，以及 Google Colab 上的 NVIDIA T4）的 CPU 和 GPU 上运行并交叉验证，不同项目属性在四个真实流上得到确认：推荐系统、系统日志、临床事件和保险理赔。引用的系统作为重新实现的机制而非完整模型进入，而基于真实骨干和真实语料的语言验证在配套研究中进行。  

## 1 问题与差距  

序列模型将不断增长的历史总结为固定大小的状态。结构化状态空间模型（S4 及其对角后继，以及 Mamba 的选择性循环）通过有界状态实现这一点，因此以线性时间运行且推理大小恒定；其代价——近期一系列工作已精确说明——是这种有界状态无法执行超过其自身维度的基于内容的检索——这是 Jelassi 等人的形式化复制分离结果和 Zoology 基准测量的经验检索-吞吐量权衡。注意力机制通过为每个标记保留一个键值条目并根据内容读取来消除限制，这就是它擅长检索的原因，但其缓存随序列长度线性增长，计算复杂度为二次方。这两者是同一轴线的两端：保留多少状态。  

在这两端之间，有一个自然的想法。如果长流包含许多重复或冗余的标记——正如真实序列那样——那么检索所需的信息由*不同*项目承载，而非每个标记，因此一个仅当输入真正新时才打开新槽位，否则将其合并到现有槽位中的存储器，将只存储不同项目的数量级，而不会更多。我们将此规则置于其历史中，而非声称其全新。作为一种聚类原语，它是古老的：确定性的“当最近簇的距离超过阈值时打开新簇，否则加入该簇”规则正是*DP-means*（Kulis 和 Jordan），即狄利克雷过程混合的小方差、最大后验极限，在此之前是经典的领导者或在线聚类规则。我们从中汲取的是分配规则；我们贡献的是将其用作深度序列模型的*内存算子本身*，其中簇集是键值缓存，其基数是检索成本。当然，狄利克雷过程已经为*经典*序列模型提供了无界潜在基数——无限隐马尔可夫模型、层次狄利克雷过程及其切换线性动态系统变体、以及无限循环切换线性动态系统，后者通过距离依赖的餐厅过程增长其离散模式计数——但这是在浅层经典骨干网络上作为潜在状态推断，而非作为深度循环模型的键值内存。  

深度学习中最近的工作各止步于我们方法之前一步，有必要精确说明差异所在。*聚类和路由注意力*（Vyas 等人；Roy 等人；以及 Reformer 的局部敏感哈希桶，Kitaev 等人）已经合并相似键以简化注意力——但它们在有限上下文的全注意力层内进行聚类，作为每批次分组，目的是近似全 softmax，其簇数是固定超参数而非随无界流中不同项目数量增长；它们也没有将机制定位在固定状态下限之下，或询问分组是否可以作为分配策略学习。*增长内存 Transformer*（Memorizing Transformers、infinite-former、kNN 语言模型、神经情景控制、Kanerva 机、乘积键记忆、现代 Hopfield 网络）通过最近邻检索、连续注意力或共轭更新增长或寻址外部内存，但它们保留每个标记的条目且不按新颖性分配。*固定预算稀疏缓存*（StreamingLLM、H2O、地标注意力）保持*恒定*预算并进行*驱逐*——按最近性或累计注意力质量——这是与增长缓存到不同项目数量相反的控制轴。*惊奇门控记忆*（Titans、EM-LLM）按新颖性增长，但通过可学习的写入门而非 DP-means 分配。因此，我们的贡献是一种综合而非新原语：DP-means 在线分配用作缓存本身，大小适应无界流的不同项目，放置在状态空间到注意力轴上，并证明可通过最小可学习门端到端恢复。我们尚未发现这种综合被报道过，并且我们有意以狭窄形式提出这一新颖性声明，因为每个组成部分都独立存在。  

##### 两种方法。我们以两种形式开发中间地带存储器，这两者是本文的主题。*方法 I，狄利克雷过程缓存*，以*固定*浓度按新颖性分配槽位，因此缓存大小跟踪不同项目数量；这是基础构造。*方法 II，惊奇自适应狄利克雷过程缓存*，使浓度*自适应*：一个类似温度的惊奇信号，由近期分配率驱动，在新内容爆发时提高槽位预算，并在流变得可预测时降低预算、合并槽位。浓度因此由模型自身的惊奇调整，在新内容到达时接纳容量，并在流稳定后回收容量，因此缓存跟踪非平稳需求同时保持有界。两者共享相同的按新颖性分配算子和读取方式，仅在于浓度是常数还是温度驱动的时间表。  

##### 贡献。围绕这两种方法，我们提出以下声明，每一项都限定于本研究的受控设置。(i) 我们识别并定位了差距：一个用于深度循环骨干网络的狄利克雷过程按新颖性分配稀疏存储器，介于固定有界状态和全注意力缓存之间。(ii) 狄利克雷过程缓存（方法 I）存储的是*不同*项目而非标记的数量级，在缓存容量一小部分的情况下匹配全注意力检索，并在检索-大小前沿上优于部署的固定预算驱逐缓存（H2O、StreamingLLM、SnapKV 和最近性），无论是在合成探针上还是在真实重尾任务上。(iii) 与状态空间骨干网络结合，它在测试的所有模型中以最低内存同时执行检索和长程整合。(iv) 分配可通过一个最小新颖性阈值门从任务损失端到端学习，因此策略无需手动指定，且起作用的因素是归纳偏差而非门容量。(v) 不同项目属性在来自无关领域的四个真实流上保持不变——推荐、系统日志、临床事件和保险理赔——无需对阈值进行每领域调优，并且在系统日志上，同一规则本身就是在线日志解析器。(vi) 惊奇自适应缓存（方法 II）在需求非平稳时以相同平均成本击败固定预算，而在需求平稳时相比匹配的固定预算无优势；其增益正是固定预算通过始终为最困难时刻配置容量而浪费的部分；这一边界的两个方面在真实流上均得到确认。  

## 2 方法  

我们以两种形式呈现缓存，它们共享相同的按新颖性分配算子和读取方式，仅浓度不同：静态的*狄利克雷过程缓存*（方法 I），其浓度为固定阈值；以及*惊奇自适应*缓存（方法 II），其浓度为适应需求的温度驱动时间表。  

##### 设置及工作示例。探针是关联检索，一种受控合成任务，旨在隔离单一能力而非对系统进行基准测试。一个流呈现 K 个不同的 (key, value) 项目，每个重复 r 次，顺序随机，因此流长度为 L = rK，而只有 K 个标记是不同的；模型读取整个流，然后收到一个查询键，必须返回该键绑定的值。键是向量，值是类别标签，冗余度 r 是有意设置的杠杆，因为它将标记数量 L 与不同项目数量 K 拉开，这正是新颖性分配缓存应跟踪的量。具体来说，假设 K = 3 个不同事实，如 (apple, red)、(sky, blue) 和 (grass, green)，每个重复 r = 2 次，打乱成一个长度为 L = 6 的流，例如 (sky, blue)、(apple, red)、(grass, green)、(apple, red)、(sky, blue)、(grass, green)，随后是查询 "apple?"，答案为 "red"。一个*槽位*是模型内存中的一个条目，保存单个键及其值。全注意力为每个标记保留一个槽位，因此存储所有 L = 6 个条目；固定状态模型根本不保留每个项目的槽位，而是将流折叠成一个有界状态，因此一旦事实数量超过其状态维度，就无法将它们分开；狄利克雷过程缓存仅在标记新时打开一个槽位，因此存储 K = 3 个不同事实，并将每次重复合并到它已产生的槽位中。因此，槽位数是我们报告的内存成本，在冗余度 r 下，缓存存储的条目比注意力少 r 倍，同时可回答相同查询。真实流当然更长且更具冗余性——同一单词、名称或标识符在文档中多次出现——这正是这个小而透明的探针旨在隔离暴露的机制。  

##### 方法 I：按新颖性分配缓存。缓存是一组槽位，每个槽位保存一个键和一个值，并带有每个槽位的使用计数器。对于输入的 (k_t, v_t)，我们测量其相对于已占用槽位的新颖性：  
nov_t = 1 - max_{i ∈ 槽位} sim(k_t, κ_i),  (1)  
其中 sim 是余弦相似度，κ_i 是第 i 个槽位键。狄利克雷过程规则然后决定分配：如果输入与所有已占用槽位不相似，则打开新表，否则加入最近的表：  
if nov_t > τ: append (k_t, v_t) (新槽位); else: merge into arg max_i sim(k_t, κ_i),  (2)  
这正是 DP-means 规则（狄利克雷过程混合的最大后验、小方差极限）：远离所有现有簇的点以阈值 τ（扮演浓度角色）所控制的力度开始新簇，接近簇的点被吸收。该规则是否真正是狄利克雷过程，而不仅仅是阈值，通过模拟确定：作为底层随机过程运行时，槽位数以 α ln N 增长，浓度作为斜率（Antoniak, 1974），这是狄利克雷过程的标志；而在从有限项目类型集合（本文始终如此）抽取的数据上，同一规则反而*饱和*于真实的不同类型数量，精确分配与数据复杂度所需一样多的槽位。由于重复与它们产生的槽位相似，它们合并而非分配，因此缓存大小收敛到不同项目数量。槽位键保留在其第一次出现时；将其更新为合并出现的运行质心（在线 k-means）在低噪声下等价，在项目间分离附近略更稳定，因此这一选择不改变结果。查询 q 通过（少量）槽位上的注意力读取：  
\hat{v} = ∑_i softmax_i( sim(q, κ_i) / θ ) ν_i,  (3)  
解码为值类别。因此读取成本由槽位数（接近 K 而非 L）决定。  

##### 可学习门。阈值规则可以用可微分门替换，以便从任务中学习分配。我们为每个标记分配一个保留概率 g_t = σ(a (nov_t - b))，其中可学习标量 a, b，新颖性特征按因果方式计算，nov_t = 1 - max_{s < t} sim(κ_s, k_t)（即使用过去槽位，而非未来）。在训练期间，g_t 视为伯努利试验的概率，在测试时硬性阈值化（g_t > 0.5 则分配）。该门不访问未来，不违反因果性，且当反向传播通过离散决策时使用 Gumbel-Softmax 重新参数化。此可学习变体保留了 DP-means 的归纳偏差——分配仅取决于相似性——但如果数据要求不同的策略，则允许调整阈值。  

##### 方法 II：惊奇自适应缓存。在方法 I 中，槽位数自动增长，但不同项目数量是无界的——缓存不会在平稳流中无限增长，但会适应任何出现的事实数量。对于需要硬上限的部署场景，我们引入一个温度 T_t，该温度跟踪近期分配率：在每个步骤，二进制分配决策 o_t = 1（如果 nov_t > τ，否则 0），并保持其指数移动平均：  
T_t = (1 - η) T_{t-1} + η o_t,  (4)  
因此 T_t 在新内容爆发期间高，在流可预测且重复占主导时低。有效预算随后随温度上升：  
M_t = M_0 + β T_t,  (5)  
每当槽位数超过 M_t 时，最少使用的槽位（使用计数器最小，每一步衰减）被驱逐：  
while |槽位| > M_t: evict arg min_i u_i.  (6)  
因此缓存在新内容爆发时变热，接纳容量，并在流稳定后冷却，通过合并回收槽位，从而保持一个跟踪当前需求的有界工作集。这是狄利克雷过程浓度的自适应调整，一个温度驱动的时间表。