Atlas: Learning to Optimally Memorize the Context at Test Time

Associative memory—the ability to map different entities or events—is an inseparable component of learning in humans (Terry, 2017) and so has motivated several recent studies to understand the state-of-the-art deep learning architectures through its lens (Ramsauer et al., 2021; Behrouz et al., 2024; Behrouz et al., 2025; Wang et al., 2025). In this perspective, memory is defined as a neural update caused by an input; the more surprising the input is, the more it affects the memory and so is memorable. Therefore, finding an effective “surprise metric” is a critical step towards designing such memory modules. As earlier discussed by Behrouz et al. (2025); Behrouz et al. (2024), almost all existing architectures use a surprise metric that updates the memory based on the current input. An event (as a sequence of tokens), however, might not consistently be surprising through a long-period of time although it is memorable. To overcome this issue, Behrouz et al. (2024) suggest breaking the surprise metric into two parts of “momentary” and “past” surprise, incorporating the cumulative surprise of past inputs when updating the memory with respect to the current input. This design, however, can miss the context by memorizing individual tokens. To this end, in this work, we present a long-term neural memory module that measures the surprise of a local (or global) context window, meaning that it learns how to memorize the (token) context at test time.

Through the paper, we use the terminology “Test Time Memorization” because the process involves storing and retrieving information strictly within the global context, without updating the model’s core learned parameters (i.e., outer-loop) or initial states from pre-training. Typically, no persistent learning or skill acquisition carries over to new, independent global context once the memory is cleared. Thus, we prefer the use of "test time memorization" over using "test time training".

In this paper, we aim to overcome the abovementioned limitations—i.e., (1) online nature, (2) limited memory capacity, and (3) less expressive memory management—by designing a long-term neural memory module with high capacity and the ability to memorize the context, instead of tokens. We further build upon these insights and present a family of strictly more powerful Transformers. More specifically:

Better Understanding of Memory Capacity and its Bottleneck. To improve the limited memory capacity, we suggest using higher-order feature mappings (e.g., polynomial feature kernels) on input tokens. We provide theoretical justifications on why deeper memory modules and/or higher-order feature mapping can enhance memory capacity—i.e., the maximum number of linearly independent key-value associations the memory can perfectly map.

New Expressive Learning Rule. To overcome the online nature of recent recurrent models, this work presents a sliding window update rule, called Omega rule, that optimizes and updates memory based on all past tokens in a given context window, not just the last. This allows the model to better manage its fixed-size memory and memorize a local context instead of individual tokens.

Strict Generalization of Transformers. Next, we show how our Omega rule formulation connects to global and local softmax attentions (i.e., Sliding Window Attention - SWA) and present a new family of Transformer-like architectures, called DeepTransformers and its sliding window variants SWDT, that strictly generalize Transformers (Vaswani et al., 2017). We further present a novel baseline of Deep Linear Attention (DLA) to demonstrate the role of deep memory.

New Memory Modules with Better Memory Management. Building upon the above improvements, we present OmegaNet, a new architecture using polynomial features on its keys and queries, while updating its memory based on Omega and gradient descent. To further enhance memory management, we introduce Atlas, which leverages the popular Muon optimizer (Jordan et al., 2024) for updating the internal memory. We show that both OmegaNet and Atlas can take advantage of parallelizable training algorithms, resulting in fast training without substantial overhead compared to the online version (i.e., context window = 1). To the best of our knowledge, Atlas is the first parallelizable recurrent architecture that optimizes the memory using the (approximation) of second-order information (i.e., has locally optimal memory module).

Improvement on Diverse Downstream Tasks. Extensive experiments validate our model designs and proposed techniques, including ablations of modern architectures. We evaluated DeepTransformers, OmegaNet, and Atlas on diverse benchmarks—language modeling, common-sense reasoning, recall-intensive, and needle-in-haystack tasks—where they outperformed modern linear RNNs, local attention (SWA), and Transformers. Furthermore, we studied the effects of memory architecture, feature mapping, memory management algorithm (internal optimizer), and Omega rule on memory module capacity and performance in long-context understanding tasks.

Proofs, additional experimental results, discussions on related work, and the details of experiments are in Appendix.

