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

2026-07-10

AI Summary

This work addresses the limitations of sequence modeling approaches: state space models are constrained by fixed state dimensions, while attention mechanisms, despite their long-range memory capabilities, suffer from quadratic computational complexity and linearly growing cache overhead. To overcome these issues, the authors propose a sparse memory mechanism based on a learnable Dirichlet process cache that allocates memory slots only when inputs exhibit sufficient novelty, thereby scaling cache size with the number of distinct items rather than total token count. The method integrates DP-means clustering with a learnable novelty threshold gating mechanism, enabling end-to-end training. Experiments demonstrate that on redundant associative recall tasks, the approach matches full attention performance with substantially smaller cache sizes and outperforms fixed-budget methods in the memory-recall trade-off, validating the efficacy of “distinct-items” caching.

Abstract

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.

Problem

• associative recall
• sequence modeling
• memory efficiency
• distinct items
• attention mechanism

Innovation

• Dirichlet-process cache
• distinct-item memory
• DP-means clustering
• state-space models
• sparse attention