Source-linked AI summary
Memory Caching: RNNs with Growing Memory
Ali Behrouz, Zeman Li, Yuan Deng, Peilin Zhong, Meisam Razaviyayn, Vahab Mirrokni
TL;DR
Transformers offer growing memory but incur quadratic complexity, while recurrent models’ fixed memory limits recall. Memory Caching caches memory checkpoints to provide a controllable capacity–complexity trade-off, and experiments report improved recurrent-model performance across language modeling, long-context, and retrieval tasks, though linearity and simple retrieval can collapse or simplify designs.
Problem
Recurrent models provide efficient subquadratic alternatives but their fixed-size memory can forget past information in recall-intensive and long-context tasks.
Method
Memory Caching segments sequences, caches recurrent memory checkpoints, and combines them using four aggregation strategies, including gated and sparse selective mechanisms.
Results
Memory Caching enhances recurrent models across language modeling, long-context understanding, and retrieval tasks, extending their effective context length.
Takeaways & Limitations
MC provides a controllable middle ground between fixed-memory recurrence and growing-memory attention, with complexity determined by the number of cached memories.
Takeaways & Limitations
Linearity or simple multiplication-based retrieval can simplify or collapse MC designs and may produce sub-optimal performance, especially across memory architectures.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Transformers provide growing associative memory but incur quadratic cost, while recurrent models use efficient fixed-size memory that can forget information in long or recall-intensive sequences. Memory Caching addresses this trade-off by caching recurrent memory checkpoints and combining them through several aggregation strategies.
- Transformers support growing-capacity associative memory for retrieval but require quadratic computation and high inference-time memory.
- Recurrent models maintain fixed-size memory efficiently, but compression can force them to forget past information in recall-intensive and long-context tasks.
- Memory Caching segments sequences and caches compressed memory states so recurrent models can directly access information from their history.
- MC offers controllable complexity of O(NL), interpolating between O(L) recurrence and O(L^2) Transformer attention.
- Four aggregation strategies comprise gated and residual memory, Memory Soup, and Sparse Selective Caching for context-aware or efficient memory retrieval.
- Experiments across linear attention, deep memory, sliding-window, and deep linear architectures show improved language-modeling, long-context, and retrieval performance.
2 PRELIMINARIES AND BACKGROUND
The paper frames attention and recurrent sequence models as memory mechanisms, then motivates linear attention as an efficient recurrent alternative and Memory Caching as checkpointed optimization states. These preliminaries establish the notation and conceptual basis for the proposed method.
- Attention: Attention acts as an associative memory over input-dependent keys, values, and queries, but requires O(L^2) operations to access all past tokens.
- Linear Attention: Linear attention replaces the exponential operator with a separable kernel, yielding an efficient recurrent formulation.
- Linear Attention: The linear-attention state Mt = Mt−1 + vtϕ(kt)⊤ functions as a fixed-size memory.
- Test-time Memorization and Nested Learning Perspective: The nested-learning perspective interprets attention and modern RNN updates as dynamic in-context memorization with different internal objectives.
- Test-time Memorization and Nested Learning Perspective: Memory Caching uses cached states as checkpoints of the memory optimization process, allowing past states to enhance subsequent processing.
3 RECURRENT NEURAL NETWORKS WITH MEMORY CACHING
Memory Caching augments recurrent models by storing segment-level memory checkpoints and retrieving information from current and cached memories. Its variants trade computational cost against effective memory capacity, including gated, aggregated, and sparse selective retrieval.
- Memory Caching framework: Memory Caching segments the sequence and caches each segment’s final memory state, allowing recurrent models to access compressed information from the history.The current online memory is used together with cached memories when computing each query’s output.
- Gated Residual Memory: Gated Residual Memory uses input-dependent gates to modulate each segment’s contribution, addressing residual aggregation’s equal treatment of cached memories.The gates can depend on the input and the segment context, with similarity-based parameterization using pooled segment representations.
- Memory complexity: Retrieval uses the current memory and cached memories from previous segments, while memory updates remain recurrent.The output computation performs forward passes over both online and cached memories.
- Memory complexity: MC retrieval has O(NL) complexity, interpolating between O(L) recurrence when N=1 and attention-like retrieval when N=L.N is the number of cached memories, determined by the segmenting scheme.
- Memory Soup: Memory Soup combines cached memory states by aggregating their parameters into a data-dependent memory for retrieval.The cached memories share the same architecture and parameter count, enabling parameter-based combination.
- Sparse Selective Caching: Sparse Selective Caching uses a router to select contextually relevant cached memories, reducing overhead for ultra-long sequences and supporting adaptive retrieval.SSC is described as a sparse unified memory in which tokens activate subsets of memory parameters for writing and retrieval.
4 DISCUSSION AND PROOF OF CONCEPT
Memory Caching applies to recurrent updates by segmenting sequences, caching compressed memory states, and aggregating them during retrieval. Its segmentation controls the trade-off between compression, recall resolution, and computational cost, while extensions to linear and deep memories connect caching to attention-like mechanisms and expose design limitations.
- Linear and deep memory: For linear memory, segment size one stores each token as a memory state and is equivalent to caching memory checkpoints rather than using independent memories.The same viewpoint motivates interpreting attention as caching past inputs and extending the design with more expressive memory updates.
- Design limitations: Linearity or multiplicative retrieval can collapse memory-fusion designs or produce sub-optimal performance, motivating richer retrieval and memory formulations.The paper notes that some linear-memory configurations simplify into gated residual memory or compressor modules, whereas deep-memory variants remain distinct.
- Segmentation trade-offs: Logarithmic segmentation can use long subsequences that risk memory overflow or short subsequences that prevent effective inner-loop optimization.Figure 3 contrasts constant and logarithmic segment sizes as illustrative examples of this design tension.
- Linear and deep memory: For deep memory, caching each token creates architectures distinct from hybrid variants, with token representations that can vary by query and neural-network bias terms.The cached memory is a 2-layer MLP, so retrieval can produce query-dependent representations rather than constant token vectors.
- Memory caching framework: Memory Caching segments sequences and caches compressed memory states, allowing recurrent models to access compressed information from the entire history.The technique is applicable to arbitrary recurrent update rules and is evaluated with Linear Attention, Titans, SWLA, and DLA.
- Complexity: Memory caching has total complexity O(L + p×N×L), where N is the number of segments and p is the memory forward-pass cost.Memory updates cost O(L), while retrieval passes over all cached memories and costs O(p×N) per token.
- Segmentation trade-offs: Segment length trades recall resolution against computation: constant-size and logarithmic segmentation cost O(L^2) and O(L log(L)), respectively.Logarithmic segmentation is more efficient but provides less resolution for long-past tokens in recall-intensive tasks.
5 EXPERIMENTS
Experiments evaluate Memory Caching across language modeling, retrieval, long-context understanding, recall, ablations, and throughput. Across these settings, MC generally improves recurrent baselines, with competitive recall and efficiency trade-offs.
- MC variants consistently improve DLA, Titans, and SWLA across downstream tasks and their average baseline performance.
- Titans + MC and DLA + MC achieve a +0.8% performance gain over Titans.
- MC-enhanced DLA and Titans consistently outperform base models on Needle-in-a-Haystack tasks and surpass Log-Linear especially at longer contexts.Log-Linear struggles when a single memory must compress very large initial segments, whereas MC distributes compression across segments.
- Transformers achieve the best in-context recall results, while MC variants remain competitive and outperform state-of-the-art recurrent models.
- All MC-enhanced variants improve performance over their base RNNs on LongBench tasks.
- MQAR results show MC models outperforming base RNNs and state-of-the-art recurrent models, with best performance per dimension value against models such as Atlas.
- Ablations show that context-aware gating, removing gating into residual memory, and using linear memory modules can each preserve or enhance performance.The study reports greater robustness to memory architecture and expressivity when memory caching is used.
- MC variants provide a throughput middle ground between Transformers and RNNs and become more efficient than Transformers as context length increases.SSC offers minimal overhead relative to its base RNN while performing comparably or better than other variants across discussed downstream tasks.
6 CONCLUSION
The conclusion presents Memory Caching as a simple recurrent-model technique that caches memory-state subsets so later tokens can attend to relevant past information.
- Memory Caching caches a subset of recurrent memory states, allowing subsequent tokens to attend directly to relevant past tokens.
- The paper reports improvements over a subset of baselines while emphasizing simplicity to isolate the effect of memory caching.
- Future work could use more expressive pooling or routing mechanisms to further enhance performance.
A RELATED WORK
Related work covers efficient recurrent and attention-based memory architectures, fast-weight and meta-learning perspectives, associative-memory foundations, and structured efficiency mechanisms.
- Linear attention replaces softmax with separable kernels to alleviate Transformers’ quadratic complexity and context-length limitations.
- Deep memory modules increase capacity or alter inner-loop objectives and optimization rules, including Titans, TTT layers, and higher-order attention variants.
- Fast-weight programmers integrate dynamic writable memory stores into recurrent neural networks, with Hebbian learning and delta rules as prominent paradigms.
- Hopfield Networks provide an early associative-memory formulation based on minimizing an energy function to store key-value pairs.
- Efficient attention research uses structured matrices and sparse or hybrid mechanisms to reduce the computational burden of token and channel mixing.
B EXPERIMENTAL DETAILS
The experimental details list language-modeling and evaluation datasets used to assess recurrent models and Memory Caching.
- Experiments use Wikitext, LMB, PIQA, HellaSwag, WinoGrande, ARC-easy, ARC-challenge, SIQA, and BoolQ.Training uses a 32K vocabulary.