Source-linked AI summary
It's All Connected: A Journey Through Test-Time Memorization, Attentional Bias, Retention, and Online Optimization
Ali Behrouz, Meisam Razaviyayn, Peilin Zhong, Vahab Mirrokni
TL;DR
Efficient sequence modeling lacks a broad framework that unifies associative-memory mechanisms and architectural choices across modern models. The paper introduces Miras, which combines memory architecture, attentional bias, retention gate, and learning algorithm choices, and evaluates three resulting models. Across reported tasks, Miras variants show stronger performance than Transformers and linear RNNs, while existing recurrent approaches retain memory-capacity and expressive-power limitations.
Problem
Existing sequence-model analyses lack a universal, general framework covering associative-memory objectives, learning algorithms, and forget gates across architectures.
Method
Miras designs sequence models by jointly choosing an associative-memory architecture, attentional-bias objective, retention gate, and memory-learning algorithm.
Results
Miras variants including Moneta, Yaad, and Memora outperform state-of-the-art sequence models across language modeling, commonsense reasoning, needle-in-a-haystack, and recall-intensive evaluations.
Takeaways & Limitations
The framework connects online optimization with test-time memorization and provides design choices for sequence models with different strengths across downstream tasks.
Takeaways & Limitations
Existing linear recurrent memories can overflow or have limited expressive power, while prior associative-memory accounts lack generality and do not explain forget gates.
Abstract
from arXiv · showhide
Designing efficient and effective architectural backbones has been in the core of research efforts to enhance the capability of foundation models. Inspired by the human cognitive phenomenon of attentional bias-the natural tendency to prioritize certain events or stimuli-we reconceptualize neural architectures, including Transformers, Titans, and modern linear recurrent neural networks as associative memory modules that learn a mapping of keys and values using an internal objective, referred to as attentional bias. Surprisingly, we observed that most existing sequence models leverage either (1) dot-product similarity, or (2) L2 regression objectives as their attentional bias. Going beyond these objectives, we present a set of alternative attentional bias configurations along with their effective approximations to stabilize their training procedure. We then reinterpret forgetting mechanisms in modern deep learning architectures as a form of retention regularization, providing a novel set of forget gates for sequence models. Building upon these insights, we present Miras, a general framework to design deep learning architectures based on four choices of: (i) associative memory architecture, (ii) attentional bias objective, (iii) retention gate, and (iv) memory learning algorithm. We present three novel sequence models-Moneta, Yaad, and Memora-that go beyond the power of existing linear RNNs while maintaining a fast parallelizable training process. Our experiments show different design choices in Miras yield models with varying strengths. For example, certain instances of Miras achieve exceptional performance in special tasks such as language modeling, commonsense reasoning, and recall intensive tasks, even outperforming Transformers and other modern linear recurrent models.
1 Introduction
The paper frames efficient sequence modeling as a search for a general design framework that explains and improves associative-memory-based architectures. Miras organizes this search around four architectural choices and motivates three novel variants evaluated across several sequence-modeling tasks.
- Transformers provide strong sequence modeling but their quadratic time and space complexity limits long-context applications.
- Existing associative-memory accounts of Transformers and recurrent models lack a universal explanation, broad generalizability, or an account of forget gates.
- Miras reconceptualizes sequence architectures as associative memories whose attentional bias learns mappings between inputs, or keys, and values.
- Miras designs architectures through four choices: attentional bias, retention gate, memory architecture, and memory learning algorithm.
- Moneta, Yaad, and Memora use alternative attentional biases and retention mechanisms, and are evaluated on language modeling, commonsense reasoning, needle-in-a-haystack, and recall-intensive tasks.
2 Preliminaries and Background
This section reviews attention, recurrent alternatives, and deep-memory sequence models through the lens of efficiency, memory capacity, and associative learning. It motivates a framework that compares architectures by their memory mechanisms while noting limitations in existing recurrent and deep-memory approaches.
- Attention: Attention acts as the associative memory of Transformers, but its computation requires at least N × d operators to produce outputs.
- (Linear) Recurrent Models: Linear recurrent models enable parallelizable, linear-time training, but additive Hebbian updates can overflow memory and limit expressive power.
- (Linear) Recurrent Models: Matrix-valued memories with Delta learning address some recurrent limitations but still face theoretical constraints and moderate practical performance.
- Miras framework perspective: Miras characterizes recent sequence models as sharing the same attentional bias and regularization, while emphasizing that these choices do not uniquely identify a backbone.
- Deep Memory Module: Titans and Test Time Training: Deep-memory architectures update an MLP memory with gradient descent, yet a framework that unifies existing sequence models and explains forget gates remains unresolved.
3 Associative Memory, Attentional Bias, and Retention
The paper frames sequence models as associative memories that learn key-value mappings through attentional-bias objectives, then interprets online updates as balancing new learning with retention. This viewpoint connects optimization choices, retention mechanisms, and existing architectures under a more general Learning-Retaining formulation.
- Associative Memory and Attentional Bias: Associative memory maps keys to values, while attentional bias is the objective that determines the memory type and which events it prioritizes.The formulation treats learning the mapping as minimizing an objective over the memory.
- Associative Memory and Attentional Bias: At test time, learning the key-value mapping is a meta-learning problem: the memory objective is optimized in an inner loop while network parameters are optimized outside it.This lets the model store incoming data in its parameters during inference.
- Optimization Viewpoints: Online gradient updates learn from each new key-value pair by descending its loss, with the update interpretable through online optimization and surprise-based memorization.The per-step loss can be viewed as measuring how surprising a token is to the current memory.
- Learning-Retaining Viewpoint: The Learning-Retaining viewpoint separates learning new concepts from retention regularization that stabilizes changes and preserves previously learned knowledge.Retention can include local control of deviations from the previous state and global control related to memory size.
- Optimization Viewpoints: Learning-Retaining is more general than FTRL under mild assumptions, because suitable retention and loss choices reproduce the FTRL update.Both viewpoints include online gradient descent as a special case with linear loss approximation and quadratic regularization.
- Retention: The paper calls forgetting mechanisms retention gates because memories may remain unerased while becoming inaccessible or not retained by the model.This terminology links model retention behavior to the paper’s interpretation of human memory.
4 Miras: Learning to Memorize with Robust and Expressive Memory
Miras designs sequence-model backbones by independently choosing memory structure, attentional bias, retention, and memory algorithm. It unifies existing architectures by expressing them as combinations of these choices, including Hebbian, delta-rule, nonlinear-regression, and attention-based instances.
- Miras Framework: Miras defines four backbone choices: memory structure, attentional bias, memory stability and retention, and memory algorithm.These choices govern the memory architecture, objective, retention regularization, and optimization procedure.
- Miras Framework: Memory structure controls the architecture and parameter domain, ranging from vectors and linear functions to MLPs or more complex memories.Restricting parameters to regions such as an L2 ball can avoid infinite values or unstable training.
- Miras Framework: Attentional bias determines how the memory memorizes context, maps inputs, and prioritizes events, while approximations can alter the optimization procedure.The framework permits alternative loss approximations beyond the original objective.
- Miras Framework: Retention regularization balances learning with the past state, and the paper identifies effective retention gates as important for long-context performance.The memory algorithm may use gradient descent, momentum, or non-parametric solutions.
- Existing Architectures: Existing Hebbian-style recurrent models arise as Miras instances using dot-product attentional bias, local retention, and gradient descent.Different retention settings recover Linear Attention, RetNet, Lightening Attention, or Mamba2.
- Existing Architectures: Delta-rule recurrent models are special Miras instances, while Titans-LMM represents a nonlinear MSE attentional bias optimized with momentum.The framework also places Mesa-layer and softmax attention within related objective-and-optimizer choices.
5 Beyond Existing Attentional Biases and Retention Gates
Miras broadens sequence-model design beyond predominantly linear or quadratic attentional biases and retention gates by exposing these choices as modular components. The section develops robust memory objectives, value-shift defenses, constrained memories, and stability-oriented retention mechanisms, then uses them to motivate fast-parallelizable model variants.
- Framework: Miras generalizes prior designs by allowing alternative attentional biases, retention gates, memory structures, and optimization algorithms.The framework treats these as independent design choices rather than fixing them to linear or quadratic forms.
- Alternative Attentional Biases: ℓp-attentional bias extends ℓ2 regression to a family of objectives whose choice can produce different memory properties.The framework considers matrix memories and derives corresponding update rules for different p values.
- Alternative Attentional Biases: Value-less associative memory maps stored keys to only −1 and +1, while sign and absolute-value operations can destabilize backpropagation.The proposed approximation Sign(x) ≈ tanh(αx) addresses the nondifferentiability involved in training.
- Alternative Attentional Biases: Huber-type attentional biases reduce the impact of outliers, with context-dependent variants selecting or mixing ℓ2- and ℓ1-based behavior.The parameter δ_t determines which objective or normalized version is used for an incoming context.
- Alternative Attentional Biases: Robust value-shift memory optimizes worst-case perturbations within an ℓ2 ball, combining standard ℓ2 error with a Δ-scaled error-norm term.Δ controls the trade-off between fitting nominal values and robustness to perturbations.
- Alternative Retention Gates: Miras also explores bounded memories, elastic-net retention, soft-thresholding, and general Lq stability mechanisms before introducing three variants designed for fast parallel training.These mechanisms target numerical stability, hard or soft forgetting, and broader retention-gate choices.
6 Experiments
Miras variants are evaluated across language modeling, commonsense reasoning, scaling, recall, and ablation settings. The results report strong performance across baselines, improved scaling with context and compute, and contributions from the tested design choices.
- 6.1 Language Modeling and Common-sense Reasoning: Miras variants outperform Transformer++, modern linear recurrent models, and hybrid baselines on language modeling and commonsense reasoning tasks.The variants are pure recurrent and attention-free; Memora and Yaad are generally stronger than Moneta, though the best model varies by task and size.
- 6.2 Scaling Pattern: All three Miras variants scale better than state-of-the-art baselines as training context increases from 2K to 32K.The paper attributes this to expressive 2-layer MLP memories and alternative retention gates and attentional biases that manage fixed-size memory capacity.
- 6.2 Scaling Pattern: At nearly the same FLOP budget, all three Miras variants outperform the reported baselines in perplexity.The result is presented as supporting the importance of powerful memory design.
- 6.4 Ablation Study: Different p values produce different memory-module performance, with p = 3 best and p = 4 worst, while context-length scaling remains nearly unchanged.The experiment varies p across {1, 1.5, 2, 2.8, 3, 3.2, 4} and context windows from 2K to 16K.
- 6.4 Ablation Study: Unlike p, q changes the context-length scaling pattern because q determines the retention gate and affects memory management.The q ablation varies q in {2, 3, 4, 5}.
- 6.4 Ablation Study: Removing retention, making δ input-independent, removing loss components or the ℓ1 condition, or replacing the MLP with a linear layer reduces Yaad’s performance.The ablation indicates that each tested architectural choice contributes to performance.
7 Conclusion
The paper presents Miras as a framework connecting online optimization with test-time memorization and uses it to design three sequence models. Its evaluations report stronger performance than Transformers and linear RNNs across various downstream tasks, while leaving exploration of other tasks as future work.
- 7 Conclusion: Miras explains the connection between online optimization and test-time memorization while clarifying standard architectural choices such as forget gates.The framework is used to design architectures intended to manage memory more effectively.
- 7 Conclusion: The three Miras variants have distinct advantages and disadvantages, yet all are reported as more powerful than Transformers and linear RNNs across various downstream tasks.The paper describes the variants as a diverse set of architectures and identifies further exploration on different downstream tasks as future work.
A Additional Related Work
The related work covers efficient linear recurrent models, associative-memory interpretations, test-time training, and efforts to unify sequence architectures. It positions Miras as distinguishing update rules and attentional biases that broader simplifications may conflate.
- Modern Linear RNNs: Modern linear RNNs pursue efficient alternatives to Transformers, using fast inference and training to address quadratic-cost and long-context limitations.Early models use data-independent transition matrices with Hebbian-like updates.
- Modern Linear RNNs: Nonlinear RNNs can be trained faster, but their recurrent structure remains challenging to parallelize at larger scales.
- Fast Weight Programs: Fast weight programs interpret linear layers as key-value associative memories with writable recurrent memory and commonly use Hebbian or delta-rule learning.
- Test Time Training: Test-time training adapts models using local learning on test data and has been applied across downstream tasks including vision and video generation.
- Hopfield Networks: Miras builds on a broad associative-memory view in which architectures learn mappings between keys and values.
- Unifying Frameworks: Miras distinguishes Hebbian-like dot-product updates from delta-rule regression updates, placing models with different attentional biases outside narrower unifying classes.
B Proof of Proposition 3.2
The proposition establishes equivalence between a learning-retaining update and an FTRL update under a Bregman-divergence retention function. The proof matches their parameter sequences by induction from identical initialization.
- Proposition 3.2: Under strict convexity, Bregman-divergence retention, and matching loss terms, the learning-retaining and FTRL update rules are equivalent.
- Proof strategy: The proof defines two parameter sequences generated by the respective viewpoints and reduces equivalence to showing they remain identical from the same initialization.
- Inductive step: The induction step compares the learning-retaining update with the FTRL update after applying the proposition's attentional-bias and retention choices.
- Conclusion: Combining the resulting identities completes the induction and establishes equality of the two update sequences.
C Experimental Setup
The experiments evaluate Miras models across language modeling, commonsense reasoning, and long-context needle-in-a-haystack tasks against recurrent, Transformer, and hybrid baselines.
- Tasks: The evaluation covers language modeling, commonsense reasoning, and long-context needle-in-a-haystack tasks.
- Baselines: The comparisons include state-of-the-art linear recurrent models, Transformers, and hybrid recurrent-attention models.
- Architectural details: Architectural details are presented in Table 5.