Attention. Attention is a critical component of Transformers that acts as their associative memory (Bietti et al., 2023; Sun et al., 2024; Behrouz et al., 2025). Given input $x\in\mathbb{R}^{N\times d_{\text{in}}}$, causal attention computes output $\mathbf{y}\in\mathbb{R}^{N\times d_{\text{in}}}$ over input dependent key, value, and query matrices $\mathbf{Q}=x\mathbf{W}_{\mathbf{Q}},\mathbf{K}=x\mathbf{W}_{\mathbf{K}},\>\text{and}\>\>\mathbf{V}=x\mathbf{W}_{\mathbf{V}}$ as:

where $\mathbf{W}_{\mathbf{Q}},\mathbf{W}_{\mathbf{K}},$ and $\mathbf{W}_{\mathbf{V}}\in\mathbb{R}^{d_{\text{in}}\times d_{\text{in}}}$ are learnable parameters, and $Z_{i}={\sum_{\ell=1}^{i}\exp\left(\mathbf{q}_{i}^{\top}\mathbf{k}_{\ell}/\sqrt{d_{\text{in}}}\right)}$ is the normalization term. Despite Transformers’ simple parallelizable training and effectiveness in recall-intensive tasks (Arora et al., 2024), their generation process and long-context scaling are significant drawbacks, as attention requires at least $N\times d$ operations per token to calculate the output (see Equation 1). Therefore, in recent years, there have been an extensive research effort to design alternative architectures. We divide and review these studies into two groups: (1) Linear shallow memory recurrent models, (2) Deep memory modules:

(Linear) Recurrent Models. Linear RNNs have recently gained attention as efficient Transformer alternatives due to their parallelizable, linear-time training and comparable performance (Sun et al., 2023; Peng et al., 2023). Early modern RNN variants, often based on Hebbian (Hebb, 2005) or Delta (Widrow & Hoff, 1988) learning rules, compress data into vector-valued or matrix-valued memory (Katharopoulos et al., 2020; Sun et al., 2023; Kacham et al., 2024a; Liu et al., 2024; Schlag et al., 2021; Lim et al., 2024). Let $\mathcal{M}_{t}\in\mathbb{R}^{d\times n}$ be the memory (where $n=1$ yields vector-valued memory), and $\mathbf{k},\mathbf{v}\in\mathbb{R}^{d}$ be the keys and values (projections of input $x_{t}\in\mathbb{R}^{d}$)). A simple general formulation for such linear RNNs is:

where $\ast$ is an arbitrary associative operator and $A_{t}$ is a data-(in)dependent diagonal or low-rank plus identity matrix (Yang et al., 2024b). Despite the efficient linear recurrent nature of these models, their memory can overflow, particularly with increasing context length. Although forget gates have recently significantly improved memory management in these architectures (Sun et al., 2023; Peng et al., 2025), their memory’s expressivity remains bounded by its linear structure.

Deep Memory Module. To overcome the limited expressivity of memory and to enhance the effective context length recurrent models, recent studies focus on a new line of architectures with deep memory modules (Irie et al., 2021; Sun et al., 2024; Behrouz et al., 2024; Behrouz et al., 2025). These architectures are built on the meta-learning perspective, where the memory is a deep MLP architecture updated by gradient descent (with momentum). Recently, Behrouz et al. (2025) present a framework to accurately unifies popular sequence models as the instances of test time memorization. That is, sequence models are associative memory modules that aim to learn the underlying mapping between given keys and values by optimizing an internal memory objective, called attentional bias. This optimization is based on an iterative optimization algorithms such as gradient descent. More formally, associative memory is defined as:

Optimizing this objective using an iterative algorithm (e.g., gradient descent) results in the memory update rule. Thus, the sequence model is a meta in-context learner with two optimization levels:

1. 1.

   Inner Loop: Where parameters of the memory module are optimized (i.e., $\bm{\theta}_{\mathcal{M}}=\{W_{1},W_{2},\dots,W_{\mathcal{L}_{\mathcal{M}},\dots}\}$). In the inner optimization loop, all other parameters from the model are considered hyperparameters and are fixed and not optimized.

2. 2.

   Outer Loop: Where all other parameters of the model are optimized, such as linear projections, MLP layers, convolutions, etc.

