DP-Means Distinct-Item Memory
Implementation & benchmark of arXiv:2607.09889 — Remembering Distinct Items, Not Tokens: A Learnable Dirichlet-Process Cache Between State-Space Models and Attention
Source paper: Remembering Distinct Items, Not Tokens: A Learnable Dirichlet-Process Cache Between State-Space Models and Attention arXiv:2607.09889 ⓘ · analyzed Aug 30, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace token-by-token KV storage after an SSM or recurrent encoder with an online allocate-on-novelty cache. A new slot is created only when the incoming key is sufficiently dissimilar from every stored key; otherwise the incoming value is merged into its nearest slot, so repeated or redundant content does not grow the cache.
Formulas
Mathematical statement
The novelty statistic is $\mathrm{nov}_t=1-\max_{i\in\mathrm{slots}}\mathrm{sim}(k_t,\kappa_i)$, where $k_t\in\mathbb{R}^d$ is the incoming key, $\kappa_i\in\mathbb{R}^d$ is the key prototype in occupied slot $i$, and $\mathrm{sim}(a,b)=a^\top b/(\|a\|_2\|b\|_2)$ is cosine similarity. The DP-means allocation rule creates a slot if $\mathrm{nov}_t>\tau$; otherwise it assigns the item to $j_t=\arg\max_i\mathrm{sim}(k_t,\kappa_i)$ and merges it into slot $j_t$. Here $v_t\in\mathbb{R}^{d_v}$ is the value and $\tau\in[0,2]$ is the novelty threshold: larger $\tau$ produces fewer, coarser clusters. A directly implementable merge is the count-weighted update $n_j\leftarrow n_j+1$, $\kappa_j\leftarrow((n_j-1)\kappa_j+k_t)/n_j$, and $u_j\leftarrow((n_j-1)u_j+v_t)/n_j$, where $n_j$ is slot usage and $u_j$ is its stored value. The count-weighted value merge is an adaptation; the paper's essential mathematical contribution is the DP-means novelty/allocation rule.
Implementation notes
Integrate the operator at the write side of a recurrent or SSM language model, immediately after the backbone produces per-token key and value tensors. For each sequence, maintain arrays Kappa[M,d], U[M,dv], Counts[M], and an occupancy count; M is only a safety cap. Pseudocode: for t in range(T): k,v = project(backbone_state[t]); k = k/(norm(k)+eps); if used == 0: allocate(k,v,1); else: sims = Kappa[:used] @ k; j = argmax(sims); novelty = 1-sims[j]; if novelty > tau and used < M: Kappa[used]=k; U[used]=v; Counts[used]=1; used += 1; else: Counts[j]+=1; Kappa[j]=((Counts[j]-1)*Kappa[j]+k)/Counts[j]; U[j]=((Counts[j]-1)*U[j]+v)/Counts[j]. At query positions, produce q from the current hidden state and read the sparse cache using softmax(q @ Kappa[:used].T / sqrt(d)) @ U[:used]; add this vector to the SSM output through a learned residual gate. The paper directly supplies cosine novelty, nearest-slot assignment, and threshold allocation. The count-weighted merge and softmax read are adaptations and should be compared with EMA value updates and hard nearest-slot reads. Log occupied slots, allocation rate, within-slot cosine variance, read latency, and retrieval accuracy. First use a small GRU or diagonal SSM on an associative-recall dataset with fourfold repeated key-value items. Compare full attention, FIFO/LRU fixed-budget eviction, and this cache at equal hidden size and cache-byte budgets. Success means full-attention-level recall with roughly fourfold fewer entries on fourfold-redundant streams, lower peak memory, and cache growth tracking distinct items rather than token count. Follow with a language-model perplexity test on text containing repeated entities.
Verification
Stage 1 · Toy mechanism gate: Passed ✓
Stage 2 · Mechanism transferred to benchmark: Confirmed ✓
Stage 2 · Practical benchmark result: Inconclusive
Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works
Stage 1 — Mechanism check agent confidence 9/10
Built a readable DP-Means cosine-novelty cache with count-weighted key/value merging and a deterministic associative-recall test. The mechanism predictions were confirmed: the allocation boundary matched acos(1−tau) within 1 degree, orthogonal distinct-item counts matched exactly across thresholds, and merge MSE followed sigma^2/n with normalized factors 0.977–1.036. On the toy repeated-item stream, the cache achieved 4x compression with the same 1.00 recall as token memory; this is a promising mechanism result, not evidence of language-model improvement.
- Agent confidence
- 9/10
- Baseline
- Token memory: 96 stored entries, recall 1.00
- Idea
- DP-Means cache at tau=0.10: 24 slots, 4.0x compression, recall 1.00; allocation-boundary max error 1.0 degree; merge MSE ratio n*MSE/sigma^2=0.977–1.036
Limitations: Only a synthetic associative-recall workload was tested; no trained GRU/SSM, language-model perplexity, softmax cache read, GPU benchmark, noisy-key clustering stress test, or FIFO/LRU eviction comparison was evaluated.
How to run: python3 dp_means_experiment.py
Stage 2 — Benchmark vs. tuned baseline (arXiv:2607.09889)
No significant difference vs. the tuned baseline (+65.1%, p=0.1618, wins 3 / 8 paired seeds).
- Benchmark
- Dynamics forecasting
- Model
- rnn_small
- Paired seeds
- 8
- Baseline mean
- 0.0618
- Idea mean
- 0.1021
- Effect (Δ)
- 0.0402 (+65.1%; negative = idea better)
- Wins
- 3 / 8 paired seeds
- p-value
- 0.1618 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±98.5%
- Mechanism
- Confirmed ✓
- Practical effect
- Inconclusive
- Baseline tuning
- swept over 3 configs
The benchmark used the mandated registered dynamics track but a reduced budget of 2 epochs and 400/100 train/test samples because the initial full protocol exceeded the wall-clock limit. The implementation uses a simple per-example Python cache loop and tests only the recurrent dynamics task; no language-model perplexity, full-length KV-cache memory benchmark, FIFO/LRU comparison, or longer training schedule was tested.
How to run: CUDA_VISIBLE_DEVICES='' PYTHONUNBUFFERED=1 /home/maxwelhelp/main/bin/python3 bench_dp_means_fast.py
Verdict computed by deterministic test code from paired-seed statistics — not by the language model.
Stage 2 — Benchmark latest Failed ✗
Agent confidence: 10/10
Built and evaluated DP-Means distinct-item memory on the structurally matched registered dynamics track with matched rnn_small systems, an equal learning-rate baseline sweep, three novelty thresholds, and 8 paired seeds. The cache mechanism behaved as predicted: mean occupied slots fell from 3.69 to 1.94 as tau increased, versus 8 recurrent tokens, confirming compression. However, best idea MSE was 0.10206 versus baseline 0.06182, delta_mean=+0.04024 with permutation p=0.16175; this is a non-significant degradation, not a win.
- Baseline
- MSE mean 0.0618174, std 0.0361234; best lr=0.01; per-seed [0.0084978, 0.0164876, 0.0829362, 0.0963160, 0.0865394, 0.0737027, 0.1046036, 0.0254561]
- Idea
- Best tau=0.05, lr=0.01; MSE mean 0.102061, std 0.0460542; per-seed [0.133704, 0.130765, 0.111567, 0.0411647, 0.0335320, 0.181023, 0.102962, 0.0817658]; mean occupied slots 3.69 at tau=0.05, 2.74625 at tau=0.10, 1.93875 at tau=0.20
Benchmark result
No significant difference vs. the tuned baseline (+65.1%, p=0.1618, wins 3 / 8 paired seeds).
- Benchmark
- Dynamics forecasting
- Model
- rnn_small
- Paired seeds
- 8
- Baseline mean
- 0.0618
- Idea mean
- 0.1021
- Effect (Δ)
- 0.0402 (+65.1%; negative = idea better)
- Wins
- 3 / 8 paired seeds
- p-value
- 0.1618 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±98.5%
- Mechanism
- Confirmed ✓
- Practical effect
- Inconclusive
- Baseline tuning
- swept over 3 configs
The benchmark used the mandated registered dynamics track but a reduced budget of 2 epochs and 400/100 train/test samples because the initial full protocol exceeded the wall-clock limit. The implementation uses a simple per-example Python cache loop and tests only the recurrent dynamics task; no language-model perplexity, full-length KV-cache memory benchmark, FIFO/LRU comparison, or longer training schedule was tested.
How to run: CUDA_VISIBLE_DEVICES='' PYTHONUNBUFFERED=1 /home/maxwelhelp/main/bin/python3 bench_dp_means_fast.py
Artifacts
- bench_dp_means.py 5.7 KB View Raw
- bench_dp_means_fast.py 3.6 KB View Raw
- bench_report.json 3.5 KB View Raw
- dp_means_experiment.py 6.1 KB View Raw
- report.md 1.3 KB View
- report_bench_2026-08-30T162506.md 5.7 KB View
- results.json 4.2 KB View Raw