🤖 machine learning

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

This paper introduces a learnable, sparse cache based on Dirichlet-process clustering that allocates memory slots only for novel items, enabling state-space models to achieve full-attention recall efficiency by tracking distinct items rather than total tokens while outperforming fixed-budget eviction strategies.

Original authors: Siddharth Pal, Viktoria Rojkova

Published 2026-07-14

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Siddharth Pal, Viktoria Rojkova

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to remember a long, boring story told by a friend who repeats themselves constantly. "The cat sat on the mat. The cat sat on the mat. The cat sat on the mat."

Most computer brains handle this in one of two extreme ways. The first way is like a super-organized librarian who writes down every single word your friend says on a new index card. If your friend talks for an hour, the librarian has a stack of cards as tall as a skyscraper. This is called "full attention." It remembers everything perfectly, but it's slow and the stack of cards gets huge and heavy.

The second way is like a short-term memory game where you can only hold a fixed number of cards, say 32. As soon as you get a new card, you have to throw an old one away to make room. This is called a "fixed-state model" (like Mamba or S4). It's super fast and light, but if your friend tells you a secret and then repeats it 100 times, your brain might get so confused by the noise that you forget the secret entirely once the stack gets too full.

The Middle Ground: The "Novelty" Detective

This paper introduces a clever third option: a smart detective who only writes down a card when they hear something new.

If your friend says, "The cat sat on the mat," the detective writes it down. If they say it again, the detective just nods and says, "I already know that," and doesn't waste paper. They only create a new slot in their memory for the distinct things they hear, not for every single time those things are repeated.

The authors call this a Dirichlet-Process Cache. It's a fancy name for a rule that says: "If this new piece of information is very different from what I already have, make a new slot. If it's similar, just update the old one."

The "Surprise" Thermostat

The paper also suggests a second, even smarter version of this detective. Imagine the detective has a surprise thermostat.

If your friend starts telling a wild, new story with lots of new characters, the detective's "surprise" level goes up. They open up more slots in their memory to catch all the new details.
Once the story settles down and they start repeating the same old jokes, the detective's "surprise" level drops. They close the extra slots and tidy up, keeping their memory small and efficient again.

This allows the memory to grow when it's needed and shrink when it's not, without ever getting stuck with a fixed limit or a runaway stack of cards.

