Memory Caching: RNNs with Growing Memory - Bridging the Gap to Transformers

TL;DR

The dominance of Transformers is largely due to their "growing memory" (KV-cache), which scales with context but at a quadratic cost. Memory Caching (MC) is a breakthrough framework that brings this "growing memory" capability to Recurrent Neural Networks (RNNs). By caching periodic snapshots of the hidden state, MC allows RNNs to look back at their own history, enabling them to match Transformer performance on recall tasks while staying efficient at $O(NL)$ complexity.

Problem & Motivation: The Fixed-Memory Bottleneck

In the world of sequence modeling, we face a fundamental trade-off:

• Transformers: Excellent at retrieval because they store everything (growing memory), but they suffer from $O(L^2)$ complexity.
• RNNs/SSMs: Highly efficient ($O(L)$) because they compress the past into a fixed-size state, but they suffer from "forgetting" when the sequence outgrows that state's capacity.

The authors identify that the fixed-memory state is the primary bottleneck preventing recurrent models from beating Transformers in "recall-intensive" tasks (like finding a specific fact in a 100k-token document). Their insight is simple: Why not cache the RNN's internal memory state at different time intervals?

Methodology: The MC Framework

The core idea of Memory Caching is to split an input sequence into segments and store the "compressed" memory state of each segment. When the model processes a new token, it doesn't just look at its current (online) memory; it also queries the cached memories of the past.

The Four Flavors of Caching

1. Residual Memory: A simple summation of past states.
2. Gated Residual Memory (GRM): Uses an input-dependent gate ($\gamma$) to decide which past segments are relevant to the current query.
3. Memory Soup: Instead of averaging outputs, it averages the parameters of the memory modules (specifically for non-linear/deep memories).
4. Sparse Selective Caching (SSC): A "Mixture-of-Experts" approach where a router selects only the Top-K most relevant cached states, drastically reducing computation.

Figure 1: The MC framework. Each token attends to its current state and a set of cached historical states.

Why it Works: The Optimization Perspective

The paper provides a beautiful theoretical intuition: Recurrence is essentially an online optimization process. The hidden state is a "model" being trained on the fly to memorize the sequence. Caching hidden states is equivalent to saving checkpoints of this optimization. By querying these checkpoints, the model can "re-read" the past without the quadratic cost of full attention.

Figure 2: Sparse Selective Caching (SSC) selects contextually relevant segments, making the memory "growing" yet "sparse".

Experiments & Results

The authors tested MC on several modern backbones, including Titans, DLA, and Linear Attention.

Key Findings:

• Needle-In-A-Haystack: In 16K context retrieval, DLA+GRM achieved significant improvements over standard DLA, nearly reaching Transformer-level accuracy.
• Throughput: SSC (Sparse Selective Caching) provides the "best of both worlds"—it maintains high throughput similar to RNNs while providing the recall capabilities of Transformers.
• SOTA Comparison: Titans + MC outperformed pure Titans by nearly 1% on average across common-sense reasoning and language modeling tasks.

Table 1: MC variants consistently elevate the performance of base recurrent models across the board.

Critical Analysis & Conclusion

Memory Caching is a powerful "plug-and-play" technique. Its greatest value lies in the SSC variant, which effectively creates a Sparse Unified Memory.

Limitations: While MC reduces the complexity from $O(L^2)$ to $O(NL)$, it still requires managing a cache of states. For extremely long sequences (e.g., millions of tokens), the overhead of storing and routing to thousands of segment checkpoints remains a challenge for GPU memory.

Future Outlook: This work signals a shift in architectural design. We are moving away from the "Attention vs. RNN" binary and toward Hybrid Memory Systems—where models use local recurrence for speed and sparse checkpointing for long-term retrieval.