Our terminology builds on this framework. Therefore, instead of full recurrent formulations, we describe models by their: (1) memory architecture, (2) internal objective (i.e., attentional bias), and (3) memory learning algorithm (optimizer). In most cases, models use matrix-valued memory with online gradient descent; for brevity in such instances, we refer to an architecture solely by its internal memory objective. For additional discussions and examples, see Appendix B.

As previously discussed, an effective long-term memory module should store past data abstractions in its parameters. However, with a fixed number of memory parameters, a key unanswered question remains: “what is the maximum number of uncorrelated (key, value) pairs that a model can store?” To answer this, we start with the simplest case: matrix memory, an $\ell_{2}$ regression loss as the attentional bias (i.e., $\ell(\mathcal{M}_{t};\mathbf{k}_{t},\mathbf{v}_{t})=\|\mathcal{M}_{t}(\mathbf{k}_{t})-\mathbf{v}_{t}\|^{2}_{2}$), optimized by gradient descent:

The above proposition indicates that matrix-valued memory with delta update rule has sub-linear capacity with respect to its number of parameters. This means that the number of independent patterns that can be stored in a fixed-size memory with size $M$ is strictly less than $c\times M$, for some $c\in\mathbb{R}^{+}$. Recent recurrent models suggest using deep memory modules to store the abstraction of the past into the parameters of a deep neural network (Irie et al., 2021; Sun et al., 2024; Behrouz et al., 2024; Behrouz et al., 2025). While these deep memory architectures can intuitively enhance the expressive power in modeling complex underlying mapping patterns between keys and values, it is still unclear that if they enhance the memory capacity.

This theorem indicates that deep memory not only improves representational power but also further boosts network capacity, with advantages growing with depth. However, the upper bound remains subquadratic in key and value dimensions, raising the question if a long-term memory module can achieve super-linear capacity.

As stated earlier, the dimension of $\mathbf{k}_{t}$s is crucial for increasing memory capacity. Simply increasing all key and value dimensions, however, significantly increase the number of parameters ($\mathcal{O}(d_{\text{in}})$ per each extra dimension) and memory usage, particularly with long contexts. To address this, building on methods from Kacham et al. (2024); Krotov & Hopfield (2016), we suggest using separable kernels $\sigma(x,y)=\phi(x)^{\top}\phi(y)$ for keys and queries. As an example of such kernels, we focus on polynomial kernels of degree at most $p$ to increase input dimensionality and thus network capacity. Given $p\in\mathbb{N}$, let $\phi_{p}(x)=[x^{\beta}]_{|\beta|\leq p}$ be a polynomial mapping of $x$ with degree at most $p$. We redefine the associative memory module in Definition 1 by replacing the inner objective of $\mathcal{L}(\mathcal{M}(\mathcal{K});\mathcal{V})$ with $\mathcal{L}(\mathcal{M}\left(\phi\left(\mathcal{K}\right)\right);\mathcal{V})$. This polynomial mapping enhances representational power by increasing the effective dimensionality of keys without additional parameter overhead for the input projections. Next, we discuss their effect on memory capacity, even with a single matrix-valued memory:

Beyond the above intuition, polynomial kernels are further motivated by two perspectives: (1) Approximating Softmax using Taylor series; and (2) Input feature gating. For the sake of clarity, we continue with linear memory and two popular attentional biases i.e., $\ell^{(1)}(\mathcal{M}_{t};\mathbf{k}_{t},\mathbf{v}_{t})=\langle\mathcal{M}_{t}\mathbf{k}_{t},\mathbf{v}_{t}\rangle$ and $\ell^{(2)}(\mathcal{M}_{t};\mathbf{k}_{t},\mathbf{v}_{t})=\|\mathcal{M}_{t}\phi(\mathbf{k}_{t})-\mathbf{v}_{t}\|^{2}_{2}$. The same process can be applied on other attentional objectives and deep memory modules. Optimizing these objectives using gradient descent in the inner loop results in the following recurrent formulas:

Kernel Attention Perspective for the Special Case of Hebbian Rule. The formulation for (Hebbian Rule) is equivalent to kernel linear attentions (Kacham et al., 2024a; Wang et al., 2025; Hua et al., 2022; Kasai et al., 2021; Katharopoulos et al., 2020; Arora et al., 2024). In this viewpoint, the role of $\phi(.)$ is to approximate Softmax or more accurately the exponential kernel. Since exponential kernel with normalization (i.e., Softmax) is not separable, it results in Transformers’ quadratic time and memory complexity. However, Transformers’ exponential feature map kernel ($\exp(\mathbf{q}_{i}^{\top}\mathbf{k}_{j})$) can be approximated using its Taylor series as:

Our polynomial feature map extends this approximation to a more general case of:

with learnable parameters $a_{i}\in\mathbb{R}$ initialized at $a_{i}=\frac{1}{i!}$, the polynomial kernel can be viewed as an expressive approximator of Softmax attention. This provides theoretical motivation for using polynomial kernels, especially when memory capacity is limited; i.e., with (i) linear memory and (ii) Hebbian learning rule. This intuition, however, further generalizes to more expressive cases using deep memory modules and more complex attentional biases (i.e., Eq. Delta Rule). That is, $\exp(\cdot)$ feature mapping has infinite dimension and provides a more powerful similarity measure of keys and queries (i.e., $\mathbf{q}_{i}^{\top}\mathbf{k}_{j}$); however, its computation with normalization can cause additional memory and time complexity to the model. Using polynomial kernels in architectures with deep memory and complex attentional bias can further enhance performance by approximating more powerful representations for keys-queries similarities (i.e., $\mathbf{q}_{i}^{\top}\mathbf{k}_{j}$). See Section 4 for additional discussions on exponential kernels and Transformers.

Input Gating Interpretation. Another perspective that motivates the use of polynomial features is their more expressive representational power in modeling complex functions compared to the simple case of $\phi(x)=x$. That is, the coefficients of $a_{i}$s can be seen as input feature gating, in which $a_{i}\rightarrow 0$ means excluding the feature map of $[x^{j}]_{|j|=i}$, and $a_{i}\rightarrow 1$ means retaining the corresponding feature. This is similar to the gating mechanisms of RNNs but on the input rather than the memory. This gating mechanism clearly provides a more representational power as the model can learn to set $a_{i}\rightarrow 0$ for all $i\neq 1$ and $a_{1}\rightarrow 1$, resulting in the simple case of $\phi(x)=x$.

As discussed earlier, one of the critical drawback of most existing recurrent models is their online nature, in which they optimize the inner objective (attentional bias) with respect to only the current input while retaining the previous state of the memory (Behrouz et al., 2025; Liu et al., 2024), i.e.,

where $\texttt{Ret}(\cdot,\cdot)$ is the retention gate. This online nature while making the optimization of the memory simpler and faster, can cause sub-optimal memorization of the context as memory is greedily memorize individual tokens. In a more general case, however, one can optimize the memory at each time stamp with respect to the entire context (input sequence), i.e.,

This strict global formulation generally presents two critical limitations: (1) Efficiency: One of the important advantages of recurrent architectures is their efficiency at longer context in both training and inference. Optimizing the memory with respect to all the past tokens (entire context), however, (i) causes additional optimization constraints at each memory update step, resulting in inefficiency at extremely large sequences, and (ii) requires caching the past keys and values at the test time, increasing the memory consumption; (2) Context Pruning: In large context tasks optimizing with all past tokens can cause sub-optimal performance mainly due to the context change (or irrelevant context) in the middle of the input sequence. This observation has resulted to design architectures with retention (forget) gate, enabling models to erase memory when past context is no longer needed (Sun et al., 2023; Peng et al., 2025; Behrouz et al., 2024; Behrouz et al., 2025; Yang et al., 2024a).

To address these limitations, we present a sliding window recurrent model that optimizes its attentional bias w.r.t. a window of past tokens. For a memory module $\mathcal{M}(\cdot)$ and window length $c\geq 1$, we optimize the memory internal objective as:

where $\ell(\mathcal{M};\mathbf{k}_{i},\mathbf{v}_{i})$ measures the predicted mapping for $(\mathbf{k}_{i},\mathbf{v}_{i})$ pair and $\gamma^{(t)}_{i}$ is the decay term for the effect of $i$-th token in the optimization process. Building upon this formulation, we present Omega rule, which is strictly more powerful than the popular Delta learning rule (Widrow & Hoff, 1988; Schlag et al., 2021):

