It's All Connected: A Journey Through Test-Time Memorization, Attentional Bias, Retention, and Online Optimization
Authors: Ali Behrouz
Meisam Razaviyayn
Peilin Zhong
Vahab Mirrokni
Venue: ICLR 2026 Poster (2026)
Keywords: Test Time Memorization
Online Optimization
Recurrent Neural Networks
Primary area: foundation or frontier models, including LLMs

Abstract: 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 with attentional bias. We define and formalize the concept of attentional bias as the internal memory objective deep learning architectures. We show that existing deep learning architectures leverage the same attentional bias based on $L_2$ loss function. Going beyond $L_2$ loss function, we present a set of alternative attentional bias configurations along with their effective approximations. We then reinterpret forgetting mechanisms in modern deep learning architectures as a form of retention regularization. Building upon these insights, we present Miras, a general framework to design deep learning architectures based on the choice of attentional bias objective, retention gate, associative memory architecture, and memory learning algorithm. Our experiments show different designs yield models with varying strengths. Furthermore, our special instances of Miras achieve exceptional performance in language modeling, commonsense reasoning, recall intensive, and time series tasks, outperforming Transformers and other modern linear recurrent models.

Published as a conference paper at ICLR 2026
IT's ALL CONNECTED: A JOURNEY THROUGH TEST- TIME MEMORIZATION, ATTENTIONAL BIAS, RETEN- TION, AND ONLINE OPTIMIZATION
Ali Behrouz"
Meisam Razaviyayn
Peilin Zhong
Vahab Mirrokni
Google Research
{al ibehrcuz, razavi yayn, mi rrokni mirrokni]egoog /egoog com
peilin. wwongl..o@eeeeoo compoiia.mmmm.....m
ABSTRACT
This paper introduces MIRAS, unified framework that reconceives neural archi- tectures (such as Transformers and modern linear RNNs) as associative- memooy modules governed online optimization. In MIRAS, each module learns key- value mappings via an attentional bias (the internal learning objective) and pre- serves past information via retention function (the memory regularizer). This perspective provides a principled reinterpretation of "forgetting" mechanisms as forms of regularization. Our framework reveals critical limitation: virtually all existing sequence models, including recent unification efforts, are constrained by dot-pree dot-product similarity or lg loss. MIRAS moves beyond this narrow focus, pro- viding a generative framework that unlocks a richer design space informed by principles from (robust) optimization and statistics. We introduce diverse alter- natives-sun as Ep norms, Huber loss, KL--based losses, and f-divergence mea- sures- leading to novel architectural designs with improved stability and robust- ness. Utilizing this expanded space, we build three novel, attention-free and par- allelizable models (MONETA, MEMORA, YAAD) that combine expressive MLP memories with these mechanisms. Empirically, different MIRAS instanti- ations trade off complementary strengths, illustrating the framework's 's capacity navigate architectural design choices. Several variants achieve strong scaling, larger effective context windows, and demonstrate results better than state-of-the- art linear RNNs across various tasks, including language modeling. commonsense reasoning, and challenging long-context recall.
INTRODUCTION
Designing efficient architectural backbones for sequence modeling vital for strengthening foun- dation models across diverse domains and data modalities such as language (Vaswani et al. Gemma-Team et al. 2024), computer vision (Dosovitskiy et al., 2020), computational bi- ology (Wang al.] 2024). and neuroscience (Behrouz al. 2024). Transformers (Vaswani al.
2017
2017
remain state of the art thanks to their in-context learning and scalability (Kaplan et al. 2020))
but their quadratic time and space complexity limits use in long-context applications (Dalal al.
2025
Liu et al., 2024b Li et al. 2024).
Recent work tackles Transformers' long contectntett limits by creating efficient recurrent alterna- tives (Schlag et al., 2021 [mith al., 2022). Unlike the Transformer's linearly growing KV cache, these models compress context into fixed-size memory, requiring better memory management for strong performance. To design more effective architectures, studies improves memory capabilities through: (1) richer learning rules, from Hebbian (Hebb) 2005) to Delta (Schlag et al. 2022 and Omega (Behrouz et al., 2025a); ; (2) advanced forget gates, from LSTM (Schmidhuber & Hochreiter, 1997) to Mamba2 (Dao & Gu 2024) and RWKV-7 (Peng et al. 2025a); (3) expressive
ory, from vector memory in RetNet (Sun et al. 2023 and LRU (Orvieto et al. 2023) to deep neural
memory in Titans (Behrouz et al. 2025c).

---

Published as a conference paper at ICLR 2026
Table 1: Overview of recent architectures in MIRAS framework perspective. Surprisingly, all mod- els are using the same type of attentional bias and regularization (forget gate). Note that these architectural choices does not uniquely identify the backbone as there other design choices (e.g., input-dependency, channel-wise parameters, etc.) as well as the use of other components.
Model
Andirectio
Ameriioaa Biss
Serania Five
Mensey Alpotitin
Merner Wilte Operation
Stallow Merrory
RetNer 2125
Vector
Die Product
L2
GD
M. MM- 1 VyAk
Mazik
Notpuranctic
M.
Maxix
Prodact
Vidi?"
Maxix
Product
(Dene to Mo.. Yeck
Marix
Product
La
+
Dst-Prddcct
Det-Prdact
WIRT
a.l
Dehoaer
Mare
C2)
Mr (. Andk
Longhom
Maric
I
Implic GDD
M, Ma- (1
(raten Bx) vil
TTT: Invo
Merr
In
(2)
Mr. M. WPMM
Garaa Delostr PioAk
Mame
Ia
CD)
Mr M. (0.01 Bad 31 Avak!
RWKX 20100
Mame
Ia
La
C2
Mr Mu-s (daglo)) and Ami
DelaPeedie
Mazs
Is
L2
MGD'
Mr Ma-s 10. (T- 8 A.D
A.
Linnar Atteaiion (SWLA)) 1071g
Stiiig Winder
Maxic
Stdia Windows Die-Prdact
L2
Mr MuMa
Meronyy
TTTMM P 2004
S-yy MIP
In
(2)
Mi M. 6.0.M (.
4-layer MLP
Monnettrt
M, MuMe-
0,... -. M. ki
One D125g
GLU
Sliling Window
L2
GD
Ei MO
Adoa X0OM
8-layer MILINGLU
Shiding Window fo
12
Miace
5, - 4,S... Sel (.L0MM-L...)))
M. m 188- (8.)
LaCT 5025
GLU
Dot Product
Muce
M, Misi 1 1H5- 18/) S... WVE.M--... Ro..))
Atentoo CMS 20256
MILPIUUU
NTP
La
Afourryy
M. ouMili WVCLMG.I...
MONETA (oum)
bilyyer MLL
I,
by
A n.Ae- W.WA stie v.), (1g A. vil
b-byyer MLP
Hiber
B, odRii
La (V4, At N..
If (MMAA--M SA,
M & frrm)
b-byyer MIPP
In
Kl
C20
U1 in Sefmoo (vy keg o W.We, v.A
using mutippl rounds of per tolenn Titass local and global resention using Gy loss. NTP. Token Prodicton.
At the core of these advancements lies a critical question: "what is the underlying design framework behind these sequence models, and how can these models be enhanced?". Taking inspiration from the broad definitions of associative memory and learning in neuropsychology literature (Okano et al.) 2000). several studies discuss connections between Transformers and Recurrent Neural Networks (RNNs) with associative memory (Hopfield) 1982, Ramsauer et al. 2021 Bietti et al. 2023). These studies, however, either: (1) lack universal explanation to fully illustrate the underlying learning al- gorithms, (2) are limited to specific definition of associative memory and lack generalizability, and (3) are unable to describe standard, widely-used components such forget gate. To address these concerns, several works have tried to unify neural architecture designs. Notably, Liu et al. (2024a) adopted an online learner viewpoint, similar to (LearningRRetaiinng Viewpoint) in our paper. Con- currently. Wang et al. (2025) adopted an online T--regression based viewpoint, which connects to the special case of (FTRL. Viewpoint) in our work with L--regression as the objective.
While these frameworks successfully unify existing models, they remain constrained by the & and dot-product paradigms, effectively making them specific instances of our broader framework. MI- RAS distinguishes itself in two critical ways: First, we provide a formal connection between these two viewpoints (Theorem (..2). Second, and crucially, MIRAS transcends the limitations of Eu- clidean spaces. Unlike prior frameworks that merely catalog existing by-based methods, MIRAS provides the generative capacity to design novel architectures with enhanced robustness and stabil- ity, addressing the sensitivity to outliers inherent in based optimization.
Contributions. Inspired by the human cognitive phenomenon of attentional bias the natural ten- dency to prioritize certain stimuli-we re-examine the foundations of sequence modeling by con- necting (online) optimization and associative memory. This perspective allows us unify existing architectures and unlock principled design space. Our main contributions are as follows:
A Unified Framework: We introduce MIRAS a comprehensive framework that reconceptualizes sequence models (including Transforners and modern RNNs) as associative memory modules governed by online optimization. MIRAS formally defines the core components of these models as Attentional Bias (the internal learning objective) and Retention (the memory regularizer).
"Miras" is the translation of "Ieegacy" in several languages including Persian, Arabic, and Turkish. We choose this name since this framework provides clear steps for future design of sequence models.
2

---

Published as a conference paper at ICLR 2026
Theoretical Reinterpretation and Critical Insights: Through the lens of MIRAS, provide inter- pretation of existing forgetting mechanisms (e.g.. gates LSTMs or Mamba2) as specific forms of regularization within online optimization frameworks (e.g., FTRL). Crucially, unification reveals a significant limitation: virtually all existing successful architectures rely narrowly on & loss dot- dot-product similarity for both bias and retention Table 1).
Expansion of the Architectural Design Space: MIRAS provides a principled foundation for moving beyond the f2 paradigm. We leverage principles from robust optimization and statistics to propose and explore novel, non-Euclidean attentional biases (e.g., Huber loss, Ep norms) and retention gates (e.g.. KL-divergence, f-divergence). leading architectures with improved stability and robustness. We specifically derive eight of these variants in |Section 4
Artentionne Architectures and Validation: Utilizing this expanded design space, we intro- duce three novel, attention-free free, and parallelizable architectures: MОNETА, YAAD, and MEM- ORA. These models combine expressive MLP-based memories with our novel optimization mech- anisms. Empirically, we demonstrate that different MIRAS instantiations trade off complementary strengths, achieving strong scaling laws and superior performance compared to state-of-the-art Transformers and linear RNNs across language modeling, commonsense reasoning. and challeng- ing long-context recall tasks.
2 ASSOCIATIVE MEMORY, ATTENTIONAL BIAS, AND RETENTION
Associative memory, a core component of human learning Terry 2017). has inspired numerous artificial neural architectures (Hopfildd 1982, [Schlag al. 2021 Behrouz et al. 2025c). Broadly speaking, associative memory is an operator that learns mappings between keys and values. To learn these mappings effectively, the memory requires an objective function that measures the quality of the learned associations and guides the learning process. In all formulations, one can set values as identity elements, resulting in formulation beyond strict design of key-value memory. In fact, while we discuss MIRAS framework based on associative memory, the concept of associative memory not limited to only the design between key-value memory modules and can include design choices such continuum memory system (CMS) (Behrouz et al. 2025b) that updates their weights based on the objective of the task at hand (e.g., next-token-predicteee in language modeling tasks).
Definition 2.1 (Associative Memory and Attentional Bias). Given set of keys K. C Rd and values V C Rd. associative memory is an operator M : K + V. Learning the mapping of associative memory is based objective L, called Attentional Bias, that determines the type of memory and its tendency to prioritize some events:
M arg min M
L(M(K);)))
When memory parameterized by W, we use M(W, k). In this setting, the optimization in equa- tion is performed over W. This learning process can be viewed meta (in-context) learning task, where the model learns how to store data into its parameters at test time (Munkhdalai et al.,
2019
Behrouz al. 2025b).
2.1 THE OPTIMIZATION PERSPECTIVE: LEARNING AND RETAINING
Definition 2.1 translates the design of sequence model into an optimization problem. In an online setting, where key-value pairs (k; vz) arrive sequentially, a straightforward approach to optimize Equation I is to use gradient descent. Given new pair, we update the memory parameters:
W. = W,- - 7/. Ve(WW...1 k;, v.),
(2)
where ((W-1ik,, v= C( v;). This update can be interpreted as adjusting the mem- ory based on a "momentary surprise" (Behrouz et al. (023c)). where the model prioritizes memoriz- ing tokens that violate the expectations of the objective L. This update rule highlights a fundamen- tal tension in sequence modeling: the need to learn from the latest information (adaptability) while remaining stable enough to retain previously memorized context (stability). We can formalize this tension by viewing the gradient descent update (2) through an optimization lens. Mathematically, Equation 2 is equivalent to the solution of the following optimization problem:
W, arg min
{ - WT-1, 70(W,11;;; v.)) + 2n.
I W,-1l12
33)
3

---

Published as a conference paper at ICLR 2026
The first term locally approximates e(W: k2 v;) at the previous state W;-1; minimizing it corre- sponds to learning the new token. The second term is a quadratic penalty that regularizes deviations from W,11; minimizing it corresponds retaining past information and ensuring stability.
2.2 THE SEARNING-RETAINEE VIEWPOINT
The formulation in equation relies specifically on linear approximations and quadratic regular- ization. However, we can generalize this concept by employing different approximations for the attentional bias and alternative functions for retention. This generalization leads to:
W₂ arg min WEW
(WW V.) + Ret, (W, (w-))
(Learning-Retaining Viewpoint)
Attentional Bias
Retention
Here, (WW k,, vz) is an approximation of E(W; ve), driving the learning of new concepts. Ret, (W, Wi-1) is the retention function, regularizing changes in W to maintain stability and pre- serve learned knowledge. This viewpoint has also been acknowledge by Liu et al. (2024a).
The retention function can be further decomposed into local and global components:
Ret, (W, We-1)
Dr (W, Wi-1)+ + G. (W). T/t
1
at Global Retention
Local Retention
The local retention D, (W, Wi-1) is typically premetric (e.g., l distance, KL divergence) that controls deviations from the immediate past state W,- The coefficient y)t acts as a meta in-context learning rate, balancing learning (larger 17)) against retention (smaller 17.). The global retention G controls the overall complexity size of the memory (e.g., weight decay).
Remark "Forgetting" Regularization. Within this viewpoint, mechanisms often termed "forget gates" (Behrouz et al. [225cc] Yang et al., 2024a) are reinterpreted not as explicit erasure mechanisms, but as specific implementations of the Retention function (regularization). The model optimizes how much of the past state to retain by balancing the regularization penalty against the learning objective. This interpretation is crucial as it provides a principled way to design novel retention mechanisms derived from optimization theory (See [Section 4), rather than relying on heuristic gating structures. Therefore, we use the term Retention Gate throughout this work. This interpretation aligns closely with human memory processes, where memories often become inaccessible due to retrieval failures rather than complete erasure (Robertson) 2002).
2.3 ALTERNATIVE PERSPECTIVE:
(FTRL)
While the Learning-Retaining viewpoint focuses on the trade-off at the current timestep, an alterna- tive perspective from online optimization considers the entire history of the sequence. The update rule in [Equatton| 2 can also be viewed as one step of online gradient descent on the sequence of losses &(W: k1, V1), e(W; k2, Vz),..., L(W; ls,, Vz),.... Online gradient descent is special case of the Follow-Thee Serlow-Tee-Rggulee ththe--eglarized--eeee
(FTRL) algorithm Sbalev-Shwarte et al. 2012 Hazan et al.) 2016). In FTRL, the goal is to minimize the cumulative loss over all past tokens, balanced by a reg- ularization term that penalizes the overall complexity of the memory. This leads to:
W = arg min WEW
Σ &iWW v.) + Tht
1 RR)
(FTRL Viewpoint)
Memory Stability
Attentional Bias
Here, (W;z))) represents an approximation (e.g., linearization) of the loss time i, and (W) is the regularization term. The [Learnin-Reeaining Viewpoint) and (FTRL Viewpoint) offer com- plementary perspectives (one local, one global) and can formally be connected:
Theorem 2.2. Let ft. = 1) and define (wW == Σ(=1 (WW;; ke, v:) + + RRWW). Assume W Bd and the function (W)) is strictly convex in W and let (x(,,)) be the Bregman divergence defined by function h(-), i.e., (W, W') h(W)= -h(W") (Vh(W"), W W'). Set Ret,W,, W') D, (W. and (,(W xx)) = (WW xx) in Ceaariing-Retiiiinng Viewpoint Then, the update rule in Learning- Viewpoint) is equivalent to the update rule (TRLL Viewpoint
4

