Memory Caching: Giving RNNs a Growing Memory Without the Quadratic Cost

Paper: 2602.24281 Authors: Ali Behrouz, Zeman Li, Yuan Deng, Peilin Zhong, Meisam Razaviyayn, Vahab Mirrokni Categories: cs.LG, cs.AI

The Gap

Transformers dominate sequence modeling because they remember everything—their attention mechanism can look back at any previous token. This “growing memory” makes them excellent at recall-intensive tasks like retrieving facts from long documents. But it’s expensive: processing a sequence of length $L$ costs $O(L^2)$ in both time and memory.

Recent recurrent alternatives (Mamba, RWKV, RetNet) promise subquadratic complexity by compressing history into a fixed-size hidden state. They’re fast—$O(L)$ complexity—but they stumble on tasks requiring precise recall of distant information. The culprit? Their memory doesn’t grow. Once you compress 10,000 tokens into a 512-dimensional vector, you’ve lost details. State-of-the-art recurrent models consistently underperform Transformers on needle-in-haystack retrieval and long-context question answering, often by 10-20 percentage points.

The Increment

Before: RNNs choose between fixed memory (fast but forgetful) or full attention (accurate but quadratic). After: Memory Caching lets RNNs checkpoint their hidden states periodically, creating a sparse growing memory that interpolates smoothly between $O(L)$ and $O(L^2)$ complexity.

Think of Memory Caching like a library with reading rooms. A standard RNN is a single reader who takes notes (the hidden state) as they go through a book sequentially—they can only remember what fits in their notebook. A Transformer is a reading room where every page of every book stays open on tables—you can glance at any page instantly, but you need a massive room. Memory Caching is a reader who photocopies key pages at regular intervals and pins them to a wall. When they need to recall something, they first check their current notes (the hidden state), then scan the pinned pages (cached states), and finally re-read from the original if needed. The wall grows with the book’s length, but much slower than keeping every page open.

Mechanically, every $k$ tokens, the RNN saves a snapshot of its hidden state. When processing later tokens, the model can attend back to these cached states. Four variants control how: Uniform caches at fixed intervals; Gated learns to weight cached states; Sparse selectively attends to a subset; Hierarchical builds a tree of caches at multiple resolutions. The key insight is that you don’t need to remember every token—just enough checkpoints to reconstruct what matters. If you cache every 64 tokens in a 4096-token sequence, you store 64 states instead of 4096, but can still recover information by attending to the nearest cache and processing forward from there.

Key Concepts

Hidden State Compression: When an RNN processes a sequence, it maintains a hidden state vector (typically 512-2048 dimensions) that’s supposed to summarize everything seen so far. Imagine you’re listening to a story and can only remember 10 facts at a time—you constantly update which 10 facts matter most. Early tokens get “compressed out” as new information arrives. For the sentence “The cat that the dog that the mouse scared chased ran,” by the time you reach “ran,” a standard RNN has largely forgotten “cat” because it’s been overwritten by “dog,” “mouse,” “scared,” and “chased.” This is why RNNs struggle with long-range dependencies—not because they can’t theoretically represent them, but because the fixed-size bottleneck forces lossy compression.

Attention as Lookup: Transformers don’t compress—they keep every token’s representation and use attention to look up relevant ones. Attention is essentially a soft database query: given a query vector (current token), compute similarity scores with all key vectors (past tokens), then retrieve a weighted combination of value vectors. The cost is that comparing against $L$ past tokens takes $O(L)$ operations per token, hence $O(L^2)$ total. Memory Caching reduces this by only keeping every $k$-th token’s representation, so you compare against $L/k$ cached states—still attention, but over a compressed index.

