MoM: Mixture-of-Memories Fixes Linear Attention Recall

Jusen Du, Weigao Sun, Disen Lan, Jiaxi Hu, Yu Cheng — Shanghai AI Laboratory, Tsinghua University, Fudan University, HKUST(GZ), CUHK
arXiv:2502.13685 Feb 2025 Linear Sequence Modeling Mixture-of-Experts Long-Context Recall

A neuroscience-inspired fix for the recall gap in efficient Transformer alternatives: instead of one fixed-size memory state, MoM routes tokens across several independent memory slots plus a shared accumulating memory — borrowing Mixture-of-Experts routing and a hippocampus-style multiplexing analogy.

Three ways to remember a sequence

Click a column to isolate it — memory strategy determines the cost/recall trade-off.
Click "Transformer", "Linear (single state)", or "MoM" above the diagram to compare.
Transformer — KV cache (O(n²), nothing forgotten) Linear model — single fixed state (O(n), lossy) MoM — routed multi-memory + shared state

The core trade-off, at a glance

O(n²)
Transformer attention cost — perfect recall, doesn't scale
O(n)
Linear models — constant memory, capacity + interference limited
4–8
Independent memory slots MoM routes tokens across
+1.7
Routing-attributable recall gain at matched params (pts)

Router mechanics, step by step

One matmul, one softmax, one top-k, one renormalization — that's the whole router.

Capacity vs. interference: why one memory isn't enough

Toggle to see how a single fixed-size state blends unrelated writes, versus routed slots keeping them apart.

Solving scatter-gather without new kernels

Routing tokens to different memories looks like a GPU-throughput killer. MoM turns it into a sequence-reordering problem instead.

Relative kernel efficiency (illustrative)

Reusing existing Triton linear-attention kernels on reordered, contiguous segments avoids writing new GPU kernels from scratch.

Recall benchmark average (FDA · SWDE · SQuAD · NQ · TriviaQA · Drop, 2K tokens)

Six linear-model baselines vs. MoM, at two parameter scales.

LongBench average

Summarization, few-shot, synthetic, code.

Perplexity: 2K-trained, extrapolated to 32K

Fineweb dataset. Lower is better — flatter curve means more graceful long-range degradation.

Memory count scaling

Fixed 0.5 activation ratio, scaling from 1 to 8 memory slots.

Headline gap vs. controlled comparison

Part of MoM's lead over weak baselines is just "built on the strongest update rule available." Toggle to isolate the routing-only effect.

Same activated capacity, one state vs. several

Expanding a single memory's dimensionality to match MoM's total capacity still loses to separated, routed slots.

Emergent memory specialization

Routing real ARC-easy tokens through a trained model: each memory leans toward a different token category.

What the evidence actually covers

All six recall benchmarks are truncated to 2K tokens — a window a single well-tuned memory state can often handle. The only test past 2K is perplexity extrapolation to 32K, which shows the model doesn't collapse, but never proves task-level retrieval of a fact planted far back in context. Titans, a peer approach to the same interference problem, is listed in the paper's taxonomy but never run as a baseline.