---

Published as a conference paper at ICLR 2026
Associative Memory
Memory Architecture
Attentional Bias
Retention Gate
Memory Algorithm
neural architecture that
Internal objective.
The retain the the memory.
The algortthm that learns the magping.
Begression Loss
7y Regularizaaion (Local Clobal)
Cradicm
2. Matrix
2. Dot Product Similarity
2. Elastic Not Regularization
with Momeotum
3. Multilayue Puresptroa (MLP)
S. Haber Luwe
XL Divwrgenne
3. Newton's a Method
Mamory
Ki-Diverpence
Dregman Divenjence
Non-parametric Solutions
Associative Memory is neural network that learns to map keys to values based on an Attentional Blas objective.
Gradient Deapmnt
Ret M
Unsternal Objectron)
(Rtuiin Post Momoctesi
Input
Memory
Attentional Blas
Retention Gate
Effect of Choioes Output
Figure The overview of MIRAS framework. MIRAS based on four critical choices of (1) memory architecture, (2) attentional bias, (3) retention gate, and (4) memory learning algorithm. In this framework, the memory architecture determines the model capacity to memorize; attentional bias is responsible for modeling the underlying mapping patterns; retention gate determines how to balance learning new concepts and the retention of previously learned concepts; and memory learning algorithm is responsible for memory management.
The proof is provided in Appendix Intuitively, Theorem 2.2 confirms that optimizing the cumula- tive loss over the entire history (FTRL), can achieved via optimizing for the immediate trade-off between learning and retention in -etrrinng-Rttiiingg Viewpoint). provided the retention function appropriately chosen. This suggests that the Lxearning-Retaining Viewpoint) is more general. We therefore adopt it the primary lens for most of our derivations in subsequent sections.
2.4 EXTENSIONS AND GENERALIZATIONS
The Learning-RRtaaiingg Viewpoint) can be naturally extended provide more flexibility. We can consider a Universal Viewpoint where the memory update relies on recent history, rather than just the immediate past state:
W. arg min (WW vi)l + Ret, (W, (W,-1) 1
(Universal Viewpoint)
Attentional
Retention
Here, the memory update relies on its recent + states and the latest k'' + 1 key-value pairs. While Bearning-Retaining Viewpoint) is special case of (Universal Viewpointh by setting k') 0), one can obtain more flexible designs by choosing higher values of k, K. For ex- ample, setting k) (0,1) can recover optimization algorithms with momentum, such as those used in Titans (Behrouz et al. 2025c), as discussed in Appendix C
In above viewpoints, instead of using the global minimizer of the update, one can approximately find the minimizer by simply using (or multiple) steps of particular optimizer. For example, one can consider more advanced optimizers (such as Adam or Muon) for updating memory (Clark) et al. 2022 Zhang al. 2025) Behrouz et al. 2025a). In the next section, summarize the design choices within the MIRAS framework.
3 MIRAS: A FLEXIBLE FRAMEWORK FOR LEARNING TO MEMORIZE
Building viewpoints presented in the previous section, present MIRAS framework that not only accurately unifies existing architectures but also provides insights on how to design the next genera- tion of neural architectures. As discussed in Section 2 learning associative memory preted as a meta-learning task, which the associative memory learns how to compress and store data into its parameters at test time Sun et al. (204). The architecture of the memory in such task is
be inter-
5

---

Published as a conference paper at ICLR 2026
particularly important as in longer contexts as the expressivity of the memory structure limit its ability to learn the underlying patterns. Therefore, the first choice to design a neural architecture is the structure of the memory. Given the structure of the memory, parameterized by set of parameters W. we aim to minimize a loss function ((W: :,,)) with a retention regularizer Ret(-) via a learning algorithm (e.g.. gradient descent). Accordingly, MIRAS requires four design choices:
1. Memory Structure:
This choice specifies the architecture of the memory. For example, this architecture can be linear function, a Multilayer Perceptron (MLP) layer, even more complex structures. We may restrict the choice of W to be within a certain region, e.g., W to lie within an L2 ball avoid infinite values or unstable learning.
2. Attentional Bias: A key design choice is the objective L(-) (and consequently its approxi- mations in different viewpoints). This choice determines how we memorize the context and prioritize the events.
3. Memory Stability and Retention: Another key choice is the retention regularizer. This choice balances learning with retention of past state. An effective retention gate is key to reliable performance long context.
4. Memory Algorithm:
Finally, this choice specifies the learning algorithm that we use to optimize the memory objective. One may use exact minimizer in our framework or use one step (or multiple steps of) particular optimizer to update the memory.
The design choices of MIRAS are summarized Figure T In Appendix D we detail how vari- existing prchitectures-inclleee Mamba2), RNNs with Delta rules (c.g., DeltaNet), and deep memory models (e.g.. Titans)-can be derived specific instantiations of the MIRAS framework. In particular, Table T] in the appendix provides a comprehensive overview of this unification. Crucially, this analysis reveals a significant limitation: almost all these methods rely narrowly on l2 or dot-product attentional biases and & retention gates. MIRAS allows other choices of attentional bias/retention gates. In the next section, we discuss how going beyond the standard choices can lead to new architecture designs.
Softmax Attention, RNNs with Hebbian rules (e.g., RetNet,
BEYOND EXISTING ATTENTIONAL BIASES AND RETENTION GATES
As discussed in the previous section, existing work focuses only on linear/quadratic choices for the attentional bias or retention gate. However, in general there could be various choices for all the three aforementioned design choices (even by going beyond Euclidean space). To illustrate the flexibility of our designed framework, this section proposes and discusses novel design choices in MIRAS.
We start by discussing novel choices of attentional biases and retention gates in Section 4.1 We only briefly present three of such variants in this subsection and we will relegate further choices Appendix E. Then, we combine our design choices in [Section 4.2 obtain three particular architectures: Moneta, Yaad, and Memora, which we further experiment on in Section 5)
4.1 NOVEL ALTERNATIVE ATTENTIONAL BIAS AND RETENTION GATES
Variant I:
e,tAttentional Bias. Attentional bias defines the "similarity metric" and measures how well memory can recall the value, given its corresponding key. Although &g loss is used in prior work, natural extension is f--norm class of objectives: We define Ey-attentional bias as:
I|M v.).v. l. -v
(4)
where p R21 and I|-IIp is the p-norm. Depending on the distribution of the data, we might want to use different values of p (see Section 5). For the sake of simplicity, let memory be matrix defining a linear mapping, i.c. MM,, k,, = W ke, the gradient descent update is:
W.- W n. V. v., v.) W - p (Sign v.) @ Wi-ik, volp-
A
(5)
where @ is the Hadamard (element-wise) product. For p = 1, the recurrence simplifies to: W. = Wi-1 - 7) Sig ve) k,T We call this variation value-less associative memory, in which we store entities (keys) but map them into two extreme class of 1 and +1 through the sign function. This behavior provides inherent robustness, the magnitude of the error does not affect the update

---