What the Paper Actually Proved (and What It Didn't)

The researchers tested this idea in a very controlled way. They didn't just guess; they ran simulations and experiments to see if it worked.

The Big Win: In their tests, this "novelty" detective could remember the story just as perfectly as the librarian who wrote down every word. But here's the kicker: when the story had a lot of repetition (like 4 times more words than unique ideas), the detective only needed one-fourth the amount of memory.
The Real-World Check: They tried this on four different types of real-world data streams: movie recommendations, computer system logs, hospital patient records, and insurance claims. In every case, the detective successfully tracked the unique items (like unique movies or unique medical codes) while ignoring the thousands of repeated entries. For example, on a stream of 150,000 insurance claims, the detective only needed to store about 3,933 unique codes to remember everything important, whereas a standard "fixed budget" system that throws things away would fail to remember the rare, important details.
The Learning Part: The authors also showed that this "novelty" rule doesn't need to be hard-coded by a human. They trained a tiny, simple gate (just two numbers) to learn this rule on its own. Surprisingly, a bigger, more complex gate actually failed to learn it. This suggests that the secret isn't having a huge brain, but having the right kind of "inductive bias"—a specific way of thinking that prioritizes novelty.

What the Paper Rules Out

The paper is very clear about what this is not.

It is not a magic fix for everything. The authors explicitly state that this is a "mechanism study" run on controlled data. They have not yet tested it on a full, real-world language model (like a chatbot) reading a massive book. That is a job for a future "companion study."
It is not better than a fixed budget if the story never changes. If the input is steady and predictable, a simple fixed-size memory works just as well. The "surprise" version only shines when the story gets chaotic and changes its mind.
It is not a replacement for the "librarian" (full attention) in every situation. If you need to read a story once and never again, the librarian might be fine. But if you need to read a long story and answer questions about it later, the detective is much more efficient.

The Bottom Line

The paper suggests that by treating memory like a "novelty detector" rather than a "word counter," we can build AI that remembers the important, unique parts of a long story without getting overwhelmed by the repetition. It's a middle ground that is cheaper than the librarian and smarter than the short-term memory game.

However, the authors are careful to say this is a promising step, not a finished product. They have proven the mechanism works on specific tasks and real data streams, but the ultimate test—using it in a giant, real-world language model—is still on the horizon. For now, they've shown that sometimes, remembering less (by ignoring the repeats) is actually the best way to remember more.

Technical Summary: Remembering Distinct Items, Not Tokens

1. The Problem and the Gap

Sequence models face a fundamental trade-off between memory efficiency and associative recall capability.

Fixed-State Models (e.g., S4, Mamba): Compress an unbounded history into a bounded state. While computationally cheap (linear time, constant inference size), they are limited by the "state dimension." They cannot perform content-addressed recall beyond roughly the size of their state, leading to a hard cap on recall performance.
Full Attention: Escapes this cap by maintaining a key–value entry for every token. This enables perfect recall but incurs quadratic compute costs and a cache that grows linearly with the sequence length.

The paper identifies a "middle ground" often overlooked: sparse caching based on distinct items. Real-world streams contain significant redundancy (repeated tokens). The authors posit that for associative recall, the critical information is carried by distinct items, not every token. A memory system that allocates a new slot only when an input is genuinely novel, and merges redundant inputs into existing slots, could store information proportional to the number of distinct items ( $K$ ) rather than the total token count ( $L$ ), achieving recall comparable to full attention with a fraction of the memory.

2. Methodology

The authors propose a Dirichlet-Process (DP) Cache that functions as the memory operator for a deep recurrent backbone. The core mechanism is an online allocation rule derived from DP-means clustering (the small-variance, maximum-a-posteriori limit of a Dirichlet-process mixture).

Core Mechanism: Allocate-on-Novelty

For an incoming token $(k_t, v_t)$ , the system calculates novelty against occupied slots:
$\text{nov}_t = 1 - \max_{i \in \text{slots}} \text{sim}(k_t, \kappa_i)$
where $\text{sim}$ is cosine similarity.

Allocation Rule: If $\text{nov}_t > \tau$ (a threshold), a new slot is created. Otherwise, the token is merged into the nearest existing slot.
Result: The cache size converges to the number of distinct items ( $K$ ) rather than the sequence length ( $L$ ).
Reading: Queries are answered via attention over the (small) set of slots, not the full token stream.

Two Approaches

Approach I: Static Dirichlet-Process Cache: Uses a fixed concentration parameter (threshold $\tau$ ). The cache size tracks the distinct items of the stream.
Approach II: Surprise-Adaptive Cache: The concentration is dynamic. A "surprise signal" (temperature $T_t$ $T_{t}$ ), driven by the recent allocation rate, adjusts the slot budget ( $M_t$ $M_{t}$ ).
- During bursts of novel content, the budget expands.
- When the stream becomes predictable (repeats dominate), the budget contracts, evicting the least-used slots.
- This allows the cache to track non-stationary demand while remaining bounded.

Learnability

The paper investigates whether the allocation rule can be learned end-to-end. They replace the hard threshold with a differentiable novelty-threshold gate ( $g_t = \sigma(a(\text{nov}_t - b))$ ).

Finding: A minimal two-parameter gate trained solely on task loss recovers the allocation rule exactly.
Implication: The success relies on the inductive bias (using a novelty feature) rather than the capacity of the gate. Over-parameterized gates without this specific bias fail to learn the rule.

3. Key Contributions and Claims

The paper claims the following contributions, scoped to controlled mechanism studies:

Gap Identification: Positions a DP-based, allocate-on-novelty sparse memory between fixed-state recurrence and full attention.
Distinct-Item Efficiency: The cache stores on the order of distinct items, matching full-attention recall while using a fraction of the memory (e.g., a 4x reduction at 4x redundancy).
Superior Frontier: Dominates fixed-budget eviction caches (e.g., H2O, StreamingLLM, Recency) on the recall-versus-size frontier, particularly on heavy-tailed distributions where fixed budgets waste capacity on frequent items and miss rare tails.
Hybrid Architecture: When composed with a state-space backbone, the model achieves both long-range integration (via SSM) and content-addressed recall (via cache) at the lowest memory cost tested.
Learnable Allocation: The allocation policy is recoverable end-to-end via a minimal learned gate, proving the rule is an emergent property of the task loss given the correct inductive bias.
Real-World Validity: The "distinct-items" property holds across four unrelated real-world domains (recommendation, system logs, clinical events, insurance claims) without per-domain tuning.

4. Experimental Results

Synthetic Associative Recall

Setup: Streams with $K$ distinct items repeated $r$ times ($L = rK$).
Performance: The DP cache achieved 1.00 recall (matching full attention) while storing only $K$ slots.
Comparison:
- Fixed-state SSM: ~0.83 recall (limited by state dimension).
- Fixed-budget eviction (H2O/Recency): ~0.70–0.92 recall (failed to cover all distinct items at equal budget).
- DP Cache: Perfect recall at the distinct-item budget.

Real-World Heavy-Tailed Tasks

MovieLens (Genre Recall): The DP cache achieved 0.99 recall using ~1,370 slots (one per distinct movie). Fixed-budget caches (H2O, StreamingLLM, SnapKV) achieved only 0.23–0.61 recall at the same budget because they prioritized popular movies and discarded rare ones.
System Logs (Loghub): The allocation rule acted as an online parser, grouping 2,000 log lines into 15 templates (ground truth: 14) with 0.89 accuracy, achieving 143x compression.
Clinical & Insurance Data: Similar results observed, with the cache storing distinct entities (drugs, diagnosis codes) rather than events, achieving ~0.99 recall.

Learnability and Cost

Learned Gate: A two-parameter novelty gate achieved 1.00 recall, matching the rule-based approach. An over-parameterized MLP gate failed (0.37 recall), confirming that inductive bias, not capacity, is the operative ingredient.
Cost Trade-off:
- Build Cost: $O(N \cdot K)$ (heavier than attention's $O(N)$ append) due to nearest-slot search.
- Read Cost: $O(K)$ (constant relative to sequence length) vs. $O(N)$ for attention.
- Conclusion: The trade-off favors the cache for long-context inference with repeated reads, where memory is the binding constraint.

Adaptive Concentration (Non-Stationary Demand)

On streams where the working set size fluctuates, the adaptive cache outperformed a fixed-budget cache at the same average cost.
It expanded capacity during "hard" phases (high novelty) and consolidated during "easy" phases, whereas a fixed budget either wasted resources (provisioned for worst-case) or failed to recall (provisioned for average).

5. Significance and Scope

The paper explicitly frames its contribution as a synthesis of existing ingredients (online DP-means, clustered attention, growing caches) into a novel memory operator for deep recurrent models.

Scope: The evidence is derived from controlled mechanism studies at modest scale (synthetic probes and real-stream statistics). The baselines are re-implemented mechanisms, not full production systems (e.g., not full Mamba or Memorizing Transformers).
Limitations: The keys in the study are well-separated by construction. The paper acknowledges that on real, overlapping data, the threshold and merge steps would require careful engineering.
Future Work: The authors state that a "real-backbone, real-corpus validation" (e.g., a Mamba backbone with this cache on enwik8 or WikiText-103) is the necessary next step and is pursued in a companion study.

Conclusion: The paper establishes that a sparse cache, sized by the number of distinct items via a Dirichlet-process allocation rule, offers a learnable, efficient middle ground between bounded state-space models and full attention, dominating fixed-budget eviction strategies on recall tasks involving redundancy and heavy-tailed distributions.