论文

记住不同的项目,而不是词元:状态空间模型和注意力之间的可学习狄利克雷过程缓存

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

模型推理KV Cache

摘要

固定状态序列模型将无界的过去压缩为有界的状态,这将它们的联想回忆限制在大致状态维度上;注意力通过在二次计算中为每个标记保留一个键值条目以及随序列增长的缓存来逃避上限。我们研究中间立场:稀疏缓存仅在输入新颖时才分配插槽,因此其大小跟踪不同项目的数量而不是词元的数量。分配规则是 DP 均值聚类规则,即狄利克雷过程混合的小方差限制,不用作潜变量推理,而是用作深度循环骨干的键值内存运算符。我们以两种形式开发它,一种是具有固定浓度的静态缓存,另一种是其浓度遵循最近的新颖率的惊喜自适应变体。在具有冗余的受控关联召回基准上,我们表明缓存与全注意力召回相匹配,同时仅存储不同的项目,它在召回与大小边界上主导固定预算驱逐缓存,并且在状态空间主干上,它以任何测试模型的最低内存回答召回查询和远程聚合。这种分配是端到端可学习的:仅针对任务损失进行训练的二参数新颖性阈值门可以准确地恢复规则,而过度参数化的门则失败,因此有效成分是归纳偏差而不是容量。证据是一系列适度规模的受控机制研究,在四个真实流(建议、系统日志、临床事件和保险索赔)上确认了不同项目的属性;在一项配套研究中,我们进行了真正的骨干、真实的语料库语言验证。