Published as a conference paper at ICLR 2026
direction, preventing extreme events (outliers) from overly influencing the memory. One simple interpretation for such behavior is the coping mechanism in human (Lofuus) 1993), in which the memory does not store the values for extreme events. This interpretation of protective memory in extreme events motivates our next variant.
Variant 2:
Huber Loss: Memory with Coping Mechanism. While f2-norm regression objective
common choice, it is known to be sensitive noise and extreme examples (outliers). To have a robust loss against outliers, we can use Huber loss as the attention bias, in which an extreme mismatch (potentially due to outlier data) does not affect the memory learning process:
Wi = W - 1/t Vl2(Wrr) kr, Vt)
if [[M(W; k.) - v.ll S 6,, T/ on V&i (W-1;; v.) - 83 Otherwise.
(6)
In this formulation, the parameter 5, decides the type of the memory (f2-norm objective or value- less) based on the context, making the memory more robust outliers.
Variant 3: Memorization Over A Scaled Probability Simplex. To avoid numerical instabilities, we can constrained the variable W₂ to lie within a scaled probability simplex. In other words, we can restrict the state to lie in the constraint set W = {W [[]]] and Wj 2 O, Vj. In this set, each point W can be viewed as a measure. Thus, we can utilize divergences over measures to define our retention in
Viewpointh For example, by choosing (W, W") as the KL divergence and setting &(W; kt, vt) = (W - W2-1, Ve(W,-1; vz)) in |Learning-Retaining Viewpoint). we get the update rule
Seaaring--Retaniigg
W₂ < c Softmax ((1 - X) 1og(W,-1) - 7) V ks, v:))
where À E 10 and r E IR+ are the syper-parameters that can be learned during training. The Softmax operator ensures that the output lies in the set W. To more discussions and extensions to general f-divergence retention gates, see SSection E.2 in the appendix.
Other Variants. In the appendix section, we derive other novel variants of MIRAS by using elastic net regularization, Bregman divergence and f- divergence retention gates, and by memory stability. We will also discuss the above variants in more details.
4.2 FOCUS VARIANTS OF MIRAS: MONETA, YAAD, AND MEMORA
Using the above basic variants and the other variants explained in Appendix Ej we now introduce three instantiations of MIRAS, each designed to explore different facets of this expanded optimiza- tion space, moving beyond the standard f2 and dot-product paradigms.
MONETA. MONETA is designed to investigate the impact of generalized norms. Given p, 9 E 2 design (p, -o-NEAA to explore the impact of generalized Ep norms for both learning and regularization. We instantiate MIRAS as follows. Memory Structure: A 2-layer MLP with expansion factor 4, GELU activation (Hendrycks & Gimpell 2016), residual connections, and Layer Norm (N(11(W22)) (22)). Attentional Bias: Ep norm. Retention Gate: A hybrid of fg retention gate 2(q-1) || WI? (see Appendixx for details) and the standard f2 regularization 1 || Memory
M(r) = x + o
1
Algorithm: Gradient Descent. The above choices result in the following recurrent formula for the memory module:
A, 3,4, 7/ Ve(W vz), and W. = A
(8)
||A, 18 We use q) (3, 4).
Notably the gradient can be calculated using Equation 5
YAAD. YAAD designed for robustness, protecting the memory from extreme events (outliers) us- ing principles from robust statistics. We design YAAD based on the Huber objective.Memong Struc- ture: MLP (same architecture as MONETA). Attentional Bias: Huber loss (Equation 6). Retention
Gate: A combination of local and global retention: Ret, W.1) = 2m. W
W
Memory Algorithm: Gradient Descent. Given these choices, we can write the resulting memory learning process as
(/ V62 k..
if [[M(k)) - vell $ 5,, 8, (W v.) -6 Otherwise.
W. = W.W,
{
(9)
7

---

Published as a conference paper at ICLR 2026
Table Performance of MIRAS's variants and basclines on language modeling and common-senae reasoning tasks. Hybrid models are marked with * The best results are highlighted
Wikd. ppl
LMB. ppl
LMB. acc
PIQA acc
Hella. acc. n
Wine. acc
ARC- acc
SIQA acc
BoolQ acc
Avg. T
acc.
params/
tokens
Transformee++
18.53
18.32
42.60
70.02
50.23
53.51
68.83
35.10
40.66
57.09
52.25
RetNet
19.08
17:27
40.78
52.02
Mamba2
16.56
12.56
45.66
71.87
55.67
55.24
72.47
37.88
40..20
60.13
$4.89
DeltaNet
17.71
16.88
42.46
70.72
50.93
53.35
68.47
35.66
40.22
55.29
$2.14
Gated DeliaNet
16.42
12.17
46.65
72.25
$5.76
57.45
71.21
38.39
40.63
60.24
55.32
Sambar
16.13
13.29
44,94
70.94
53.42
55.56
68.81
36.17
39.96
62.11
$4,00
Gated DeltaNet-HH2
15.91
12.55
48.76
72.19
56.88
57.77
71.33
39.07
41.91
Titans (LMM)
15.60
11.41
49.14
73.09
56.31
59.81
72.43
40.82
42.05
60.97
$6.82
MONETA (ours)
15.52
56.14
59.09
YAAD (ours)
15.18
11.89
47.23
72.81
56.46
59.02
72.14
40.05
40.73
61.86
$6.39
MEMORA (ours)
15.90
12.04
48.67
53.99
57.36
40.19
55.87
Note that for improving the expressive power, in all architectures, we decouple the learning rate n) and the retention gate rate a, resulting in independent parameter Br E [0, 1].
MEMORA.
Finally, MEMORA is designed to ensure stable updates by constraining the memory to a probability simplex and utilizing divergence-hard -haseed retention. Memory Structure: MLP (same as
MONETA), constrained to the scaled probability simplex. Aftentional Bias: dot-product loss. Reten-
tion Gate: KL-ddivergence for local retention and Shannon entropy for global retention (Appendix E).
Memory Algorithm: Closed-form solution. These choices lead (see equation 30
Wi Softmax (3, log((,))))) Ti V (-: v:)
(10)
Architecture Backbone. For the architectural backbone, we fully follow recent studies (Behrouz
ctt al., 2025c Yang et al., 2024a)): We replace attention modules with our variants of MIRAS in Llama's s macro architecture with MLPs with SwiGLU(, activation, rotary positional encodings (RoPE) Su et al. al. (2024), and RMSNorm (Zhang & Sennrich 2019). We incorporate a 1D depthwise- separable convolution layer after each of the query, key, and value projections. For training stability, we also use & normalization to o and kt. The output of this module normalized and gated with a linear layer (Mehta al. 2023). For all input-dependent parameters like n/. Br. and 8,, define them as the linear projection of the input. The architectures are illustrated in Figure 5)
Parallelizable Training. We build upon the work of Behrouz et al. (002cc); Sun et al. (0224) and use a hybrid recurrence of linear and non -on-linear& by chunking the sequences into small subsequences. While the use of MLP memories and non Euclidean optimization introduces non non-linearities in the recurrence, the hybrid chunking strategy (Appendix F) ensures that the training remains highly par- allelizable. Inside each chunk, the recurrence is effectively linearized, and non-linear operations (e.g., the normalization in MONETA or Softmax in MEMÓRA) are applied only at chunk bound- aries. This maintains competitive training throughput while offering O(1) complexity per token during inference.
EXPERIMENTS
Experimental details (resp. additional experiments) are in Appendix G (resp. Appendix H).
5.1 LANGUAGE MODELING AND COMMON-SENSE REASONING
Behrouz et al., 2025c) and first focus on the per- plexity in language modeling and commonsense reasoning tasks. The results for MEMORA, YAAD, МONETA and baselines with size of 1.3B are reported in Table 2 (Full results of 340M and 760 in Table 7). All of our variants outperforms all the baselines including Transformer++, modern linear recurrent models and hybrid methods. The superior performance compared to hybrid models is par- ticularly important all our variants are pure recurrent (attention-free) Among the three variants of MIRAS, while MONETA achieves slightly weaker performance than MEMORA, and YAAD, the other two variants are close and depending on the task and model size, the best model can vary.
We follow recent studies YYang et al. 2024a c
5.2 SCALING PATTERN
To evaluate the scaling pattern of models and for comparing them with baseline, in this section, we plot their performance with varying the model size and the context window.
8

---

Published as a conference paper at ICLR 2026
25
22
22
Mamoa2
23
320
11
16
1000
I0C
Figure 2: Scaling patterns when increasing (Left) model size, (Middle) sequence length (model size 340M) (3) (Right) sequence length (model size = 760M) on C4 dataset.
Context Length. We first vary the training context length from 2K to 32K for two version of our model with size 340M and 760M. The results are reported in Figure 2 (Middle and Right). All three variants of MIRAS scales better than state-of sfte--of-the-ar art baselines when increasing the context length. We attribute this superior performance to: (1) expressive memory architecture. Contrary to baselines like Mamba2 and GSA that uses vector- and matrix-valued memory, our variants use 2- layer MLPs with more expressive power. (2) The choice of retention gate and attentional bias. While TTT also uses MI.P memory, it shows weaker scaling. This highlights a crucial finding: expressive memory alone insufficient; requires correct design choices (e.g. attentional bias, refention, and optimization algorithm) to be effectively utilized. All of our three variants go beyond the standard l-- based attentional biases and retention gates. These robust choices prevent memory corruption from outliers or noise, leading to better utilization of the fixed capacity, especially in long contexts.
Model Size. We also report the #FLOPs perplexity of models and baselines in Figure 2 (Left). All three variants outperforms all baselines given almost the same budget of FLOPs. These results, again support the importance of powerful memory design.
5.3 NEEDLE IN HAYSTACK
To evaluate the effective context window of our models and baselines, we use needle-in-a-haystack task. In this task, we evaluate the model on retrieving piece of information (i.e., the "needle") from long distractor texts (i.e., the "haystack"). We focus on the Single NIAH (S-NIAH) task from RULER benchmark (Hsich et al.) (024) and evaluate our models and baselines on sequences with length 1K, 2K, 4K, and 8K. The results are reported in Table 3, All our variants outperforms the baselines by a considerable margin. Interestingly, MONETA shows superior performance when the data is synthetic noise (S- (S-NIAH-PK). This highlights the advantage of MONETA'S 's E,-attentional bias and Ey retention (with (p. q) = (3, 4)). Unlike the & objectives used baselines, these higher-order norms inherently more robust to noisy inputs, preventing the distractor texts from corrupting the memory state. This validates the effectiveness of exploring non- non-Euclidean design choices via MIRAS.
Table 3: Performance of MONETA, YAAD, MEMORA, and base- lines on NIAH task from RULER benchmark. The best results with highest accuracy are highlighted.
Table 4: Ablation on the architec- ture of MEMORA and MONETA.
S-NIAH-PK
S-MAHHN
S-NIAH-W
Model
Average
2K
4K
SK
2K
4K
8K
IK
2K
4K
Mambal
98.6
61.4
31.0
98.4
55.8
14.2
62.2
42.2
4.2
$2.0
Variant
MEMORA
MONETA
DeltiNet
96.8
98.8
98.6
47.2
15.4
12.8
85.2
46.2
20.0
$7.9
Full Architecture
31.52
52.12
Gated DeltaNet
89.8
91.8
26.4
86.4
82.6
24.4
75.8
Retention
49.75
linear memory
50.11
50.26
wlo RoPE
51.28
51.71
MONETA
99.4
98.8
98.8
99.4
99.4
92.8
92.2
X8.3 2
70.8
93.5
YAAD
99.2
98.6
94.4
99.8
98.6
93.2
91.8
89.6
67.4
92.9
MEMORA
99.2
98.8
98.4
99.2
93.2
92.4
70.4
92.1
5.4 ABLATION STUDY
We perform ablation studies 760M model to validate if different design choices discussed through the paper are positively contributing to better performance. Additional ablations are in Appendix H

---

Published as a conference paper at ICLR 2026
The Effect of Design. To evaluate the architectural design choices, we perform ablation study on MEMORA, and MONETA in Table 4 as well as on YAADIN Table 5| The first row, reports the performance of full architecture, while (1) the second row removes the retention (i.e. 3 = 1), and (2) third row replaces the MI.P with a linear layer. In Table 5 (3) forth row makes 5 input independent, (4) the next row removes f2-loss from the Huber loss, and (5) the last removes the t condition. These results indicate that all design choices are contributing to the performance of the model.
The Effect of p on Performance. We first evaluate the effect of p on the performance of MONETA. We vary the value of 2,, 2.8, 3, 3.2, 4} and context window from 2K to 16K. The results are reported in Figure 3 Interestingly, there no monotone pattern when increasing the value of p and the best performance is achieved when p = 3, while = 4 achieves the worst performance. Also, although different values of p results in different memory modules with varied performance, the scaling pattern when increasing the context length is almost the same.
The Effect of q on Performance. We evaluate the effect of by varying its value in (2,3, 4, 5]. Interestingly, contrary to p. the value of q can change the scaling pattern when increasing the context length. The main reason for this observation is that the value of q determines the retention gate and a powerful retention gate can improve the memory management, resulting in better performance in longer context.
Table 5: Ablation study on the components of YAAD.
Model YAAD
Avg. LM
53.98
2.2
12
Retention Gate
50.63
linear memory
51.57
Contoc: Length
Cuntoe: Lengh
Input-dependent å
52.19
Figure 3: The effect of parameters p and q on the performance with different context length.
fo-loss
52.86
floso
53.04
6 LIMITATIONS AND DISCUSSIONS
While MIRAS framework provides a principled approach to designing efficient sequence models with strong performance in language modeling and long-context recall (e.g.. Needle-in-a-Hee Hays-aaystcck), our variants still lag behind quadratic Transformer models in complex in-context retrieval tasks (e.g., Table 9}. This highlights a known challenge for recurrent models, suggesting that while optimized memory management can replace attention in many scenarios, perfect in-context recall remains an area for future improvement, mainly due to the fact that is more a challenge for the fixed-size memory of RNNs rather than its memory update mechanism. One might combine the design choices in MIRAS with techniques like Memory Caching (Behrouz et al. 2026) to achieve recurrent mod- els with growing memory and effective memory management. Furthermore, while the framework connects learning in non-convex memories (MLPs) to online optimization, providing theoretical guarantees on convergence and learning in this setting remains challenging.
CONCLUSION
This paper presents MIRAS, general framework that explains the connection of online optimization and test time memorization. MIRAS framework can explain the role of several standard architectural choices in the literature (e.g., forget gate) and helps design next generation of architectures that are capable of managing the memory better. Building upon our framework, by using novel attentional bias and retention gates, we present three novel sequence models called YAAD, MONETA, and MEMORA-ead of which with its own (dis)advantages. Our experimental evaluations show that all these variants outperform various baselines in various downstream tasks.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Simran Arora, Brandon Yang. Sabri Eyuboglu, Avanika Narayan, Andrew Hojel, Immanuel Trum- mer, and Christopher Ré. Language models enable simple systems for generating structured views of heterogeneous data lakes. arXiv preprint arXiv:2304.094. 2023.
Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, James Zou, Atri Rudra, and Christopher Re. Simple linear attention language models balance the recall-throughput tradeoff. In Forty-first International Conference on Machine Learning, 2024. URL https:
//openrevis ew. net. forum? d=e 3DcpH3
Ali Behrouz, Parsa Delavari, and Farnoosh Hashemi. Unsupervised representation learning of brain activity via bridging voxel activity and functional connectivity. In Forty-forsl International Conference on Machine Learning, 2024. URL https://openreviee net ftrumoormmmmm
nO12EpLyhl
Ali Behrouz, Zeman Li, Praneeth Kacham, Majid Daliri, Yuan Deng, Peilin Zhong, Meisam Raza- viyayn, and Vahab Mirrokni. Atlas: Learning to optimally memorize the context at test time. arXiv preprint arXiv:2505.23735 2025a.
Ali Behrouz, Meisam Razaviyayn, Peilin Zhong, and Vahab Mirrokni. Nested learning: The illusion of deep learning architectures. In The Thirty-ninth Annual Conference on Neural Information Pro- cessing Systems, 2025b. URL htttttp:://peeeeeiiiieeeeee www.:eeeooooonnddmddeeeeeeeeeeee
Ali Behrouz, Peilin Zhong, and Vahab Mirrokni. Titans: Learning to memorize at test time. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025c. URL
Ali Behrouz, Zeman Li, Yuan Deng, Peilin Zhong, Meisam Razaviyayn, and Vahab Mirrokni. Mem- ory caching: Rnns with growing memory, 2026. URL
org. artp.r/////bb////2//
24281
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transformer: memory viewpoint. Advances in Neural Information Processing Systems, 36: 1560-1588, 2023.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, al. Piqa: Reasoning about physical com- monsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432-7439, 2020.
Christopher Clark, Kenton Lee, Ming-War Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. BoolQ: Exploring the surprising difficulty of natural yes/no questions. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Lan- guage Technologies, Volume (Long and Short Papers), pp. 2924-2936, Minneapolis, Min- nesota, June 2019. Association for Computational Linguistics. doi: 10. 00.18653/11/N11-11000 URI. https:://claatholl www.:hhho.oooo.o./o11....
Kevin Clark, Kelvin Guu, Ming-Wer Wei Chang, Panupong Pasupat, Geoffrey Hinton, and Mohammad Norouzi. Meta-learning fast weight language models. arXiv preprint arXiv:2212.02475. 2022.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv: 1803.05577 2018.
Imre Csiszar. On information-type measure of difference of probability distributions and indirect observations. Studia Sci. Math. Hungar., 2:299- 0099-118,177 1967.
Karan Dalal, Daniel Koceja, Gashon Husscin, Jiarui Xu. Yue Zhao Youjin Song, Shihao Han, Ka Chun Cheung, Jan Kautz, Carlos Guestrin, al. One-minute video generation with test-time training. arXiv preprint ar 25x4:2504.00299 2025.
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:2405.22 1060, 2024.