The Interpolation Trade-off: Memory Caching’s elegance is in its parameterization. Set cache interval $k=1$, and you recover full Transformer attention ($O(L^2)$). Set $k=\infty$, and you get a pure RNN ($O(L)$). Any $k$ in between gives you a point on the Pareto frontier between speed and recall capacity. If you cache every 32 tokens, you’re spending $\sim 3\%$ of a Transformer’s compute (since $32^2 / 1024^2 \approx 0.001$ for a 1024-token context, but you still process the full sequence recurrently). This lets practitioners dial in their desired speed/accuracy trade-off based on deployment constraints—something neither pure RNNs nor pure Transformers offer.

Expert Assessment

Problem significance: This addresses a real bottleneck. The ML community is actively searching for Transformer alternatives that scale to million-token contexts (think: entire codebases, long videos, scientific papers). Current recurrent models are 2-5x faster but lag 5-15% in accuracy on long-context benchmarks. Closing that gap matters for production systems where both speed and quality are non-negotiable. The affected community is large—anyone doing document understanding, code analysis, or long-form generation.

Method maturity: This is a solid proof-of-concept with room to grow. The core idea is simple (almost obvious in hindsight), which is a strength—it’s easy to implement and integrate into existing recurrent architectures. However, the paper doesn’t deeply explore hyperparameter sensitivity (how does performance degrade as $k$ increases?) or provide theoretical analysis of what information is preserved vs. lost. The gated and sparse variants add learnable components, but it’s unclear if the gains justify the added complexity. The hierarchical variant is mentioned but not thoroughly evaluated. Deployment-wise, caching adds memory overhead (storing $L/k$ states) and complicates batching—these practical concerns aren’t addressed.

Experimental rigor: The experiments are reasonable but not exhaustive. Language modeling and long-context QA are appropriate testbeds. The in-context recall tasks (needle-in-haystack) are particularly telling—MC variants close the gap with Transformers from 20-30% to 5-10%, which is meaningful. However, baselines are limited to a few recent recurrent models (Mamba, RWKV); comparisons with hybrid architectures (e.g., Transformer-XL, Compressive Transformers) are missing. The datasets are standard but not adversarial—it would be interesting to see performance on tasks specifically designed to break recurrent models (e.g., counting, parity checks). One red flag: the paper doesn’t report wall-clock time or memory usage, only complexity analysis. Real-world speedups depend on implementation details and hardware.

Verdict: Weak accept — The idea is sound and the results are promising, but the paper feels like an early-stage exploration rather than a definitive solution; deeper analysis and broader comparisons would strengthen the contribution.

Takeaways

Checkpointing as a universal pattern: The core insight—periodically saving state snapshots to enable non-local lookups—applies beyond RNNs. Gradient checkpointing in backpropagation, key-value caches in autoregressive generation, and even version control systems (Git) use the same principle: trade storage for the ability to “rewind” without recomputing everything. If you’re designing any stateful system with sequential dependencies, ask: “What if I cached intermediate states?”

Pareto frontiers over binary choices: Instead of debating “RNN vs. Transformer,” Memory Caching offers a spectrum. This framing is valuable in other domains—rather than choosing between two extremes (batch vs. online learning, centralized vs. federated training), look for parameterized interpolations that let users navigate the trade-off space. The best solution is often not at the endpoints.

Sparse attention is underexplored: The sparse variant (selectively attending to a subset of caches) hints at a broader opportunity. Most attention mechanisms are dense—every query attends to every key. But in long contexts, most attention weights are near-zero. Learning which caches to attend to, or dynamically adjusting cache granularity based on content, could yield further gains. This connects to recent work on learned sparsity (e.g., Mixture of Experts, dynamic routing) and suggests that “smart indexing” is a fruitful research direction.

Implementation simplicity matters: Memory Caching’s appeal is partly that it’s a drop-in modification—you don’t need to redesign the architecture, just add a caching mechanism. When proposing new techniques, consider the “integration tax.” A 10% improvement that requires rewriting your entire training pipeline is less valuable than a 5% improvement that’s five lines of code. Practitioners will adopt the latter.