让模型只记住不同项目,而非所有输入 tokens,提升记忆效率
Remembering Distinct Items, Not Tokens: A Learnable Dirichlet-Process Cache Between State-Space Models and Attention

- 用狄利克雷过程聚类规则动态分配记忆槽,仅当输入新项目时才存
- 在关联回忆任务中达到注意力模型效果,但内存仅存不同项数量
- 可端到端训练,适合需要长程记忆的序列建模场景
固定状态序列模型将无限历史压缩为有限状态,导致关联回忆能力受限于状态维度;注意力机制虽突破限制,但需为每个 token 存储键值对,计算开销为二次方且缓存随序列增长。本文探索中间路径:构建稀疏缓存,仅当输入为新项目时才分配存储槽,缓存大小与不同项目数成正比。采用狄利克雷过程均值聚类(DP-means)作为分配规则,不用于隐变量推断,而是作为深层循环主干的键值记忆操作符。提出两种形式:固定浓度静态缓存和基于近期新颖率自适应浓度的变体。在含冗余的受控关联回忆基准测试中,该缓存实现与全注意力相当的召回性能,同时仅存储不同项目;在召回-大小权衡曲线上优于固定预算淘汰缓存;在状态空间主干上,以最低内存完成回忆查询与长程聚合。分配规则可端到端学习:仅两个参数的新颖性阈值门控在任务损失下即能精确恢复规则,而过参数化门控则失败,说明关键在于归纳偏置而非容量。证据来自小规模控制机制研究,四个真实数据流(推荐、系统日志、临床事件、保险理赔)验证了不同项目特性;真实主干与真实语料的验证将在配套研究中展开。
原文摘要 · Abstract (English)
Fixed-state sequence models compress an unbounded past into a bounded state, which caps their associative recall at roughly the state dimension; attention escapes the cap by keeping a key-value entry for every token, at quadratic compute and a cache that grows with the sequence. We study the middle ground: a sparse cache that allocates a slot only when an input is novel, so its size tracks the number of distinct items rather than the number of tokens. The allocation rule is the DP-means clustering rule, the small-variance limit of a Dirichlet-process mixture, used not as latent-variable inference but as the key-value memory operator for a deep recurrent backbone. We develop it in two forms, a static cache with a fixed concentration and a surprise-adaptive variant whose concentration follows the recent novelty rate. On a controlled associative-recall benchmark with redundancy we show that the cache matches full-attention recall while storing only the distinct items, that it dominates a fixed-budget eviction cache on the recall-versus-size frontier, and that on a state-space backbone it answers both a recall query and a long-range aggregate at the lowest memory of any model tested. The allocation is learnable end to end: a two-parameter novelty-threshold gate trained on the task loss alone recovers the rule exactly, whereas an over-parameterized gate fails, so the operative ingredient is the inductive bias rather than capacity. The evidence is a family of controlled mechanism studies at modest scale, with the distinct-items property confirmed on four real streams (recommendation, systems logs, clinical events, and insurance claims); a real-backbone, real-corpus language validation is pursued in a companion study.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。