---

Published as a conference paper at ICLR 2026
Soham De, Samuel L Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Murant Al- bert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, et al. Griffin: Mix- ing gated linear recurrences with local attention for efficient language models. arXiv preprint arXiv:2402.1 19427, 2024.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010.11929 2020.
Dheeru Dua, Yizhong Wang, Pradeep Dasigi, Gabriel Stanovsky, Sameer Singh, and Matt Gardner. Drop: A reading comprehension benchmark requiring discrete reasoning over paragraphs. arXiv preprint arXiv:1903.0011 2019.
Gemma-Teane Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi. Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, et al. Gemma: Open models based on gemini research and technology. arXiv preprint arXiv:2403.0822 08295, 2024.
Riccardo Grazzi, Julien Siems, Jörg KH Franke, Arber Zela, Frank Hutter, and Massimiliano Pontil. Unlocking state-tracking in linear mns through negative eigenvalues. arXiv preprint arXiv:2411.12537 2024.
Klaus Greff, Rupesh K Srivastava, Jan Koutnk, Bas R Steunebrink, and Jürgen Schmidhuber. Lstm: A search space odyssey. IEEE transactions on neural networks and learning systems, 28(10): 2222-22333 2016.
Albert Gu and Tri Dao. Mamba: Linear-tine sequence modeling with selective state spaces. In First Conference on Language Modeling, 2024. URL /ttl:/ttss////eeereeiee..ttt//t/ fwt.r//meeeeeee//fdddddo tEYskw VY2
Trevor Hastie, Robert Tibshirani, Jerome Friedman, et al. The elements of statistical learning, 2009.
Elad Hazan et al. Introduction to online convex optimization. Foundations and Trends@ Opti- mization, 233-4):157-325, 2016.
Donald Olding Hebb. The organization of behavior: neuropsychological theory. Psychology press, 2005.
Dan Hendrycks and Kevin Gimpel. arXiv:1606.0844 08416, 2016.
Gaussian error linear units (gelus).
arXiv preprint
Donald E Hilt and Donald Seegrist. Ridge, a computer program for calculating ridge regression estimates, volume 236. Department of Agriculture, Forest Service, Northeastern Forest Experi- ment 1977.
Arthur E Hoerl and Robert W Kennard, Ridge regression: applications to nonorthogonal problems. Technometrics, 12(1):69-82, 1970.
John J Hopfield. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the national academy of sciences, 7955)25554-2558 1982.
Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, and Boris Ginsburg. RULER: What's the real context size of your long-context language models? In First Conference Language Modeling, 2024.
id=kIoBbe76Sy
Peter Huber. Robust estimation of a location parameter. In Breakthroughs in statistics: Methodol- ogy and distribution, pp. 492-518. Springer, 1992.
Kazuki Irie. Robert Csordas, and Jürgen Schmidhuber. Practical computational power of linear transformers and their recurrent and self-referentiil extensions. arXiv preprint arXiv:2310. 16076, 2023.
12

---

Published as a conference paper at ICLR 2026
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.0836/ 2020.
Angelos Katharopoulos, Apoory Vyas, Nikolaos Pappas, and François Fleuret. Transformers mns: Fast autoregressive transformers with linear attention. In International conference on ma- chine learning, pp. 5156-5165, PMLR, 2020.
Aniruddha Kembhavi, Minjoon Seo, Dustin Schwenk, Jonghyun Choi, Ali Farhadi, and Hannaneh Hajishirzi. Are you smarter than a sixth grader? textbook question answering for multimodal machine comprehension. In Proceedings of the IEEE Conference on Computer Vision and Pattern recogmition, pp. 4999-5007, 2017.
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, et al. Natural questions: benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:453 466, 2019.
Aonian Li, Bangwei Gong, Bo Yang, Boji Shan, Chang Liu, Cheng Zhu, Chunhao Zhang, Congchao Guo, Da Chen, Dong Li, et al. Minimax-01: Scaling foundation models with lightning attention. arXiv preprint arXiv:2501.08313 2025.
Chengxuan Li, Di Huang. Zeyu Lu, Yang Xiao, Qingqi Pei, and Lei Bai. A survey on long video generation: Challenges, methods, and prospects. arXiv preprint 2Xiv:2443..14077 2024.
Bo Liu, Rui Wang, Lemeng Wu, Yihao Feng, Peter Stone, and Qiang Liu. Longhorn: State space models are amortized online learners. arXiv preprint arXiv:2207. 14207, 2024a.
Nelson F Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models use long contexts. Transactions of the Association for Computational Linguistics, 12:157-172 2024b.
Colin Lockard, Prashant Shiralkar, and Xin Luna Dong. Openceres: When open information extrac- tion meets the remi-structured web. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Vol- ume I (Long and Short Papers), pp. 3047-3056, 2019.
Elizabeth F Loftus. The reality of repressed memories. American psychologist, 48(5):518, 1993.
Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language model- ing via gated state spaces. In The Eleventh International Conference on Learning Representations, 2023. [ttps::t://///
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mix- ture models. In International Conference on Learning Representations, 2017. URL https:
//openreview.neee nwt. By
Tsendsuren Munkhdalai, Alessandro Sordoni, Tong Wang, and Adam Trischler. Metaleamed neural memory. Advances Neural Information Processing Systems, 32, 2019.
Hideyuki Okano, Tomoo Hirano, and Evan Balaban. Learning and memory. Proceedings of the National Academy of Sciences, 97(23): (2:3):12400 1403-122044 2000.
Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pas- canu, and Soham De. Resurrecting recurrent neural networks for long sequences. International Conference Machine Learning, pp. 26670-26698. PMLR, 2023.
Denis Paperno, German Kruszewski, Angeliki Lazaridou, Ngoc Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernandez. The LAMBADA dataset: Word prediction requiring broad discourse context. In Katrin Erk and Noah A. Smith (eds.), Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Vol- ume 1: Long Papers), pp. 1525-1534, Berlin, Germany, August 2016. Association for Com- putational Linguistics. doi: 10.18653/v1/P166 1144. URL https: hwtp//aaannntto.......ooooo
P P1
13

---

Published as a conference paper at ICLR 2026
Guilherme Penedo, Hynek Kydlcek, Loubna Ben allal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, and Thomas Wolf. The fineweb datasets: Decanting the web for the finest text data at scale. In The eiiht Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2024. UtLhttps::/ /ttpss////pereeeeeeeee net. /forumrr id==nSCkn22aa
Bo Peng. Ruichong Zhang, Daniel Goldstein, Eric Alcaide, Xingjian Du, Haowen Hou, Jiaju Lin, Jiaxing Liu, Janna Lu, William Merrill, Guangyu Song, Kaifeng Tan, Saiteja Utpala, Nathan Wilce, Johan Wind, Tianyi Wu, Daniel Wuttke, and Christian Zhou-Zheng. RWKV-7 "goose" with expressive dynamic state evolution. In Second Conference on Language Modeling, 2025a. URL http:ppps///pppeeeeiieeee net wtttteeuttt/eeeeeeeoeeee http://oorrrn1dddAAPAAAAAAA
Bo Peng, Ruichong Zhang, Daniel Goldstein, Eric Alcaide, Haowen Hou, Janna Lu, William Merrill, Guangyu Song, Kaifeng Tan, Saiteja Utpala, et al. Rwkv-7" goose" with expressive dynamic state cvolution. arXiv preprint arXiv:2503, 14456, 2025b.
Michael Poli, Armin W Thomas, Eric Nguyen, Pragaash Ponnusamy, Bjorn Deiseroth, Kristian Kersting, Taiji Suzuki. Brian Hie, Stefano Ermon, Christopher Re, al. Mechanistic design and scaling of hybrid architectures. arXiv preprint arXiv:2403.17844 2024.
Yury Polyanskiy and Yihong Wu. Information theory: From coding to learning. Cambridge univer- sity press, 2025.
Zhen Qin, Songlin Yang, Weixuan Sun, Xuyang Shen, Dong Li, Weigao Sun, and Yiran Zhong- HGRN2: Gated linear RNNs with state expansion. In First Conference on Language Modeling, 2024.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140):1-6) 2020.
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100,000+ questions for machine comprehension of text. arXiv preprint arXiv: arXiv: /606.5250, 2016.
Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Lukas Gru- ber, Markus Holzleitner, Thomas Adler, David Kreil, Michael Kopp, Günter Klambauer, Jo- hannes Brandstetter, and Sepp Hochreiter. Hopfield networks is all you need. In International Conference on Learning Representations, 2021. URthttps: htt.t:pp:///penneiiee..... net forum?
1.= T Cd
Lee T Robertson. Memory and the brain. Journal of dental education, 66(1):30 2002.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adver- sarial winograd schema challenge at scale. Communications of the ACM, 64(9)999-106, 2021.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In International Conference on Machine Learning, pp. 9355-9366. PMLR, 2021.
Jürgen Schmidhuber and Sepp Hochreiter. Long short-term memory. Neural Computation MIT- Press. 1997.
Shai Shalev-Shwarzz et al. Online learning and online convex optimization. Foundations and Trends@ in Machine Learning, 4(2):177-19, 2012.
Julien Siems, Timur Carstensen, Arber Zela, Frank Hutter, Massimiliano Pontil, and Riceardo Grazzi. Deltaproduct: Increasing the expressivity of deltanet through products of houscholders. arXiv preprint www:iiv:202..029900 2025.
Jimmy T.H. Smith, Andrew Warrington, and Scott Linderman. Simplified state space layers for se- quence modeling, In The Eleventh International Conference Learning Representations, 2022. net. httphttps: htattps://openreviii htt.://.eee.foooooeeeeeeee forrmm i id=Ai BHw3Axqks
Jianlin Su. Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding. Neurocomputing, 568: 127063, 2024.
14

---

Published as a conference paper at ICLR 2026
Yu Sun, Xinhao Li, Karan Dalal, Jiarui Xu, Arjun Vikram, Genghan Zhang, Yann Dubois, Xinlei Chen, Xiaolong Wang, Sanmi Koyejo, et al. Learning to (leam test time): Rnns with expressive hidden states. arXiv preprint arXiv:2407.04620, 2024. 2024.
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang. and Furu Wei. Retentive network: A successor to transformer for large language models. arXiv preprint arXiv:2307.08622 2023.
W Scott Terry. Learning and memory: Basic principles, processes, and procedures. Routledge, 2017.
Robert Tibshirani. Regression shrinkage and selection via the lasso, Journal of the Royal Statistical Society Series B: Statistical Methodology, 58(1):267-288 1996.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie- Lachaux, Timothec Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E. ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Ad- vances Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL httpp::///rocceddss....eeee neurips........................m cc/paper ftles//aperr20117 Paperp
Johannes Von Oswald, Maximilian Schlegel, Alexander Meulemans, Seijin Kobayashi, Eyvind Niklasson, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Max Vladymyrov, al. Uncovering mesa-optimization algorithms in transformers. arXiv preprint arXiv:2309.0588 05858888, 2023.
Johannes von Oswald, Nino Scherrer, Seijin Kobayashi, Luca Versari, Songlin Yang, Maxi- milian Schlegel, Kaitlin Maile, Yanick Schimpf, Oliver Sieberling, Alexander Meulemans, et al. Mesanet: Sequence modeling by locally optimal test-time training. arXiv preprint arXiv:2506..0223,, 2025.
Ke Alexander Wang, Jiaxin Shi, and Emily B Fox. Test-time regression: unifying framework for designing sequence models with associative memory. arXiv preprint arXiv:2501.12352. 2025.
Yingheng Wang, Zichen Wang, Gil Sadch, Luca Zancato, Alessandro Achille, George Karypis, and Huzefa Rangwala. Long-context protein language model. bioRxiv, pp. 2024-10, 2024.
Songlin Yang, Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule. arXiv preprint Xiiv:2412.066644 2024a.
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficien training. In Forty-first International Conference on Machine Learming. 2024b.
Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transform- ers with the delta rule over sequence length. Advances in Neural Information Processing Systems, 37:115991-155522 2024c.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. HellaSwag: Can machine really finish your sentence? In Anna Korhonen, David Traum, and Lluis Marquez (eds.). Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 4791-4800, Florence, Italy, July 2019. Association for Computational I Linguistics. doi: 10. 18653/v1/P19 1472. hR
https: //tp//aaaannthooo htt.tsss//aalaaattolllllo www.hhoooo....oo httpholoooolooooo1111177777
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Infor- mation Processing Systems, 32, 2019.
Tianyuan Zhang, Sai Bi, Yicong Hong, Kai Zhang, Fujun Luan, Songlin Yang, Kalyan Sunkavalli, William T Freeman, and Hao Tan. Test-time training done right. arXiv preprint arXiv:2505.2388 23884, 2025.
15

