△ Mechanism confirmed, baseline not beaten 2026

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

Usefulness8/10
Difficulty4/10
Novelty6/10

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

$$\mathrm{nov}_{t}=1-\max_{i\in\mathrm{slots}}\mathrm{sim}(k_t,\kappa_i),\qquad \mathrm{sim}(a,b)=\frac{a^{\top}b}{\|a\|_2\|b\|_2}.$$
$$\text{if }\mathrm{nov}_{t}>\tau:\ \text{append }(k_t,v_t)\text{ as a new slot};\qquad\text{else }j_t=\arg\max_{i}\mathrm{sim}(k_t,\kappa_i)\text{ and merge into slot }j_t.$$
$$n_{j_t}\leftarrow n_{j_t}+1,\qquad \kappa_{j_t}\leftarrow\frac{(n_{j_t}-1)\kappa_{j_t}+k_t}{n_{j_t}},\qquad u_{j_t}\leftarrow\frac{(n_{j_t}-1)u_{j_t}+v_t}{n_{j_t}}.$$
$$r(q)=\operatorname{softmax}_{i\in\mathrm{slots}}\left(\frac{q^{\top}\kappa_i}{\sqrt d}\right),\qquad \operatorname{read}(q)=\sum_{i\in\mathrm{slots}}r_i(q)u_i.$$

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

Mechanism confirmed, baseline not beaten

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
Limitations:

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
Limitations:

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

Implementation overview ⬇ Download all as ZIP 7 files · code, reports and structured results