Following Behrouz et al., 2025, this update rule can be extended to $q$-Omega rule (or other variants) by replacing $\ell_{2}(\cdot)$ with $\ell_{q}(\cdot)$. In the extreme cases of (1) $c=1$: the update rule becomes online (Delta rule); and (2) $c=\infty$ or context length: the update becomes global optimization w.r.t. all past tokens. In this formulation, parameters $\gamma^{(t)}_{i}\in[0,1]$ act as hard (direct) gates for the past tokens. That is, $\gamma^{(t)}_{i}\rightarrow 0$ means that the model directly prunes the optimization of $i$-th token in the local context, while $\gamma^{(t)}_{i}\rightarrow 1$ means fully incorporating the optimization of memory for $i$-th token in the local context. In our design, we use input-dependent parameters for $\gamma^{(t)}_{i}$, providing in-context pruning ability. Note that, the design of sliding window recurrence allows such flexibility as for each token we need a constant number of gates; i.e., $\{\gamma^{(t)}_{i}\}_{i=1}^{c}$. Using input-dependent gates for the global optimization (Equation 7), however, can result in significant parameter increase and memory usage, diminishing the advantages of recurrent models.

OmegaNet. We now present OmegaNet, a novel sequence model that updates its memory using Omega rule. To enhance the memory capacity of OmegaNet, we use polynomial kernels on $\mathbf{k}$s and $\mathbf{q}$s. Accordingly, optimizing the objective in Equation 9, results in an update rule of OmegaNet as:

or in the spacial case of linear memory:

From the memory perspective, Omega rule (OmegaNet) does not measure the surprise of a token, but the surprise of a local context based on the context-aware combination of individual tokens within the context.

Beyond Gradient Descent. The concept of Omega rule and “test time memorization of context” can simply be extended to optimizing the objective in Equation 9 with any arbitrary optimizer, even beyond simple gradient descent. We use two extreme cases for $c$ as the illustrations. In the first case, we let $c=1$, $\gamma^{(t)}_{i}=1$, and use gradient descent with momentum as the optimizers, resulting in the following update rule:

This update rule is equivalent to the long-term neural memory in Titans (Behrouz et al., 2024). In the second case, using a linear memory $\mathcal{M}$, letting $\gamma^{(t)}_{i}=1$, and $c$ be equal to the context length, the memory update process is equivalent to optimizing the (regularized) least-squares problem:

Von et al. (2023) suggest directly optimizing the above objective and use Sherman-Morrison formula (Sherman & Morrison, 1950) to recursively calculate the inverse term in the optimal solution. Despite the optimality of memory, such direct solutions comes with the cost of non-parallelizable training and also are limited to only the linear matrix-valued memory setup. Furthermore, as discussed earlier, the global nature without any direct hard gating terms (i.e., $\gamma^{(t)}_{i}$s) can force the model to not prune the context, damaging the performance in longer sequences.

While Omega rule provides a more general and expressive formulation for the design of memory modules than Hebbian or Delta learning rules, its applicability to large-scale models relies on its efficiency in training. To this end, we discuss a fast parallelizable training algorithms that does not add any significant computational overhead with the online counterpart version (i.e., $c=1$). A naive implementation requires materializing $c$ gradients $\nabla\ell\in\mathbb{R}^{d_{\text{in}}\times d_{\text{in}}}$, which can result in a significantly higher memory footprint and I/O cost when $d_{\text{in}}$ is large. Also, to fully utilize hardware accelerators such as TPUs and GPUs, it is important to tensorize computations and maximize the use of matmul operations. Motivated by recent work (Behrouz et al., 2024; Sun et al., 2024), we propose a simple sliding window masking strategy that supports efficient parallel training while avoiding substantial memory overhead. Specifically, we partition the input sequence with length $L$ into chunks of size $b\geq 1$, each of which is represented by $\mathbf{S}_{i}=\{\mathbf{x}_{(i-1)b+1},\dots,\mathbf{x}_{ib}\}$. Then for each chunk, we calculate the gradients with respect to the last state of the previous chunk. For the sake of clarity, we first assume $\gamma^{(t)}_{i}=\eta_{t}$ for all positions in the sequence. When the chunk size is $b=1$, the update rule is:

where $\mathcal{M}_{t}$ is the model state at step $t$, $\alpha_{t}$ and $\eta_{t}$ are the weight decay and learning rate parameters respectively, and $(\mathbf{k}_{i},\mathbf{v}_{i})$ denote the input pair at position $i$. In practice, we strike a balance between the fully recurrent form and the fully parallel form by dividing the sequence into smaller chunks. Within each chunk (intra-chunk), we apply parallel computation, while across chunks (inter-chunk), we adopt a recurrent computation scheme. We now define $t^{\prime}=t-\operatorname{mod}(t,b)$. That is, for time steps $t$ such that $t^{\prime}\leq t<t^{\prime}+b$, the update rule within each chunk becomes:

In our implementation, for $G_{t}$, we follow the same gradient computation approach as described in Titans (Behrouz et al., 2024) but additionally apply a sliding window mask $M_{s}$ during the broadcasting operation (e.g., using einsum). When $c=1$, the sliding window mask $M_{s}$ reduces to the identity matrix. For $c>1$, $M_{s}$ is an identity matrix except that the $c-1$ positions immediately preceding each diagonal entry are also set to 1. This allows gradient contributions from a window of size $c$, enabling efficient computation without materializing all gradients inside the chunk.

Connection to Sliding Window Attention. Softmax attention block can also be reformulated as a non-parametric solution to the $\ell_{2}(\cdot)$ regression with Nadaraya-Watson estimators (Zhang et al., 2022; Fan, 2018):

where $L$ is the sequence length. While this formulation optimizes the memory $\mathcal{M}$ with respect to the entire sequence length, one can limit the optimization process to the past $c$ tokens, resulting in:

which is equivalent to the sliding window attention (SWA). This connection provides an important insight on the difference of attention and recurrent models: Not only attention is a non-parametric solution (contrary to the parametric nature of recurrent models), it globally optimizes its internal objective (attentional bias), while most recent modern recurrent models are online learners (Sun et al., 2024; Yang et al., 2024; Behrouz et al., 2025; Peng et al., 2025)⁰⁰ 0 Two of the exceptions are Titans (Behrouz et al., 2024) and Mesa-layer (Von et al., 2023), where Mesa-layer optimizes the memory with respect to all past tokens (comes with the cost of slow training), and Titans optimizes the memory with respect to all past tokens but with an implicit decay term (i.e., the result of the momentum) for each past token, maintaining parallelizability. . Our formulations of sliding window RNN and Omega rule fill this gap by optimizing the memory with respect to a context window of past tokens based on parametric methods, effectively memorizing the context instead of individual tokens.

Deep Linear Attention. As a novel baseline, we present Deep (Gated) Linear Attention (DLA) that replaces a matrix-valued memory in (gated) linear attention (Katharopoulos et al., 2020; Yang et al., 2024a) with a deep neural network (e.g., $k$-layer MLP). As discussed earlier in (Hebbian Rule), using dot product similarity as the internal attentional bias results in linear attention. Thus, leveraging recent deep memory modules (Sun et al., 2024; Behrouz et al., 2024; Behrouz et al., 2025), we optimize the memory using gradient descent with dot product attentional bias:

where $\ell(\mathcal{M}_{t-1};\phi(\mathbf{k}_{t}),\mathbf{v}_{t})=\langle\mathcal{M}_{t-1}(\phi(\mathbf{k}_{t})),\mathbf{v}_{t}\rangle$ and $\phi(\cdot)$ is a polynomial kernel. The training of DLA can simply be parallelized using the hybrid of linear and non-linear chunk-wise training, the same as Sun et al. (2024); Behrouz et al. (2024) and our discussion in Section 3.3.

Sliding Window Linear Attention. Building upon the above intuition and the connection of our formulation to SWA, we present Sliding Window Linear Attention (SWLA) block. Following the formulation of linear attention in associative memory perspective (Behrouz et al., 2025), we use dot product similarity (i.e., $\ell(\mathcal{M}_{t};\mathbf{k}_{i},\mathbf{v}_{i})=\langle\mathcal{M}_{t}(\mathbf{k}_{i}),\mathbf{v}_{i}\rangle$) as the attentional bias and optimize the loss function using gradient descent. For the sake of clarity, we use a linear memory here to derive the closed form:

In the online case ($c=1$) and $\phi(\cdot)=(\cdot)$, this recurrence is the same as linear attention (Katharopoulos et al., 2020).

We first recall the formulation of softmax attention in Transformers (i.e., Equation 1):

which its $\exp(\cdot)$ kernel is not separable and so cannot be written as a recurrence. Following the discussion in Kacham et al. (2024a), one can see $\exp(\cdot)$ kernel (compared to polynomial kernel $\phi_{p}(\cdot)$) as a feature map that maps the input into an infinite dimension. That is, we define:

where $x^{\otimes p}=x\otimes x^{\otimes(p-1)}$ is a “self-tensoring” operator with Kronecker product (Kacham et al., 2024a) and so:

Based on the above kernel, we can reformulate the attention (see Equation 21) as: (we remove $1/\sqrt{d_{\text{in}}}$ term for the sake of simplicity)

This formulation, provides another important insight on the differences of attention and (kernel) recurrent models: Softmax attention as an associative memory has an unbounded memory and so can better memorize larger context into its parameters. Building upon this insight, we present DeepTransformers by replacing polynomial kernel with $\phi^{*}(\cdot)$ kernel in Deep Linear Attention formulation (Equation 19), resulting in unnormalized formulation of:

In the special case of linear memory, we can derive the closed form for the above formulation as:

which matches the output of the unnormalized Transformers. Therefore, DeepTransformers are strict generalizations of Transformers with softmax attention (Vaswani et al., 2017).

Our above formulation of DeepTransformers is based on the (Hebbian Rule), which is also used in original Transformers. However, as discussed earlier, using more powerful memory management and learning rules in associative memory modules can further enhance their performance. To this end, we extend the above formulation by replacing the Hebbian rule with our Omega learning rule, resulting in an unnormalized formulation of Deep Omega Transformers (Dot):

We now discuss special instances of Dot to provide further intuition on its generalized formulation.

Linear Memory. This setup results in the following unnormalized formulation:

Online Case with $c=1$. We now let $c=1$:

The above (unnormalized) formulation can be seen as the generalization of Transformers with Delta Rule. Therefore, due to the unbounded memory, Dot not only appends the new keys and values (similar to original Transformers), but it also replaces the new value with its predicted value from the previous state.

In this section, we discussed how the training process of Atlas can be parallelized. For the sake of clarity, we assume $c=1$. Generalizing the process to arbitrary value for $c$ follows the procedure in Section 3.3. We use the same process as we discussed in Section 3.3 and so chunk the sequence and compute all the gradients with respect to the last state of the previous chunk. Accordingly, using the recurrence of Atlas with momentum but without , we have:

Since $t^{\prime}$ is the last state of the previous chunk, we can calculate all the gradients before hand and so we let $u_{t}=\nabla\ell(\mathcal{M}_{t^{\prime}};\mathbf{k}_{t},\mathbf{v}_{t})$. Therefore, we have:

Now by expanding the second recurrence, we have:

where $G$ is the gradient matrix, $E$ and $\Theta$ are diagonal matrices with value $\theta$ and $\eta$, and $\odot$ is broadcasting.

The main advantage of the above formulation (chunk wise recurrence) is that the recurrence of momentum is independent of the state of memory. That is, we can calculate all the momentum terms in the beginning of the chunk using the above formulation. Now in the Muon case, we want to use Newton-Schulz algorithm on the momentum terms, which results in:

Since the calculation of all $\mathcal{S}_{t}$s can be done in parallel, the calculation of $\texttt{Newton-Schulz5}(\cdot)$ can also be done in parallel.

Architectural Backbone. As for the architectural backbone, we follow the recent modern recurrent models (Behrouz et al., 2024; Arora et al., 2024; Yang et al., 2024b; Allen-Zhu, 2025) and use linear layers to project keys, values, and queries, followed by short convolution layers with size 4. We apply normalization on keys and queries to stabilize the training. We also follow Behrouz et al. (2024) and use two hybrid variants of MAL and MAG for our Atlas model. The architectures are illustrated in Figure 3. For models with deep memory architectures we use 2-layer MLP with residual connections:

We further extend this memory architecture, which is commonly used in recent studies (Behrouz et al., 2024; Irie et al., 2021; Behrouz et al., 2025), to gated MLP layer as: