Memory Caching: RNNs with Growing Memory

Abstract: Transformers have been established as the de-facto backbones for most recent advances in sequence modeling, mainly due to their growing memory capacity that scales with the context length. While plausible for retrieval tasks, it causes quadratic complexity and so has motivated recent studies to explore viable subquadratic recurrent alternatives. Despite showing promising preliminary results in diverse domains, such recurrent architectures underperform Transformers in recall-intensive tasks, often attributed to their fixed-size memory. In this paper, we introduce Memory Caching (MC), a simple yet effective technique that enhances recurrent models by caching checkpoints of their memory states (a.k.a. hidden states). Memory Caching allows the effective memory capacity of RNNs to grow with sequence length, offering a flexible trade-off that interpolates between the fixed memory (i.e., $O(L)$ complexity) of RNNs and the growing memory (i.e., $O(L^2)$ complexity) of Transformers. We propose four variants of MC, including gated aggregation and sparse selective mechanisms, and discuss their implications on both linear and deep memory modules. Our experimental results on language modeling, and long-context understanding tasks show that MC enhances the performance of recurrent models, supporting its effectiveness. The results of in-context recall tasks indicate that while Transformers achieve the best accuracy, our MC variants show competitive performance, close the gap with Transformers, and performs better than state-of-the-art recurrent models.

• The paper introduces Memory Caching, a mechanism that augments RNNs with scalable, tunable memory capacity for improved long-context recall.
• It employs segmentation and various aggregation methods like Gated Residual Memory and Sparse Selective Caching to balance computational efficiency and memory usage.
• Empirical results show enhanced recall and common-sense reasoning, significantly reducing the memory overhead compared to Transformer models.

Memory Caching: Towards RNNs with Adaptive, Growing Memory

Motivation and Context

Modern sequence modeling is dominated by Transformer-based architectures due to their dynamic, context-length dependent memory via self-attention. However, attention’s $\mathcal{O}(L^2)$ complexity with sequence length $L$ yields quadratic computational cost and exorbitant memory overhead for long contexts, motivating investigation into sub-quadratic alternatives—particularly recurrent neural architectures. RNNs possess constant memory complexity but suffer from catastrophic forgetting under long contexts given the fixed-size memory constraint, resulting in substantial performance degradation in recall-intensive settings.

The paper "Memory Caching: RNNs with Growing Memory" (2602.24281) introduces Memory Caching (MC), a mechanism for strong, tunable, and scalable memory capacity in RNNs. MC interpolates between fixed-size RNN memory ($\mathcal{O}(L)$) and unbounded Transformer memory ($\mathcal{O}(L^2)$), explicitly controlling the computation-memory trade-off while retaining recurrent inductive biases. This bridges the representational gap between memory-efficient recurrence and high-capacity non-local attention, enabling robust long-context modeling and in-context recall.

Figure 1: The Overall Memory Caching Method: Each token attends to both its online memory and a set of cached memories from prior segments, effectively increasing context capacity.

The Memory Caching Framework

Segmentation and Memory Growth

MC divides the input sequence into variable-length segments. At the end of each segment, the current recurrent memory state is checkpointed and cached. During inference, queries at each time step attend both to the online (current) memory and all (or a subset of) previously cached memory checkpoints. By adjusting segment length and the aggregation/retrieval scheme, the system’s effective memory and computational complexity is made tunable:

• Segmenting at every token ($N=L$) recovers attention (full memory, high cost).
• Segmenting at maximal length ($N=1$) recovers classic RNNs (minimal memory, low cost).
• Intermediate segmentation yields subquadratic complexity ($\mathcal{O}(NL)$).

This adaptive mechanism enables RNNs to expand their effective context capacity without incurring prohibitive resource demands.

Memory Aggregation Mechanisms

The paper formalizes several aggregation strategies for combining cached memories:

• Residual Memory: Simple summation of all cached and online memories.
• Gated Residual Memory (GRM): Input- and segment-dependent gates (computed via contextual similarity between queries and segment summaries) modulate contributions from each cached state, increasing selectivity.
• Memory Soup: Inspired by weight interpolation, cached parameter sets are data-dependently interpolated to create time-dependent retrieval modules in non-linear settings.
• Sparse Selective Caching (SSC): Employs a contextual router (akin to a Mixture-of-Experts) to select the $k$ most relevant cached memories for each query, enabling controlled sparsity and efficient memory utilization.

  Figure 2: Sparse Selective Caching (SSC) routes token queries to a subset of the most contextually relevant cached memory states for increased efficiency and sparsity.

Tradeoffs via Segmentation

The segmentation schedule controls the granularity of compression and computation:

• Constant-length segments provide a straightforward interpolation between quadratic (attention-like) and linear (RNN-like) complexity.
• Logarithmic segmentation yields log-linear complexity but can compromise recall if early context is excessively compressed.

  Figure 3: Impact of constant vs. logarithmic segmentation on memory capacity and context recall; logarithmic schemes can result in suboptimal compression or overflow.

Empirical Results

The authors instantiate MC atop several strong RNN variants, including Sliding Window Linear Attention (SWLA), Deep Linear Attention (DLA), and Titans (momentum-augmented RNN), evaluating across language modeling, retrieval (Needle-In-A-Haystack), LongBench, and in-context recall tasks.

Key findings include:

• Consistent performance improvements: MC-enhanced RNNs (with GRM/SSC/soup) close much of the recall and common-sense reasoning gap to full attention models, outperforming other state-of-the-art recurrent and hybrid baselines over a range of tasks and scales.
• Robust long-context retrieval: On Needle-In-A-Haystack and in-context recall benchmarks, MC models approach or surpass Log-Linear and attention hybrids, outperforming fixed-memory RNNs especially at large context lengths.
• Superior efficiency: SSC delivers the most favorable tradeoff between throughput and recall, retaining efficiency gains over Transformers as sequence length increases.

Figure 4: Training throughput comparison: MC variants (especially SSC) provide substantial efficiency gains versus quadratic-cost attention.

Practical and Theoretical Implications

MC provides a unified design axis for recurrent sequence models to dynamically trade off memory, throughput, and recall—precisely controlling the inductive bias between global context capacity and computational tractability. This enables:

• Scalable training and inference for long-context LMs on resource-constrained hardware.
• Competitive in-context and associative recall without explicit token memorization.
• Efficient deployment in streaming, low-latency, and edge inference scenarios where fixed memory budget is required.

Theoretically, the results highlight that, under suitable aggregation and routing, RNNs can match much of self-attention’s recall/long-range capacity, challenging the notion that quadratic context scaling is a necessary condition for high-performance autoregressive modeling. Further, the modularity of MC supports composability and integration with advanced associative memory or meta-learning objectives.

Forward-Looking Directions

Potential research avenues sparked by this work include:

• Learned adaptive segmentation and dynamic sparsity routing for further performance-optimized memory scaling.
• Composition with retrieval-augmented architectures and external memory modules for hybrid context reasoning.
• Investigation of MC in the context of meta-learning and data-dependent continual/lifelong learning.

Conclusion

The Memory Caching framework constitutes a principled mechanism for augmenting RNNs with scalable, context-dependent memory via checkpointed state reuse and selective aggregation. This approach fundamentally alters the computational context-recall tradeoff, equipping recurrent models with flexible, competitive long-context capabilities at controllable computational cost and memory usage—crucially expanding their applicability for modern, resource-constrained sequence processing (2602.24281).