---

Published as a conference paper at ICLR 2026
Tianyuan Zhang, Sai Bi, Yicong Hong. Kai Zhang. Fujun Luan, Songlin Yang, Kalyan Sunkavalli, William T. Freeman, and Hao Tan. Test-time training done right. In The Fourteenth International Conference on Learning Representations, 2026. URLhttps:/ openreview wwt. net//orumm forum?
Td=
16

---

Published as a conference paper at ICLR 2026
A PRELIMINARIES AND BACKGROUND
In this section, we review the related studies and background concepts that we use through the paper.
Attention. Attention as the backbone of Transformers is a critical component that acts as their associative memory (Bietti et al., 2023). Given input x E R&xxda causal attention computes output y E xx based on Softmax over input dependent key, value, and query matrices:
Q zW == WW.K
V=
(11)
Y1 exp qq) k, vdm) v
(12)
Ei=1 exp (q. kt/ V
where Wo. WK, and Wv E Rdd are learnable parameters. While Transformers achieve significant improvements compared to traditional Recurrent Neural Networks (RNNs)-such as LSTM (Schmidhuber & Hochreiter [997), their complexity that requires at least N x d opera- tors to calculate the output has been the main motivation for researchers to think about alternative architectures. We divide and review the research efforts to design alternative architectures into two groups: (1) Linear shallow memory recurrent models, (2) Deep memory modules.
(Linear) Recurrent Models. For many years, non-linear (gated) recurrent neural networks had been the de facto architectural backbones in deep learning (Greff et al., 2016). Their recurrent nature. however, results in non-pant paraaalllllli training. making their large scale training infeasible.
To this end, in recent years, linear RNNs as alternatives to both Transformers and non-linear RNNs attract much attention mainly due their parallelizable and tineartimme training while maintaining
competitive performance (Yang et al., 2024c Sun et al. 2023 Peng et al. (025a)). Earlier variants of linear RNNs (Yang et |l.]| 20.2bb Sun al., 2023) De et al. [0044). which mostly are based on Hebbian learning rule (Hebb, 2005). , aim to compress the data into their vector-valued (or matrix- valued) memory (Katharopoulos et al. 2020 [Snn et al. 2023 Yang et al.) 2024b) De et al., 2024] Liu et al. 2024a). Let M₁ E Ra* be the memory (n k, v E R" are keys and values (i.e., projection of input 2, E RR), a simple general formulation for such linear RNNS can be written as:
means valuee memory), and
M. A. + M + v
(13)
where is an arbitrary associative operator and Ar is a data-((n)dependent diagonal matrix or scalar (Yang et al., 2024c). Despite the efficiency that comes with the linear recurrent nature of these models, the memory can overflow mainly due to the additive (without replacement) nature of Hebbian learning rule, resulting in limited memory capacity and limited expressive power in in-context learning tasks. Moreover, the vector-valued memory of these architectures can limited their ability to learn/memorize large context window, mainly due to the limited expressive power of memory to learn the underlying patterns of data (Behrouz et al. 2025c Sun et al. 2024
To address the above mentioned limitations, recurrent models that use matrix- valued memory with Delta learning rule has gained popularity in recent years (Schlag et al., 202 Yang et al. 2024c). Despite significant advantages, even these delta-rule-basee recurrent models face theoretical limitations (Irie et al.) 2023) with moderate performance in practice (Yang al. 2024c). Recently, several studies aim to improve the performance of such models by adding scalar or channel-wise forget gate mechanisms (Yang et al. 2024a Peng al.] (025b).... using negative eigenvalues (Grazzi et al. 2024), and multiple learning steps (Siems et al. 2025). They. however, still suffer from performance drop in long context, mainly due to the less expressive memory architectures (Behrouz
ct al. 2025c).
Deep Memory Module: Titans and Test Time Training. To overcome the limited memory and to extend the effective context length of deep sequence models, more recent studies focus on a new gen- eration of architectures with deep memory module (Behrouz et al.) 2025c Sun et al. 2024). These architectures are built on the meta-learning perspective, where the memory is an MLP architecture that is updated using gradient descent (with momentum) (Behrouz et al. 2025c Sun et al.) 2024). [unn et al. (2024) further provide a unifying perspective that how linear and softmax attention are respectively parametric and non-parameteric solutions of (kernel) regression loss but consider other
17

---

Published as a conference paper at ICLR 2026
modern linear RNNs outside of this class of models. Recently, in concurrent work ours, Wang ct al. (2025) show that with additional simplification of modern RNNs (e.g., RetNet (Sun et al, 2023 Mamba (Dao & Gu 2024) they approximately place in the same class of models that inter- nally optimize regression loss. It, however, still remains unanswered that "What is the underlying design framework behind these sequence models that can accurately unify existing architectures?" Moreover, the role of forget gates and its alternative choices in modern sequence models is surpris-
ingly less explored.
To clarify the relationships among existing architectures, several recent works have sought unifying perspectives. Liu et al. (2024) adopt an online-learner view, closely aligned with our (Learning- Retaining Viewpoint) while the concurrent work of Wang al. (2025) frames the problem as online regression, which corresponds to our (FTRL Viewpoint)
Our approach formally links these two viewpoints (Theorem 2.2). Unlike those studies, which restrict themselves to l2 and dot- product losses, MIRAS extends beyond these standard choices: it supports non- oon-Euclidean losses and regularizations, enabling new architectural designs. This makes MIRAS a comprehensive framework that both (i) explicitly interprets retention/forget gates forms of regularization and (ii) generalizes the attentional-bias objective beyond simple regression losses.
B PROOF OF THEOREM 2.2
Here we present the proof of Theorem 2.2 For the sake of completeness, let us first re-state this Proposition.
Theorem 2.2 Let 1)+ = 17 and define he (W) E (W;; k, v;) + 1 R(W). Assume W R
"
and the function (w,) is strictly convex in W and let (.. be the Bregman divergence defined by function ln-), i.e., (W,
h(W)- h(W"") (Vh(W") W W") Set Ret,(W, W")
D (R,,, (") and (W; xx) = (W; xx) in Searning-Retainiing Viewpoint Then, the update rule
in Learning-Retainig Viewpoint is equivalent to the update rule in FTRL Viewpoint)
Proof. Let { . be the sequence of parameters obtained by
FTRL Viewpoint) and
{W1, W..... be the sequence of parameters obtained by Geearning-Retainiigg Viewpoint) To show both update rules are equivalent, it suffices to show that the above two sequences are the same if they are initialized at the same point. We prove this statement by induction. First of all, since both sequences are initialized at the same point, the induction base satisfied (i.e. Wi W, Now, assume by induction hypothesis that
W = WT
(14)
To complete the induction, we need to show W. W., To this end, notice that, by Learning Retaining Viewpointh. we have
We = arg min W
(WW, ki 0,) + Ret, (W, W.-1)
Using the choice of the Attentional Bias and the Retention function in the Proposition, we obtain
Σ (, ky v.) + 1 R[(W) - [-1
W. = arg min W
&W ky, vz) +
Σ (WW v.)
(15)
1 R(W,
Wi ki, 0;) + V.(...)), W - We-1
Ignoring the constant terms and using the induction hypothesis equation 14
we get
W. = arg min W
(W kr, v.) (, ki, vz) + R(W)
(16)
Σ
k., v;) + 1 VR((W..), W- -W-11
n
18

---

Published as a conference paper at ICLR 2026
On the other hand, recall that {W1, W2,
is obtained by (TRLL Viewpoint)
Therefore, we have
W arg min
((W; k; uz) + RWW)
Thus, we have
[-1
Σ (W ki, u.)
VR(W--1)=0 =
(17)
"
Combining equation 17
and equation 16 we obtain
D...k. v.) R(W).
W. arg min
W
71
This implies W₂ We which completes the proof.
C VIEWING TITANS AS (UNIVERSAL VIEWPOINT)
Here we diseuss how Titans in BBerrouz et (2025c) can viewed as a special instantiation of the Universal Viewpointh. Let (k, (0, 1). Set
&(W; (ks, vi) (W W Ve(W,-I, vt))
and
)I²
W (1- W -mo -o(1-a.)
Ret, (W, (W-WWk
20,
in (Universal Viewpoint...
Then, it is not hard to verify that the update rule for W, can be given
as
W, = (1 az + W n1(1 o..) O.W. vo).
This dynamics is equivalent to
W; = - (1-a))
S S1 - Ve
which is essentially the gradient descent update with momentum used in Titans of Behrouz et al. 2025ch.
D UNIFYING VARIOUS EXISTING METHODS UNDER MIRAS FRAMEWORK
In this section, we discuss how various existing architectures fit into MIRAS framework. To facilitate the discussion, we recall Figure for comprehensive presentation of MIRAS.
Next, we discuss how various existing architectures can be unified under MIRAS.
RNNs with Hebbian Rule. The first generation of modern recurrent architectures (e.g.. Linear attention (Katharopoulos et al., 2020), RetNet (Sun et al., 2023), Mamba (Gu & Dao 2024). and GLA Yang et al. (2024b)) are based on Hebbian-like (e.g.. gated Hebbian) learning rule (Hebb, 2005 We let attentional bias be the dot product similarity. That is, given a memory M E 2
and k, v E Rdd, we define & := 2(M,,,, vi) and local retention as (M, |IM aM,-1l1 Using Equation Learning-Resn Retaining Viewpoint and gradient descent the optimizer (i.e., memory learning algorithm), the memory update rule is:
M₁ = -aMMM + vek,
(18)
When (1) a = 1, memory update is equivalent to Linear Attention (LA) (Kaharrppoulos et al. 2020 ;
(2) IR is learnale parameter, resulting architecture is either lightening attention (n 1) (Li
et al. 2025) or RetNet (n = 1) (Sun et al. (023); and (3) a E IR are data-dependent learnable parameters, resulting sequence model is Mamba2 (Dao & Gu 2024
19

---

Published as a conference paper at ICLR 2026
Associative Memory
Memory Architecture
Attentional Bias
Retention Gate
Memory Algorithm
neural architeeture that
The memory Internal objective.
The retain the the memory.
The algortthm that learns the magpingg
Vector
Begression Loss
fy Regularizaaion (Local or Clobal)
Cradiemt (CI)
2. Matrix
2. Dot Product Similarity
2. Elastic Not Regularization
2. with Momeotum
3. Multilayue Puresptroa (MLP)
S. Haber Lose
XL Divwrgenne
3. Newton's as Method
Mamory
Xi-Diverpence
Dregman Divenjence
Non-parametric Solutions.
Associative Memory is neural network that learns to map keys to values based on an Attentional Blas objective.
Gradient Deapmnt
Ret, (Reuiin Past Momoctesi
Unttrrnal Objectoon)
Input
Memory
Attentional Blas
Retention Gate
Learning
Effect of Choices Output
Figure 4: The overview of MIRAS framework. MIRAS based on four critical choices of (I) memory architecture, (2) attentional bias, (3) retention gate, and (4) memory learning algorithm. In this framework, the memory architecture determines the model capacity to memorize; attentional bias is responsible for modeling the underlying mapping patterns; retention gate determines how to balance learning new concepts and the retention of previously learned concepts; and memory learning algorithm is responsible for memory management.
RNNs with Delta Rule. To improve the memory management and to enhance the memory capacity
of the above group, several studies suggest using delta rule (Schlag et al. 202I) as the learning
algorithm in recurrent neural networks (e.g.. DeltaNet (Schlag et al.] [022 Longhorn (Liu et al.
2024a), and RWKV7 (Peng et al. 2025b)). In this part, we recall that where M E Rdx delta rule
equivalent to optimizing MSE objective [lM,k, vellž with RM,
[_ a.M,-1l1
as local retention, and stochastic gradient descent as optimizer: (7) is defined in Equation Learning Retaining Viewpointh
Mr = a (I-1kk, M. v.
(19)
When ( = 1, memory update is equivalent to DeltaNet (Schlag et al. 2021 ; and (2) ag R are data-dependent learnable parameters, resulting sequence model Gated DeltaNet (Yang al,, 2024a) (when m = 1). Therefore, RNNs with delta rule are special instances of MIRAS.
Beyond Delta Rule. As discussed earlier, while delta rule with its value replacement strategy is more powerful than Hebbian-like learning rules, it suffers from theoretical limitations (Irie et al.) 2023) and achieves moderate performance in practice (Yang al. 2024c). Therefore, several studies have focused on update rules beyond delta rule. Recently, Titans (Behrouz et al.) (202cc) suggests using non-linear MSE objective of (k,) vellž with both local and global retention of D2 |W,-W,,1 1. and G, = || I|2 and optimize with gradient descent with momentum Therefore, Titans-LMM is special instance MIRAS, where we use the abovementioned attentional bias and retention regularizations, and gradient descent with momentum as the optimizer. Another way to obtain Titans under MIRAS explained in Appendix C
Another example of such models is Mesa-layer (Von Oswald et al.) 2023 von Oswald et al. 2025
in which the model uses Ei= [[.)((k)) vill3 as the attentional bias objective with M| as the retention regularization. Since these models uses Newton's method optimize such an objective, they provide a more expressive update rule than delta rule. We further discuss set of new learning algorithms beyond delta rule in Section 4
"The retention gate (forget gate) in Titans different from Mamba2 and Gated DeltaNet that we discussed above. The main difference comes from the case of full memory erase. While Mamba2 gating removes the entire memory and treats the next token as the first ever seen data, Titans use a "cold start" strategy and use the previous state of the memory to measure the surprise of the incoming token before fully erasing the memory.
20

---

Published as a conference paper at ICLR 2026
Attention. As discussed by Sun et al. (2024), softmax attention is non con-aarmettriic solution of b2- MSE loss function (i.e., [W - v with Nadaraya-Wart Watson estimator. Therefore, softmax attention (i.e.. Transformers) is instance of MIRAS, when we find the non-parameteric solution to the MSE loss with Nadaraya-Wand Watson estimator, without retention.
All in all, as illustrated in Table many existing methods can be unified under MIRAS.
E BEYOND EXISTING ATTENTIONAL BIASES AND RETENTION GATES
Here we provide the details of the alternative attentional biases and retention gates discussed in Section 4, We first propose several novel possible choices of attentional biases and then we discuss novel choices for retention gate.
E.1 ALTERNATIVE ATTENTIONAL BIASES
Variant 1: Ep-Attentional Bias. As discussed in the main body, attentional bias defines the "sim- ilarity metric" and measures how well memory can recall the value, given its corresponding key. Although & regression loss often is a natural choice, it is sensitive to noise in the data. A natural extension is to use Ey-norm class of objectives. That is, let M be the memory, & be the keys, and be the values, we define lp Ep-attentional bias as:
((M(W,k)):v)) = (IM())) vrlp.
(20)
where p I221 and I|l1pp is the p-norm. Although depending on the distribution of the data, we might want to use different values of (see Section 5), different values of p can result in memory architectures with interesting properties. For the sake of simplicity, let memory be a matrix, i.e., W E Rmd and M(W, k2) Wk,, the closed form can be derived as:
W. W 7/ Ve(W, k, v.) W 11-1-P1 p) (Sign(W,,1ks vv) O W2
velp-1) k
Let = 1, the recurrence is simplified as:
We = W - Th Sign(W. - Vz) k.
(21)
which means that the memory has only two values of -1 and 1. We call this variation value-less associative memory, in which we store entities (keys) but map them into two extreme class of -1 and +1.
One of the critical challenges to use the above update rule is in the backpropagation process, in which Sign(.) and 1 are non-differentiable and so might cause unstable training. To overcome this issue, use Sign(x:) taanh (ax), and |2|| v²2 c, as the smooth approximators of these functions.
Remark 1.
One simple interpretation for such behavior (i.e., value-less memory) is similar to the coping mech- anism in humans (Lofius) 1993), in which the memory does not store the values for extreme events. This interpretation of protective memory in extreme events motivates our next variant.
Variant 2: Huber Loss: Memory with Coping Mechanism. While mnrmm objective is a com- mon choice for many statistical and machine learning tasks, it is known to be sensitive to outliers and extreme samples. This sensitivity extends the use of Eg loss for attentional bias. To address this and drawing motivation from robust regression literature, we suggest utilizing the Huber loss- type (Huber) 1992, Hastie et al. 2009) as the attentional bias, thereby reducing the negative impac of the outlier data on the memory learning process.
We can apply Huber-type loss in three different ways: The first approach to define the summation of the Huber loss across different coordinates as the total loss, i.c.,
E(W; kr, v.) Σ (H(M(W,)))) ve.)),
where M(W, kt)y and Ve.j denote the j-th coordinate of M(W, kt) and v. respectively. The func- tion H(-): : IR + RR is the Huber loss defined
H(a) = {
t
if |a| < 8 a > 8.
(22)
8
-1al-10)
- (0)
21

---

Published as a conference paper at ICLR 2026
Utilizing this attentional bias can lead to various memory update rules. For example, for the matrix form memory M(W, k.) = Wkt, the update rule given by
W = Wt-1 ?1t (IVE: - vo)kk) @ (I( - (4l56))1
+ (6,Sign(Wk;; - -v)) (I(WKk v.l > 6,)1
(23)
In this formulation, the parameter ó, decides the type of the memory used for each block of mem- ory (62-norm objective value-less) based on the context, making the memory more robust to outliers.
The second approach is define the Huber-type loss based on the & loss over all coordinates, i.e,
((W; kr v.) = H([M(WW, ke) vellz).
For simplicity of derivations, assume matrix memory M(W, kt) = Wke Then using gradient descent for updating memory leads the memory update rule
{ [M kb) vz) koT if [IM(W,-1, - vell2 S de. (MCW, k.) i v.)
W = W 1)
(24)
k.
Otherwise.
[MM k.)
Again, in the form equation 24 the parameter ó, decides the type of the memory used (o2-norm
objective or normalized version) based on the context, making the memory more robust to out- liers.
Finally, in the third approach, we present a smooth mixture method, in which the memory decides for incoming data it is better to use & or & attentional bias:
ne Vlz (W v,, vz). me & V kr v)) Otherwise.
if [[(()))) vell S 6,,
W W.
(25)
The role of parameter 5, is the same as above.
Variant 3: Memory Robust to Value Shifts. Following the robustness requirement discussed in the previous section, we aim to design a memory mechanism that exhibits resilience against small shifts in the value parameter. A natural approach in this context is to employ robust optimization formulation. Specifically, we define the loss function as the worst-case &g distance between the predicted memory output and the perturbed true value:
L(M(W, k.): = max
[IM - (v. + 8v.)
(26)
[|ōv, ISS
This formulation seeks the memory parameters W that perform well even under the adverse local perturbation of the true value v, within la ball of radius A. To solve the maximization problem in equation 26 we find the optimal perturbation dv; By solving this problem with respect to dv., we arrive at:
8v A -M(W, k.) + V. [M(W, k.) vellz
Substituting this optimal perturbation back into the loss function equation 26 we obtain the robust loss:
C(M(W, k.) = I.M(W, v. + AIIM(W,k)) k v2l12 + 2 A²
This robust loss function is combination of the standard Eg loss and a term proportional to the & norm of the error, scaled by the robustness parameter A. The value of A thus controls the trade-off between fitting the nominal data and ensuring robustness against value perturbations.
For simplicity of the derivations, let us consider a constant value for A, an Euclidean retention gate Ret, (W, W,-1) = -l Wl, and an attentional bias term i(W; v,) = (W Wi-1, VW(W,-;; ke, vz)). Furthermore, simplify the memory operation, assume linear
22

---

Published as a conference paper at ICLR 2026
matrix memory model M(W, k,) = Wk, Under these assumptions, derive the memory update mechanism using gradient descent on the robust loss:
W (MM(W--, -v) + A M(WI-1,ku)
Wi Wi-1
k
|MW,,- Vrll2
In this update rule, the parameter A, which governs the influence of the robustness term, can also be treated as a learnable parameter, allowing the model to adapt its robustness based on the observed data.
E.2 ALTERNATIVE RETENTION GATES AND MEMORY STABILITY
Variant 4: Memorization Over A Scaled Probability Simplex Via f-Divergence. A common technique in learning to prevent numerical instabilities and exploding values is to restrict the search space to a bounded domain. Following this principle, to avoid numerical instabilities, we can con- strained the variable W, to lie within (scaled) probability simplex. In other words, the state to lie the constraint set
restrict
W (W | c and W5z O, Vj.
In this set, each matrix W can be viewed as measure. Thus, in [Larning- Rettinng Viewpoint we can utilize divergences over measures to define our premetric. For example, we can use f- divergence measure (Polyanskiy & Wu 2025 Def 4.9), (Csiszar 1967) to define D.(-,)). More specifically, let f() be a smooth strictly convex function from R to I with f(1) = Then, we can define the f- divergence between W and W" as
W
D,WW, = Σ W f Il
W
It is known that f-divergence is zero if and only if W Theorem 2.3). Using the above premetric as the retention gate and setting E(W; kG, Vo) Learning-Retainngg Retaining Viewpoint). we get the update rule
W (Polyanskiy & Wu 2025
(W- - Wi-1, Ve(W(11; kg, v.)) in Learning
W. = W.-1 3 Gg - Ve(W,-1 k2, Vz)).
(27)
Here g(-) is the inverse of the mapping f', i.e. ((f(()) Hadamard (elementwise) product, and Çe should be chosen such that |Welll c. Notice that since the function f(-) is strictly convex and smooth, its derivative is strictly increasing and hence g(-) is well defined. Conversely, for any strictly monotone function g(.). we can find its inverse function g (which is strictly increasing) and define f(r) = const + - g
Vr; the operator O denotes the
( (r')drr' The term const should be chosen such that f(1) = Then the update rule in equation 27 can be interpreted by the f-divergence regularization, explained above. Therefore, one can directly choose continuous monotonically increasing function g(-) and use equation 27 for memory update.
Specializing to KL divergence. Let us further make the above update rule explicit by using special function f. we choose f(r) = 7 In()), then the f-divergence becomes the widely used KL also utilize the Shannon entropy as the global retention by regularizing deviations from uniform distribution, i.e., (W) E5 Wii log( (W,).. Combining these choices of the local and global retention gates, we obtain the overall retention gate
divergence measure DW,,, Wr-1) En Wa log ( (Iwi),) In addition,
(e, W,-
Σ Wir log jl
W (W:)
+
1
Σ
War log(Wai)
T/.
Ol.
Choosing the attentional bias Z(W; ke, vz) -W WW-1, VE(W. 11; kr, vz)) and the above reten- tion gate will lead to the update rule
W. = arg min { W,1, Ve(W,115 v.)) + 1 Σ W
W5 log
W3
+
1
Σ
W log (W5))
(We)gt
a1
(28)
s.t.
Σ 11
W c. W3 0, Vjl
(29)
23

---

Published as a conference paper at ICLR 2026
Attaching the Lagrange multiplier to the first constraint, the KKT conditions implies
(Ve(W.-1;ke vel + (n) + as
1 ) (1
log Wii)
log (WW-1)))) M = 0, Vj. (30)
where fe should be chosen such that En W Rearranging the terms and defining X = 1/a Vi I/a,t1/w we get the update rule 1/act1/m
Wr t e Softmax ((1 - A²) log(W,-1) - - Ve(W,-;;; v.))
(31)
where A (0, 1) and v/ R+ are the hyper-parameters that can be learned during training. The Softmax operator ensures that the output lies in the set W.
Notice that the above calculation is done assuming the elements of W form probability vector (normalized to c). One can easily get update rules for the case that W is a stochastic matrix (just by decomposing the problem across columns of W). In such case keys and values can each be a probability vector and W plays the stochastic matrix mapping keys to values. Moreover, we can even go beyond matrix memory and obtain update rules for other forms of parameters such as when W is neural network (or when the parameter W is normalized per slice).
Variant 5: Elastic Net Regularization: Hard and Soft Forgetting. Elastic net is a powerful and popular tool in regression analysis to balance the feature selection capabilities of LASSO (Tib- shirani 1996) and bias reduction properties of Ridge regression (Hilt & Seegrist, 1977 Hoerl & Kennard 1970p. It has been widely used in different applications due to its ability to handle high- dimensional data and mitigate the effects of multicollinearity. Given this success, natural question is what happens if we use this regularization scheme in our context.
Viewpoint> design our memorization scheme. As mentioned in (Leaaning-Retainingg Viewpoint). the loss function le(W; kr, Ve) is approximation of the original function ((.), measuring our goodness-of-fit. Regularizing this loss with elastic net regularizer, we obtain the approximation
Let start based on (Learning-Retanl
(W kg v.) = (W- W,-1, ke
with a global retention of G, (W) 1 || + [|| To fully specify the update rule of Learning-Retaiin Viewpoint) we also need to specify the premetric functions Dr(:,). For the sake of keeping the update rule simple (and parallelizable), we can choose
DWW,, Wi-1) = wlW W
These choices of the attentional bias and retention gate leads the following update rule:
W S. (AW (V k,,
(32)
where 7
B+n" ( = n,,, and S,y the soft thresholding operator, applied element- wise. For each element, this operator is defined as
nB a)
S2) sign(2) max (0, 1
In other words, for large values of z, (z) makes closer zero by 7 amount. If it already in the 7-vicinity of zero, then it makes it zero (hard forget).
Equation equation 32 can be viewed by E 1), and a hard forgetting (if is smaller than 7). The hyperparameters 7, X, and can be learned. Notice that since the shrinkage operator not differentiable, we can approximate with its smooth approximation. For example, we can use (x) |=|<arctan(s/1)
combination of soft forgetting (obtained by multiplying W
x/2
Variant 6: Elastic Net Regularization: Forgetting via Soft-thresholding, The elastic net regu- larizer can also be used in the (FRLL Viewpoint) In particular, in (TRLL Viewpoint) we can set
1 R,( = || ||
7)
24

---

Published as a conference paper at ICLR 2026
and use :x, (W ( W,- ve(W;-1;::)) Assuming initialization Wo = 0, these choices of attentional bias and retention gate leads to the update rules:
A. A,-1 - Vo)
W. = S (At)
(33)
Here (.) is the soft-thresholding operator with parameter n/a, which can be smoothly as ex- plained in Variant 1.1.
Variant 7: General L, Memory Stability. Existing work is based on the retention gate choices (W,,, Wt-1) = |W w R(W) [ However, one can choose other choices of retention gate. For example, in (FTRL Viewpoint we can choose Lg norm as the regularizer R(W). More specifically, for 1 < q 22 we can set
1
R(W) 2n(9 - 1) W11..
7
Using this retention gate and choosing (WW; kes v1) = (W W,-1, VE(W kg. ve)) in A. where أ-ب and A, =
FTRL Viewpointh leads to the update rule We
A
Ei=1 V kG, ve): see -Shalev-Shwarte et al,, 2012 Section 2.6). Here, denotes the Hadamard (esment-wise)) product and is the element-wise absolute value operator. Assuming Wo = O, this update rule can be recursively written as:
A
A = = - vo), and W.
12-2" | lb
Variant 8: Bregman Divergence Retention Gate.. Another natural choice is Bregman divergence as retention gate, leading to a mirror descent-type algorithms. In particular, given smooth strictly convex function f()) E + R, we can define the function F(W) E3 f(W Based this choice of function F, we define the Bregman divergence
D,WW, W') = F(W) - F(W") - (W', W W'
parametric function. W. Ve(W, v.))
Utilizing this retention gate and choosing (W k; vt) Viewpoint). we obtain the update rule
(W
W - W
v.) F"WW--))
Here, F' is the mapping obtained by applying f'(-) (the derivative of f) element wise to all entries of its input matrix argument. The function g is the inverse of the mapping F'(+), i.e. ((F"(W) W.
If choose f(r) r", then F'WW) becomes the identity mapping and so is g. Therefore, the above update becomes simple gradient descent with no nonlinearity involved in the update rule. However, other choices of (() introduces additional nonlinearity in 9(-), which can enhance the expressivity of our memory. For example, we can choose the function f(-) so that its derivative becomes the inverse sigmoid function, i.e., (T( = In (F= with f (0,1) + R. Since f'(-) is strictly increasing, then the function ((() (and hence F(-)) is strictly convex. Therefore, the Bregman divergence is well defined. Moreover, the inverse of the function f'(.) becomes the sigmoid function,
exp(r) with g : + (0, 1). Then, the update of the memory becomes i... g(T) = (T) (+exp(T)
W. W In ( W W. )
where a the sigmoid function operated element-wise on the entries of W, and the division operator I-W; is also performed element-wise wise. This update rule guarantees that the elements of W2 remains within the interval (0, 1).
25

---

Published as a conference paper at ICLR 2026
F PARALLELIZABLE TRAINING AND EFFICIENT IMPLEMENTATION OF MIRAS' VARIANTS
While the design of MIRAS'S variant are theoretically well-motivated, their recurrence is non-linear, potentially make their straightforward training slow for large scales. In this section, we build upon the work of Beerouu et al. (2025c); [Sun et al. (2024) to make the training parallelizable. The main idea is to divide the sequence into chunks with size (usually is 16 or 64) and calculate the gradient for all tokens in the current chunk with respect the last state of the memory in the previous chunk. That is, we use VE(Me; k2, V2) instead of VE(M,-; k, ve), where e' is the last state in the previous chunk.
Given the above trick, we can calculate all gradients at once and make the recurrence inside each chunk linear. However, to fully take advantage of accelerators, we need to reformulate the process matrix multiplication. For MONETA, for the sake of clarity, assume algorithm as Behrruz et al. (2025c) and expand the recurrence as follows:
2. We follow the same
Mr = -.MM-1 Te
Ve(MM Ki, Vs)
Br Ti 3,
= BBMo Σ
Ve(Me;ki.
(34)
where t' = mod(t, b). and B, IT; aj. For the sake of clarity, we focus the first chunk, i.e., and so l' 0, and explain the process for the case that M, W, is linear. The process for 2-layer MIPs and other chunks is similar. Using ép loss function, we have:
ve(Wo: ki v.) p (Sign(Wk: v.) O Wk. val=-1) R.
n 3
Ve(Wo; = PEg Bg Sign(Wi v.) (WoK (WoK-VP-J) VP--)KK
(35)
where E, [n
L For the sake of stablity in To] and B, is defined analogously on 3.
training. we use S Sign(x) x tanh (axx) and |z|| = VI² + e, where > 0 is small number (i.e., e Le 6). 6). As discussed before, the case that q # appears as a normalization term on the memory. Similar to Titans (Behrouz et al., 2025c) and TTT (Sun et al. 2024). we do not apply this non-linearity inside each chunk and instead use it at the end of each chunk.
The process is the same for other two variants: (1) YAAD: We calculate the gradient of both & and f2 loss and use a masking based on [[,((;;)) -vll S &g. (2) MEMORA: update has two non-linear part, i.e., softmax and log. As discussed above, we apply the softmax at the end of each chunk. Therefore, for the log function, we can calculate all the gradients of each chunk at first and then expand the recurrence with respect to the log of weights. Again, this process make the inside chunk recurrence linear and inter-chunk recurrence non-linear.
G EXPERIMENTAL SETUP
We perform experimental evaluation on the language modeling MMerity et al. 2017 Paperno et al. 2016), common-senss reasoning (Bisk et al. 2020 Zellers al. 2019 Sakaguchi al. 2021 Clark et al. 2018 2019), and long context needle-in-haysee pare our models with the state-of-the-art linear recurrent models, Transtormers, and hybrid models (recurrent + attention). More specifically we compare with Transformer++ (Touvron et al. 2023).
tasks (Hsich et al. 2024h. We com-
RetNet (Sun al. (002)). Gated Linear Attention (GLA) YYang et al.) 2024b), Mamba (Gu & Dao, 2024), Mamba2 (Dao & Gu 2024), DeltaNet (Yang et al. 2024c). TTT (Sun et al. (024). and Gated DeltaNet (Yang et al. 2024a). By "Transformer+** we refer to standard Transformer based on the LLaMa architecture (Touvron et al.) 2023), incorporating modern practices like RMSNorm, SwiGLU, and RoPE.
We train our models with training context window of size 4096 using either Fibe dataset (Penedo et al. 2024) (for LM and common-senss reasoning tasks) or C4 dataset (Raffel et al. 2020) (for scaling patterns). We use model sizes of 120M, 340M, 760M, and 1.3B parame- ters. We train small models (120M and 340M) on 15B tokens sampled from the dataset, the medium size model (760M) 30B tokens, and the large model on 100B tokens.
26

---

Published as a conference paper at ICLR 2026
Table 6: Architectural Details.
Model
Block
Dim
Head
Peak LR
Token
170M
12
3e-3
15B
350M
24
1024
1.5e-3
15B
780M
24
1536
1.25c-3
30B
69
SWICLD
Layer
SCILAG Laget
Coro
LEWE
MIKAS Lagsr
Input
Aschitecture
Variants
Hyybed Architecture
MIRAS Block Design
Figure 5: Visualization of the MIRAS' variant architecture, their hybrid counterpart with SWA, and block design of MIRAS layer.
H ADDITIONAL EXPERIMENTAL RESULTS
H.1
LANGUAGE MODELING
The full results for experiments language modeling and common-sense reasoning tasks are ported in Table 7 Similar to 1.3B scale, our models achieve higher average accuracy compared to modern recurrent models.
H.2 EFFICIENCY EVALUATIONS
In this section, we evaluate the training and inference throughput of MIRAS'S variants with state-of- the-art sequence models, including Transformers. In particular, in 8K context window, the training throughput (10 T/s) of Transformers, Mamba, DeltaNet, and Titans are 48, 33, 39, and 37, respec- tively. MIRAS'S variants of MEMORA, YAAD, and MONETA have training throughput of 34, 36, 37 (103 T/s), which is compatible and on par with sfate-of-the-art recurrent neural networks. is notable that this throughput achieved without any specially design kernel. Therefore, in summary: (1) Comparing to modern sequence models such as Mamba and DeltaNet (which also take advan- tage of optimized kernels), MIRAS'S variants show competitive speed and are fast enough to be able to be scaled to larger scales; (2) Comparing to Titans, MIRAS'S variants do not add significant computational overhead, despite they having more expressive attentional biases.
H.3 MAD BENCHMARK
Next, we evaluate our models' performance and baselines' on MAD benchmark, which is synthetic benchmark for evaluating the performance of sequence models in memorization. recall, compres- sion, and copying tasks (Poli et al. 2024). The results are reported in Table 8 All MIRAS'S variants achieve higher accuracy compared to baselines. Particularly in memorization, atively higher rate of improvements, which highlights the importance of going beyond conventional attentional biases.
models show rel-
H.4 IN-CONTEXT RETRIEVAL TASK
In this section, we evaluate performance of MIRAS'S variants and baselines on in-context re- call tasks, which is one of the most challenging benchmarks for recurrent neural networks. In this section, we follow Arora et al. (2024) and evaluate the models on SWDE (Lockard et al. 2019), NQ (Kwiatkowski et al. 2019). DROP (Dua al. 2019), FDA (Arora al. 2023), SQUAD (Ra- jpurkar et al.) 2016).. and TQA (Kembhavi et al. (017). The results are reported in Table 9 Trans- formers still achieve the best results, outperforming all the recurrent models in in-context recall
27

---

Published as a conference paper at ICLR 2026
Table 7: Performance of MIRAS' variants and recurrent- and Transformer-based baselines on lan- guage modeling and common-sense reasoning tasks. Hybrid models are marked with The best results are highlighted.
Model
Wiki. pel
LMR.
LMB. ace
PIQA acc
Hella. acc_n
Wino.
ARC-e acc
ARC-c acc_n
SIQA acc
BoolQ acc
Avg.
ppl
parnms
Transfrmerr+
31,52
41.08
30.76
62.98
34.76
50.53
45.21
24.05
36.81
58.24
42,92
RetNet
GLA
28.51
43.02
28.73
6-1005
35.96
50.00
$4.19
24.29
37.13
58.39
44.09
Mamba
30.83
63.79
49.24
60.07
43.59
DellaNet
28.65
47.30
28.43
63.52
35.95
49.63
52.68
25.37
37.96
58.79
44.04
TTT
27.44
34.19
30.06
63.97
35.71
50.08
53.01
26.11
37.32
59.83
44.51
Gated DeltaNet
27.01
34.11
63.08
55.28
26.77
59.54
MONEIA (ours)
26.19
29:31
35.70
63.99
39.23
52.04
55.96
27.
37.29
60.22
46.44
(ouns)
60.29
MEMORA (ours)
27.16
30.44
33.68
65.21
39.17
51.23
$3.40
27.99
34.1
$9.29
45.51
760M params 30B
Transformer++
25.21
27.64
35.78
66.92
42.19
51.95
60.338
32.46
39.51
60.37
48.69
RetNet
52.09
63.
32.78
57.92
48.46
Mamba2
22.94
28.37
33.54
67.90
42.71
49:77
63.48
31.09
40.06
58.
48.34
DeltaNet
24.37
24.60
37.06
66.93
41.98
50.65
64.87
31.39
39,88
59.02
48.97
24.17
23.51
34.74
67.25
43.92
50.99
64.53
33.81
16
59.58
47.32
Gated Debaaeet
21.18
22.09
35.54
68.01
44.95
50.73
66.87
33.09
39.21
59. 14
49.69
Sambu"
51.08
Gated DeltaNet--H2"
19.88
20.83
39.18
68.95
48.22
52,57
67.01
35.49
39.39
61.11
51.49
MONETA
53.01
36.09
63.18
YAAD (ours)
20.99
21.57
37.85
69.
50.02
53.93
61.78
36.27
41.01
63.34
53.98
MEMORA (ours)
22.28
22.31
38.19
67.82
49.30
53.28
63.57
36.15
40.94
62.96
51.52
1.3B params 100B tokens
Transfomer++
18.53
18.32
42.60
70.02
50,23
53.51
68.83
35.10
40.66
57.09
52.25
RetNet
19.08
17.27
40.52
70.07
49.16
54.14
67.34
33.78
40.78
60.39
$2.02
Mamba2
16.56
12.56
45.66
71.87
55.67
55.24
72.47
37.88
40.20
60.13
54.89
DeltaNet
70.72
55.29
Gated DeluaNet
16.42
12.17
46.65
72.25
55.76
57.45
71.21
38.39
40.63
60.24
55.32
Samba"
13.29
70.94
55.56
36.17
62.11
Gated DeltaNet-H22
15.91
12.55
48.76
12.19
56.88
57.77
71.33
39.07
41.91
61.55
56.18
MONKTA (ours)
15.52
11.47
47.88
73.16
56.14
59:09
72.53
40.32
41.91
61.18
56.52
YAAD (ouns)
15.18
11.89
47.23
72.81
56.46
59.02
72.14
40.05
40.73
61.86
56.39
MEMORA (ours)
15,90
12.04
48.67
73.10
55.99
57.36
71.55
37.92
40.19
61.34
55.87
Table Performance MIRAS' variants, and baselines on the synthetic benchmark of MAD (Poli et al. 2024). Our models achieve higher accuracy compared to all the baselines, including Trans- formers.
Compression
(Noisy) ICR
Fuzzy ICR
Selective Copying
Memorization
Average
Transfommrs
49.4
100
48.2
95.9
83.8
75.46
Gated DeltaNet
44,8
100
32.5
96.2
81.7
71.04
Tilass
49.6
100
49.7
99.4
83.5
76.44
YAAD (ours)
$1.9
100
50.2
99.6
85.7
77.28
MONETA (surs)
51.1
100
48.9
99.6
85.4
77.00
MEMORA (ours)
50.5
100
48.7
99.6
85.1
76.78
tasks. Our variants of MIRAS, however, show competitive performance and improve the gap of recurrent models with Transformers.
The performance of MIRAS' variants compared to baselines. While still Transformers achieve the best results in in-context recall tasks, our design of more expressive attentional bias can potentially reduce the performance gap with Transformers in future.
Table 9:
SWDE
NQ
DROP
FDA
SQUAD
TQA
Average
Transformers
84.9
23.0
28.4
72.5
48.1
64.4
53.55
Gated DeltaNet
19.1
33.4
39.6
59.7
40.28
Titans
65.1
20.7
27.2
37.3
42.6
61.0
42.31
YAAD (ours)
66.2
20.9
27.2
38.1
42.7
61.3
42,73
MEMORA (ours)
65.5
20.5
26.9
38.2
43.0
61.2
42.55
MONETA (ours)
64.9
20.7
27.1
37.9
42.5
61.0
42.35
28


Reviews:
[{"review_id": "FzOtBwUZTZ", "reviewer": "Reviewer_xTsB", "summary": "This paper proposes a unified online optimization view that interprets mainstream sequence models as associative memory updates jointly driven by attentional bias and retention, and it treats forget gates as a form of regularization. Building on this framework, the authors move beyond dot product and Euclidean objectives to introduce three attention free and parallelizable architectures that outperform strong baselines on language modeling and long context retrieval while showing greater robustness.", "questions": "1. The paper provides sensitivity curves or heat maps such as p, q, Huber cutoff thresholds, retention weights, learning rates and step sizes for performance and stability. Is there a robust default configuration or an automatic parameter tuning strategy?\n\n2. How does the framework perform on multimodal long sequences? Is task-specific bias/retention configuration required to earn benefits?\n\n3. Can MIRAS be combined with sparse attention or local attention? Under the same computing budget, what are the trade-offs between pure MIRAS, pure attention, and hybrid solutions?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. This paper proposes a unified online optimization perspective, taking \"attention bias - retention\" as the core paradigm, and systematically reinterpret the forgetting gate as regularization. Further breaking away from the reliance on dot product and binary norm, introducing robust loss and divergence is conceptually enlightening and has promotional value.\n\n2. On benchmarks such as language modeling and long context retrieval, thorough comparisons were made with strong linear RNN baselines to stably obtain gains and demonstrate robustness against noise and outliers. While maintaining O(1) memory overhead and good throughput at the training end, the experimental setup and indicators are relatively comprehensive.", "weaknesses": "1. The comparison mainly focuses on several linear RNNS and long context benchmarks, and has not yet made a \"same-scale, same-training budget\" head-to-head comparison with the current first-line inattentional/near-attentional models. It is recommended to supplement the system comparison under the same number of parameters, the same number of training steps and data quota, and report the significance and confidence intervals.\n\n2. Although various forms of bias and retention are presented, there is a lack of strict item-by-item ablation to quantify the independent contribution of each design decision. It is recommended to incorporate grid ablation, negative control (reverting to the two-norm/dot product), and cross-dataset reproduction to verify the robustness of the conclusion.\n\n3. It is enlightening to link \"learning-retention\" with FTRL, but when memory is non-convex (such as with MLP), the guarantee of convergence and stability is not clear. It is suggested to provide applicable conditions, failure cases and stability diagnosis", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0gmJnBHq3E", "reviewer": "Reviewer_6YPG", "summary": "This paper proposes MIRAS, a framework that reinterprets sequence models with memory as associative-memory modules governed by online optimization. The framework introduces two key components: (i) Attentional Bias (internal learning objective) for capturing new data, and (ii) Retention (regularization term) for retaining old knowledge. The paper then investigates various loss functions and regularizers (e.g., ℓp, Huber, KL, f-divergences), resulting in three derived architectures (MONETA, MEMORA, and YAAD) claiming improved robustness, stability, and scalability compared to existing linear RNNs and attention-based models.", "questions": "- Minor formatting error: lack of spaces between sentences, e.g, in L20 and 338\n- The framework in Fig.1 introduces 4 components, but it looks like only 2 are examined. Have you tried different memory architectures and memory algorithms?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The paper presents a clear and systematic formalization of memory update in sequence models within a unified optimization-based lens.\n- The connection between the Learning–Retaining and FTRL viewpoints is well-written and could be theoretically valuable.\n- Empirical evaluations contain multiple model variants and tasks.", "weaknesses": "- The paper lacks conceptual novelty. The proposed MIRAS framework primarily rephrases existing viewpoints (e.g., online optimization, meta-learning, and energy-based associative memory) under new terminology. Sec. 2.1, 2.2, and 2.3 review the known formulation as cited in the paper. Sec. 2.3 introduces an \"alternative\" view which considers the entire history of sequences, but basically presented before in concrete form in [1, 2]. \n- The proposed “alternative attentional bias” and “retention gates” correspond to standard online learning formulations with user-defined approximations and regularizers. There are no clear insights/proofs why some choices are better than others.\n- The experimental evaluation covers only a small subset of recent linear RNNs and attention models, with little explanation for the chosen baselines. Comparisons with stronger and more recent models, such as attention-based architectures like Titan [3] (cited by the paper) or recurrent approaches like SHM [4], would be necessary to support the claimed advantages.\n- Many reported improvements (0.3–1.0 points) in Table 1 are small and potentially within noise margins, especially since no variance or statistical significance is reported. Table 2 shows better improvement, but the set of baselines is limited. Ablation study, the selection of p=3,q=4 as “optimal” appears arbitrary, and the paper does not provide insight into why higher-order norms should improve robustness beyond qualitative intuition.\n\n[1] Munkhdalai, Tsendsuren, Alessandro Sordoni, Tong Wang, and Adam Trischler. \"Metalearned neural memory.\" Advances in Neural Information Processing Systems 32 (2019).\n\n[2] Bartunov, Sergey, Jack Rae, Simon Osindero, and Timothy Lillicrap. \"Meta-Learning Deep Energy-Based Memory Models.\" In International Conference on Learning Representations, 2019.\n\n[3] Behrouz, Ali, Peilin Zhong, and Vahab Mirrokni. \"Titans: Learning to memorize at test time.\" arXiv preprint arXiv:2501.00663 (2024).\n\n[4] Le, Hung, Dung Nguyen, Kien Do, Sunil Gupta, and Svetha Venkatesh. \"Stable Hadamard Memory: Revitalizing Memory-Augmented Agents for Reinforcement Learning.\" In The Thirteenth International Conference on Learning Representations, 2025.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "i8pYVoxgUP", "reviewer": "Reviewer_suGh", "summary": "The purpose of this paper is to unify prior work on the optimization and regularization of Transformers and recurrent networks with associative memory and to use the insight gained thereby to improve on the state-of-the-art linear RNNs. The authors create a formal theory that synthesizes disparate viewpoints, design a generic framework – based in part on the aforementioned formal theory – for the architecture of neural language models, and use the framework to guide the design of candidate linear RNNs. The results reported indicate that the authors' candidate networks are competitive with or superior to existing linear RNNs, albeit far from competitive with the conventional Transformer w.r.t. in-context recall, as reported in Appendix H.4.", "questions": "1. Given the substantial gap between Transformer and linear RNNs and Miras-based models on the in-context retrieval task (see Table 9), would you characterize the theoretical and empirical work presented here as providing some improvements over linear RNNs without the promise of bridging the gap with Transformers on this task?\n1. Are all results of single or multiple runs?\n1. Which dataset and metric are used in the ablations (see Tables 3 and 4)?\n1. Some of the results are for \"Transformer++\" and some (in the appendices) are for \"Transformers\". By \"Transformer++\", do you mean the enhanced Transformer named \"Transformer++\" so-named in [Gu and Dao 2023](https://arxiv.org/abs/2312.00752)? Since [Thapak & Hore 2020](https://arxiv.org/abs/2003.04974) also uses the name \"Transformer++\", the paper's readability and future value would be improved by being explicit what is meant by \"Transformer++\". As it stands, the citation for \"Transformer++\" is a LLaMa paper which doesn't explicitly articulate the definition of \"Transformer++\".\n1. In Table 1, the most competitive model vis-a-vis Moneta is Gated DeltaNet-H2. Why is Gated DeltaNet-H2 not shown in Figure 2? Is there something about Gated DeltaNet-H2 that prevents it from being compared to Moneta in the setting of Figure 2?\n1. Where is the code? If there's already a link to the code somewhere in the paper, this reviewer would much prefer it being clearly presented early in the paper.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "This paper has many positive attributes. It is written quite well (some typos noted below), presents a unifying framework of the optimization and regularization of neural language models, proposes a number of architectural variants based on the framework, and contains ample experimental results and ablations. The overall results – reported in the main body of the paper and in the appendix – are competitive or superior, especially needle-in-the-haystack.", "weaknesses": "* Given the scope of this work, this reviewer believes the presentation would be improved by the inclusion of a brief limitations section in the main body of the paper. The results of Table 9 in Appendix H.4 should be mentioned there.\n* Code does not appear to be provided. This reviewer considers the lack of code to be a blocker to acceptance.\n* It’s not clear whether the reported results are from single or multiple runs. The lack of reporting of any sort of variance across runs leads this reviewer to believe the reported results are for single runs.\n* The section about ablations is not clear about which dataset and metric are used during the ablations (see Tables 3 and 4).\n* In general, a succinct and comprehensive performance comparison accounting for latency and memory at both train- and test-time -- covering all models used in every evaluation -- would be useful in the main body of the paper.\n\nTypos:\n* Line 20: there's no space between \"loss\" and \"Miras\".\n* Line 185: extra space before comma after \"(Learning-Retaining Viewpoint)\".\n* Line 266: \"As discussed in the previous section, existing work focus\" -> (\"As discussed in the previous section, existing work has focused\" or \"As discussed in the previous section, existing work focuses\")\n* Line 426: \"use needle-in-haystack task\" -> \"use the need-in-a-haystack task\"\n\nIn one focused reading of the paper, I noticed quite a few additional typos -- missing articles, subject-verb disagreements, and singular-plural errors -- elsewhere in the manuscript. I suggest the authors perform a close reading of their own to find them.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CVIxkWaLXl", "reviewer": "Reviewer_VGJG", "summary": "This study proposed MIRAS, a unifying framework that reinterprets sequence models like Transformers and RNNs as associative memory modules governed by online optimization. MIRAS decomposes model behavior into two parts, i.e., Attentional Bias and Retention. The authors propose a broader design space by generalizing beyond various losses. To justify the effectiveness of MIRAS, authors conduct experiments on language modeling, commonsense reasoning, and long-context tasks.", "questions": "I have provided my concerns in the Weaknesses part.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. MIRAS provides a new perspective to view and design sequence models through the lens of online optimization and associative memory.\n\n2. MIRAS achieves great performance compared to baselines, especially on long-context tasks.", "weaknesses": "1. The training and inference speed and GPU memory of MIRAS are not reported. Can the efficiency be maintained on larger models, e.g., 7B+ parameters?\n2. More ablation studies are needed to justify the effectiveness of the proposed method, such as the impact of different p/q values and different memory algorithms.\n3. The significance test is not reported in the experiments.\n4. MIRAS achieves inferior performance in some tasks (see Table 1, Table 7, and Table 9). Thus, an in-depth analysis should be provided in such cases.\n5. The proposed models, i.e., MONETA, YAAD, and MEMORA, are the incremental combinations of existing ideas, e.g., MLP memory, Huber loss, and softmax gating, rather than fundamentally novel architectures.\n6. The implementation code is not provided to check and reproduce.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)