Test-time regression: a unifying framework for designing sequence models with associative memory

Ke Alexander Wang $^{}$ $^{1}$, Jiaxin Shi $^{\dagger}$ $^{2}$, Emily B. Fox $^{1}$ $^{2}$ [email protected] $^{}$ Correspondence to [email protected]

$^{\dagger}$ Now at Google Deepmind

$^{1}$ Department of Computer Science, Stanford University

$^{2}$ Department of Statistics, Stanford University

Abstract

Sequence models lie at the heart of modern deep learning. However, rapid advancements have produced a diversity of seemingly unrelated architectures, such as Transformers and recurrent alternatives. In this paper, we introduce a unifying framework to understand and derive these sequence models, inspired by the empirical importance of associative recall, the capability to retrieve contextually relevant tokens. We formalize associative recall as a two-step process, memorization and retrieval, casting memorization as a regression problem. Layers that combine these two steps perform associative recall via ``test-time regression'' over its input tokens. Prominent layers, including linear attention, state-space models, fast-weight programmers, online learners, and softmax attention, arise as special cases defined by three design choices: the regression weights, the regressor function class, and the test-time optimization algorithm. Our approach clarifies how linear attention fails to capture inter-token correlations and offers a mathematical justification for the empirical effectiveness of query-key normalization in softmax attention. Further, it illuminates unexplored regions within the design space, which we use to derive novel higher-order generalizations of softmax attention. Beyond unification, our work bridges sequence modeling with classic regression methods, a field with extensive literature, paving the way for developing more powerful and theoretically principled architectures.

Executive Summary: Sequence models such as Transformers and their efficient recurrent alternatives have proliferated rapidly, yet they have largely been developed through separate, often empirical lines of work. This fragmentation obscures their shared mechanisms and limits systematic improvement. A key empirical observation is that an architecture’s ability to perform associative recall—retrieving contextually relevant information from prior tokens—strongly predicts its language-modeling performance. The paper addresses this by introducing a single, principled framework that explains why many existing models succeed and how to design better ones.

The document sets out to derive sequence layers explicitly engineered for associative recall. It formalizes recall as a two-step process: memorization of key-value associations followed by retrieval. Memorization is cast as solving a weighted regression problem over the input tokens at inference time; retrieval is then simply applying the resulting regressor to a query. Layers that perform both steps in a single forward pass are termed “test-time regression” layers. Any regression method can be used, so the framework reduces architecture design to three explicit choices: the relative weight of each association, the function class of the regressor, and the optimization algorithm used to fit it.

Under this lens, linear attention and its feature-mapped variants arise as (suboptimal) linear or kernel least-squares regression that ignores correlations among keys. Adding geometric forgetting weights recovers gated linear attention and state-space models such as Mamba. First-order gradient methods on the same objective recover fast-weight programmers (DeltaNet) and online learners. Softmax attention itself emerges as locally constant nonparametric regression; query-key normalization is mathematically required for this interpretation to hold. The same construction also yields previously unexplored higher-order generalizations based on local linear regression.

Empirically, the derived layers do perform implicit regression in one forward pass. On a non-stationary autoregressive prediction task, methods that account for key covariance (recursive least squares) or adapt locally (nonparametric regressors) substantially outperform linear attention. On the multi-query associative-recall benchmark, a single test-time regression layer plus one short convolution suffices to solve the task perfectly when memory capacity is adequate, showing that sequence length itself is not the limiting factor—memory capacity and key construction are.

These results matter because they replace ad-hoc architectural tweaks with a mathematically grounded recipe drawn from a mature regression literature. They explain the practical value of techniques such as output normalization and query-key normalization, identify why linear attention underperforms on correlated inputs, and open concrete avenues for new, more expressive layers. The framework also highlights that hardware-efficient implementations of second-order or higher-order regressors remain an open engineering challenge.

Actionable next steps include: (1) developing hardware-aware algorithms (analogous to FlashAttention) for the higher-order nonparametric variants; (2) exploring nonlinear parametric regressors at test time, incorporating momentum and decoupled weight decay; and (3) systematically testing whether the short-convolution-plus-single-layer construction scales to realistic language-modeling regimes. The main limitations are that the framework currently covers only query-key-value architectures and that the most expressive variants are not yet practical at scale without further optimization work. Within these bounds, the derivations and supporting experiments provide a coherent, evidence-based map for future sequence-model design.

1. Introduction

Section Summary: The introduction explains that sequence models have become central to machine learning across many fields, yet they have developed in isolation with incompatible notations that obscure their shared strengths. It highlights associative recall—the ability to store and later retrieve relevant information from context—as a key driver of strong performance in language modeling, and argues that this capability can be systematically achieved by treating the forward pass as a regression problem solved on the input tokens themselves. The paper therefore presents a unifying “test-time regression” framework that derives many existing architectures, including linear attention, state-space models, and softmax attention, simply by varying the choice of regression weights, parameterization, and optimization method.

Sequences play a vital role in modern machine learning by providing a powerful abstraction: any computational task can be viewed as transforming one sequence into another ([1]). This sequential perspective has spread across diverse domains, including natural language processing ([1, 2, 3]), computer vision ([4, 5]), time series analysis ([6, 7, 8]), and computational biology ([9, 10, 11]), highlighting the importance of building generically applicable sequence layers ([12]).

This development has produced a diversity of architectures, each with its own unique characteristics and performance trade-offs. While these architectures have achieved considerable success, they have largely emerged through separate lines of investigation. Such a fragmented and often empirically-driven approach to model development limits our ability to systematically understand and improve design choices. Moreover, the idiosyncratic notations of each architecture obscures their underlying connections ([13]). Given the wide variety of sequence models, a natural question to ask is whether there is an underlying principle that explains why some sequence models work better than others.

One empirical discovery that ties together disparate architectures is the strong correlation between an architecture's associative recall ability and its language modeling performance ([14, 15]). Associative recall is the act of retrieving contextually relevant information based on an association with a query. ([15]) gives the following example of how associative memory and retrieval can improve language modeling. Consider the example sentence from [16]: "Hakuna Matata, it means no worries for the rest of your days. Hakuna ". To predict the next word, we can first memorize the previous occurrence of "Hakuna" and its associated value "Matata"; the next time we encounter "Hakuna" again we can retrieve its associated value from memory as our prediction. Notice that we can perform this task using only in-context information and we can make an accurate prediction even if we have never encountered this strange phrase before. Indeed, transformer-based language models have been discovered to exhibit this kind of behavior via "induction heads" that emerge during training ([14]).

Given the empirical importance of associative recall, how can we systematically design neural network layers that can perform associative recall (AR)? In this paper, we introduce a simple but principled framework for deriving sequence layers designed to perform associative recall, which we call "test-time regression layers". From the "Hakuna Matata" example, we see that associative recall has two-steps: memorization and retrieval. Our crucial observation is that we can implement the memorization step by solving a weighted regression problem. We can then generate an output by applying the regressor to a cue/query token, retrieving the most relevant token from our associative memory. Combining both memorization and retrieval into the forward pass of a sequence layer results in a layer that performs regression over the input tokens, a procedure that we call "test-time regression", visualized in Figure 1. Our terminology reflects that the regressor is regenerated with each forward pass and depends only on the input tokens rather than a fixed training dataset.

Since any regression method can be used, our framework provides a general recipe to derive a large class of sequence layers.

Under our framework, an AR-based sequence layer is a mathematical consequence of choosing (1) the regression weights, (2) the regressor parameterization, and (3) the optimization algorithm for finding the regressor. In this paper, we consider a few choices of these three "ingredients", and show how to derive many recently proposed classes of sequence layers, illustrating the generality of our framework. Our derivations reveal that linear attention ([17]), its feature-mapped variants (e.g. ([18, 19, 20, 21, 22, 23])), its gated variants ([24, 25, 26, 27, 28, 29, 30, 31]), state-space model layers ([32, 33]), fast-weights layers ([34, 35, 35]), online learning layers ([36, 37, 35, 38]), and softmax self-attention ([12]) are all test-time regression layers, implicitly performing memorization followed by retrieval in their forward passes, despite being developed from different perspectives. Figure 1 previews how existing architectures are instantiations of test-time regression layers within our framework.

Our derivations also lead to new understandings of existing sequence layers. We show that layers based on linear attention underperform because they fail to account for the correlation between tokens. We also show that query-key normalization, an important technique in stabilizing the training of large language models ([39, 40]), is mathematically necessary to ensure that softmax self-attention is a proper local constant regressor. Finally, we propose a higher order generalization of softmax attention, motivated by local linear regression.

Outline of our paper.

We start in Section 2 by examining a few of the most prominent classes of sequence layers. Although each class of layers was developed from distinct motivations, their similar computation pattern hints at a common unifying theme that underlies their effectiveness. We then introduce our test-time regression framework in Section 3 to formalize the connection between associative memorization and regression. Our perspective provides a systematic approach to designing sequence layers that produce its outputs via associative recall. In Section 4, we show that indeed all of the aforementioned classes of sequence layers can be understood from a single unified perspective using the principles of test-time regression and associative recall. We demonstrate the generality of our framework by deriving these layers simply by varying how we minimize the regression objective. Section 5 empirically validates that test-time regression layers implicitly perform regression over its input tokens with a single forward pass. Then, in Section 6 we examine how to construct effective key-value pairs for associative recall in next-token prediction tasks. We discuss prior work related to associative recall and memory in Section 7. We finish in Section 8 by discussing the broader implications of viewing sequence models through the lens of regression, including future directions for new architectures.

2. Background on existing sequence modeling layers

Section Summary: Sequence modeling layers typically convert input tokens into queries, keys, and values, then combine them through different mixing operations to produce outputs. Softmax attention performs this mixing via similarity-weighted sums but scales quadratically with sequence length, while linear attention achieves linear-time computation by maintaining and updating a compact matrix-valued state that accumulates outer products of keys and values. Later variants add a forgetting factor to this state or adopt equivalent state-space formulations, allowing the model to selectively discard older information and better handle long sequences.

The goal of sequence modeling is to transform a sequence of input tokens $\mathbf{x}_1, \ldots, \mathbf{x}_T$ into output tokens $\mathbf{y}_1, \ldots, \mathbf{y}_T$ ([1]). A common design pattern, first proposed by [12] is to transform the input tokens into key-value pairs $(\mathbf{k}_1, \mathbf{v}_1), \ldots, (\mathbf{k}_T, \mathbf{v}_T) \in \mathbb{R}^{{D_k}}\times \mathbb{R}^{{D_v}}$ and queries $\mathbf{q}_1, \ldots, \mathbf{q}_T \in \mathbb{R}^{D_k}$ before applying operations that "mix" the sequences together to produce the output tokens. Typically, each $\mathbf{q}_t, \mathbf{k}_t, $ and $\mathbf{v}_t$ is a linear transform of the corresponding input token $\mathbf{x}_t$, which can be thought of as different "views" of $\mathbf{x}_t$. Most subsequent sequence layers have kept this pattern but propose other ways of "mixing" queries, keys, and values. We begin by reviewing a few classes of existing sequence layers that follow this pattern before unifying them under our test-time regression framework.

Softmax attention.

Given input tokens $\mathbf{x}_1, \ldots, \mathbf{x}_T$, and their query, key, value vectors, softmax attention ([12]) computes each output token by

$ \mathbf{y}t = \sum{i=1}^t \mathbf{v}_i \frac{s(\mathbf{k}_i, \mathbf{q}t)}{\sum{j=1}^t s(\mathbf{k}_j, \mathbf{q}_t)} $

for $t\in [T]$ where $s(\mathbf{k}, \mathbf{q}) = \exp(\mathbf{k} ^\top \mathbf{q}/\sqrt{{D_k}})$. Each output $\mathbf{y}_t$ is a weighted sum of values ${\mathbf{v}i}{i=1}^t$, weighted by the similarity between the query $\mathbf{q}_t$ and the keys ${\mathbf{k}i}{i=1}^t$ associated with the values. This mechanism can be seen as a continuous-valued lookup table, using $\mathbf{q}_t$ to lookup an interpolated value for $\mathbf{y}_t$. Though highly effective, self-attention requires $\mathcal{O}\left(T^2 \right)$ time to autoregressively compute the entire length $T$ sequence even after caching the keys and values, which can be prohibitively slow for generating long sequences.

Linear attention.

To address the computational challenges of self-attention, there has been a wealth of methods that seek to develop more efficient alternatives; see the survey by [41]. One particular class of more efficient sequence models has its origins in a recurrent variant of attention, known as linear attention ([17]). To achieve this efficiency, linear attention uses the standard inner product $\mathbf{k} ^\top \mathbf{q}$, rather than using the exponentiated inner product $s(\mathbf{k}_i, \mathbf{q}_t)$, producing outputs of the form $\mathbf{y}t \propto \sum{i=1}^t \mathbf{v}_i \mathbf{k}_i ^\top \mathbf{q}_t$ up to a normalization. This change allows all $T$ outputs to be recurrently computed in $\mathcal{O}\left(T \right)$ time, an asymptotic speedup, by updating a ${D_v}\times {D_k}$ "matrix-valued state" $\mathbf{M}t = \sum{i=1}^t \mathbf{v}_i \mathbf{k}_i ^\top$, following the recurrence

$ \mathbf{M}t = \mathbf{M}{t-1} + \mathbf{v}_t \mathbf{k}_t ^\top, \quad \mathbf{y}_t = \mathbf{M}_t \mathbf{q}t = \sum{i=1}^t \mathbf{v}_i \mathbf{k}_i ^\top \mathbf{q}_t. $

Each rank-1 outer product $\mathbf{v}_t \mathbf{k}_t ^\top$ can be thought of as writing a new key-value pair into memory. Although [17]'s original formulation of linear attention included a normalization like softmax attention, we drop the normalization here since later works found that this normalization method causes training instabilities ([42, 33]) .

To increase expressiveness, [17] also proposed a variant of linear attention, inspired by the kernel trick ([43]), that replaces the inner product $\mathbf{k} ^\top \mathbf{q}$ with an inner product in a feature space $\phi(\mathbf{k})^\top\phi(\mathbf{q})$ where $\phi$ is a feature map $\phi: {D_k}\to {D_\phi}$. This modification results in the recurrence

$ \mathbf{M}t = \mathbf{M}{t-1} + \mathbf{v}_t\phi(\mathbf{k}_t)^\top, \quad \mathbf{y}_t = \mathbf{M}_t \phi(\mathbf{q}t) = \sum{i=1}^t \mathbf{v}_i\phi(\mathbf{k}_i)^\top\phi(\mathbf{q}_t). $

[17] used $\phi = 1 + \operatorname{ELU}$, but numerous alternative choices of feature maps have since been proposed, including the cosine function ([19]), deterministic parameter-free projections ([34]), polynomial maps ([21, 22, 44]), random features ([18, 45, 23]), and more.

State-space models and linear attention with forgetting.

Since the recurrent state $\mathbf{M}_t$ of linear attention has finite size (${D_v}\times {D_k}$ numbers), $\mathbf{M}_t$ can only store a bounded amount of information. In, $\mathbf{M}_t$ weighs all key value pairs equally across timesteps, trying to compress all key-value pairs equally well into its finite-sized state. As a result, the performance of linear attention can degrade as the sequence length increases. To address this limitation, various works have proposed adding a forgetting factor $\gamma_t$ to allow $\mathbf{M}_t$ to retain only the most recent information, decaying out key-value pairs over time. This produces the recurrence

$ \mathbf{M}t = \gamma_t \mathbf{M}{t-1} + \mathbf{v}_t \mathbf{k}_t ^\top, \quad \mathbf{y}_t = \mathbf{M}_t \mathbf{q}_t. $

At each step $t$, all previous values in $\mathbf{M}_{t-1}$ are attenuated by a factor of $\gamma_t\in[0, 1]$ before the new value is stored. When $\gamma_t$ depends on $\mathbf{x}_t$, this forgetting factor can be seen as a forgetting gate, akin to that of an LSTM ([46]). This class of linear attention with forgetting includes RetNet ([24]), gated linear attention ([30]), Gateloop ([26]), mLSTM ([31]), LRU ([25]), HGRN ([28]), and RWKV-6 ([29]). See Table 1 of [30] for how each model implements $\gamma_t$.

The forgetting factor in linear attention is also equivalent to another class of sequence layers known as state-space models (SSMs), which includes Mamba ([32]) and Mamba-2 ([33]). State-space models originated from a line of work on optimal online orthogonal polynomial projections ([47, 48, 49]) that enabled the architecture to retain information over long sequences. However, many recent works ([50, 51, 52]), including the Mamba-2 followup ([33]), have since shown that SSM layers are equivalent to linear attention with forgetting; the forgetting factor maps onto the time-varying step size and transition matrix used in SSMs.

Fast-weight layers and test-time training layers.

An alternative to attenuating all past information in $\mathbf{M}_{t-1}$ is to forget only the values whose keys are most similar to the new key $\mathbf{k}_t$. This was the motivation behind DeltaNet ([34]), which is a revival of classic fast weight layers ([53]). DeltaNet follows a recurrence

$ \begin{aligned} \mathbf{M}{t} &= \mathbf{M}{t-1} - \underbrace{(\mathbf{M}{t-1}\mathbf{k}t)}{\mathbf{v}{t}^{\text{old}}}\mathbf{k}_t ^\top + \underbrace{[\beta_t \mathbf{v}t + (1-\beta_t) \overbrace{\mathbf{M}{t-1}\mathbf{k}_t}^{\mathbf{v}t^{\text{old}}}]}{\mathbf{v}_t^{\text{new}}}\mathbf{k}t ^\top \ &= \mathbf{M}{t-1}(\mathbf{I} - \beta_t \mathbf{k}_t \mathbf{k}_t ^\top) + \beta_t \mathbf{v}_t \mathbf{k}_t ^\top \end{aligned} $

where the first equation follows the notation of [34] in defining $\beta_t\in [0, 1]$, $\mathbf{v}_t^{\text{old}}$, and $\mathbf{v}_t^{\text{new}}$. Like previous layers, the output at step $t$ is also $\mathbf{y}_t = \mathbf{M}_t \mathbf{q}_t.$ Comparing the recurrence of DeltaNet (the section) to that of linear attention (the section), we can interpret the update rule as rewriting the value that was bound to key $\mathbf{k}_t$. The update rule deletes an old value $\mathbf{v}t^{\text{old}}= \mathbf{M}{t-1}\mathbf{k}t$ retrieved from the previous timestep's memory $\mathbf{M}{t-1}$ but using the new key $\mathbf{k}_t$. The layer then updates the recurrent state with a new interpolated value $\mathbf{v}_t^{\text{new}}$, binding it to $\mathbf{k}_t$.

The iterates $\mathbf{M}_t$ are also known as "fast-weights", originally proposed by [53]. These weights/iterates are "fast" because they change with each timestep, rapidly adapting to the input tokens, in contrast to the "slow-weights" in the rest of the neural network, such as the linear projections weights, which are only adjusted at training time. The traditional "slow-weights" can be seen as the weights of a metalearner ([54]), learning to use the "fast-weights" to learn from the input sequence ([37]).

Online learning layers.

The recurrences we have reviewed so far all share a common high-level behavior: at each timestep $t$, the recurrent state must be adapted to incorporate information from the latest key-value pair $(\mathbf{k}_t, \mathbf{v}_t)$. This property can be seen as making an online sequence of decisions, balancing between retaining historical information and adapting to new information.

Motivated by this sequential decision making perspective, [36] proposed to derive recurrences by minimizing an online sequence of objectives:

$ \mathbf{M}t = \operatorname{\arg!\min}\mathbf{M}\ d(\mathbf{M}_{t-1}, \mathbf{M}) + \ell_t(\mathbf{M}) $

where $d(\mathbf{M}{t-1}, \mathbf{M})$ is a measure of the discrepancy between $\mathbf{M}$ and the previous iterate $\mathbf{M}{t-1}$, encouraging retention of older information, and $\ell_t$ is a loss function that encourages adapting to new information. Once again, the output is computed by $\mathbf{y}_t = \mathbf{M}_t \mathbf{q}_t$. Recurrent layers derived via are called "online learning layers" or "online learners". Recent examples of such layers include Longhorn ([36]), Gated DeltaNet([35]), and Titans ([38]), each of these layers was derived from specific choice of functions $d$ and $\ell_t$. See Table 2 from [35] for each layer's choice of $d$ and $\ell_t$.

3. Test-time regression as a framework for designing sequence layers

Section Summary: The section introduces test-time regression as a unifying lens for sequence layers, framing their core job as associative recall: at each step, a model memorizes key-value pairs by solving a regression problem, then retrieves an output by applying the fitted function to a new query. This two-step process directly explains how layers achieve in-context learning and generalizes across linear, nonlinear, and multi-head designs. The authors distill the approach into a simple design recipe based on choosing association weights, the function class, and the optimization method, allowing new layers to be derived systematically rather than through isolated tweaks.

Having discussed a number of existing sequence layers, each motivated by an atomic architectural modification, we now present a unified framework to clarify their relationship to each other. Our approach is motivated by empirical findings that a model's associative recall ability strongly correlates with its performance and in-context learning capabilities ([14, 15]). We formalize the two step process of memorization and retrieval, presenting a principled recipe for deriving sequence layers with mathematically grounded recall abilities.

Step 1: Memorization as regression.

From our everyday experience, we already have an intuition for how associative memory should behave; for example, hearing a friend's name should trigger a mental impression of that friend. This pairing of a "cue" and a "response" is known as an association. To reflect terminology in deep learning ([12]), we will refer to cues as "keys" and responses as "values". With this intuition, we can define a mathematical model for associative memory. Given key-value pairs $(\mathbf{k}_1, \mathbf{v}_1), \ldots, (\mathbf{k}_t, \mathbf{v}_t) \in \mathbb{R}^{D_k} \times \mathbb{R}^{D_v}$ to be memorized, an associative memory system is a function ${m}: \mathbb{R}^{D_k} \to \mathbb{R}^{D_v}$ such that ${m}(\mathbf{k}_i)\approx \mathbf{v}_i$ for $i=1, \ldots, t$ ([55, 56]). Such a mapping ${m}$ between two sets of related objects is also known as a hetero-associative memory in classic signal processing and neurocomputing literature ([57]).

Our main observation is that finding a memory map that satisfies these constraints reduces to finding a vector-valued regressor ${m}$. We can then store a set of key-value associations into an associative memory system by solving the weighted regression problem

$ \text{Memorize associations:}\quad \quad {m}t \approx \operatorname{\arg!\min}{{m}\in \mathcal{M}}\ \frac{1}{2}\sum_{i=1}^t\gamma_i^{(t)}\lVert \mathbf{v}_i - {m}(\mathbf{k}_i)\rVert _2^2, $

where the importance of each association is controlled by adjusting the relative weights ${\gamma_i^{(t)}}_{i=1}^t$ Classic works on associative memory are restricted to linear memory maps, corresponding to linear regression ([56], Chapter 6.6), but here we allow $\mathcal{M}$ to be any function class. We refer to this relationship between regression and associative memory as the "regression-memory correspondence".

Step 2: Memory retrieval as function application.

Having derived a regression-based procedure to memorize a set of associations into ${m}$, we can now retrieve values from memory by applying ${m}$ to a new vector, which we call a "query" $\mathbf{q}$, producing output

$ \text{Memory retrieval:} \quad \quad \mathbf{y}_t = {m}_t(\mathbf{q}_t). $

The retrieval performance depends on how well is minimized. When attains zero-loss, ${m}_t$ memorizes all associations perfectly, and we get perfect retrieval: ${m}_t(\mathbf{k}_i) = \mathbf{v}_i$ for $i= 1, \ldots, t$. For certain function classes $\mathcal{M}$, e.g. smooth functions, ${m}$ can also extrapolate outside of the existing set of keys. In those cases, ${m}_t(\mathbf{q})$ instead retrieves a value based on the proximity of $\mathbf{q}$ to the existing set of keys.

Sequence layers that memorize then retrieve.

Having formalized memorization and retrieval via our regression framework, we can now combine the two to produce sequence layers that directly performs associative recall as a forward pass. Recall that we are given ${\mathbf{q}i}{i=1}^T$, ${\mathbf{k}i}{i=1}^T$, and ${\mathbf{v}i}{i=1}^T$, and we seek to produce a set of outputs ${\mathbf{y}i}{i=1}^T$ with our sequence layer. The forward pass of a test-time regression layer first memorizes a set of key-value pairs and then produces $\mathbf{y}_t$ by retrieving a value from its associative memory. More specifically, in causal modeling, at each timestep $t$, we first solve to get a regressor ${m}_t$ that memorizes the prefix ${(\mathbf{k}_i, \mathbf{v}i)}{i=1}^t$; then we apply ${m}_t$ to $\mathbf{q}_t$ to get $\mathbf{y}_t = {m}_t(\mathbf{q}_t)$. This naturally generalizes to non-causal tasks where ${m}_t$ can depend on suffix key-value pairs ${\mathbf{k}_i, \mathbf{v}i}{i > t}$ as well. Our procedure is compatible with multihead design patterns inspired by multihead attention ([12]), multi-query attention ([58]), and grouped query attention ([59]). For example, multihead test-time regression can be implemented using $H$ independent sets of queries, keys, and values, resulting in $H$ independent associative memory maps at each timestep.

A recipe for designing your own sequence layer.

Although succinctly summarizes the memorization step of our framework, it leaves many details unspecified. For example, over which class of regressors $\mathcal{M}$ should we consider? Additionally, which associations should we prioritize memorizing? Finally, which optimization procedure should we use to minimize the loss? The regression-memory correspondence thus allows for a vast design space. We partition this wide landscape into three simple design choices:

1. the relative importance of each association, specified through weights in,
2. the function class $\mathcal{M}$ over which we search,
3. the minimization algorithm,

giving us a formulaic "recipe" for deriving sequence layers that can perform associative recall.

4. Deriving existing architectures from regression

Section Summary: The section shows how a single regression-based framework can derive and unify many recent sequence models by treating each new input as a key-value pair to be memorized and then solving a test-time regression problem to produce the output. It organizes existing layers into categories such as plain linear attention, gated linear attention (including state-space models), fast-weight programmers, online learners, and standard softmax attention. The text then works through the first case in detail, showing that linear attention is simply a hardware-friendly approximation to linear least-squares regression that drops the covariance term and therefore requires extra normalization to remain stable.

To demonstrate how our framework unifies our understanding of sequence models, we succinctly derive a swath of recently proposed sequence layers. We categorize existing sequence layers into four distinct classes: linear attention (possibly with feature maps), linear attention with gating (which includes state-space models like Mamba), fast-weight layers, online learners, and softmax self-attention.

Notation.

We first explain the mathematical notation we will use in our derivations. We use bold lowercase letters to refer to vectors (e.g. $\mathbf{x}, \mathbf{k}, \mathbf{v}$) and bold uppercase letters to refer to matrices (e.g. $\mathbf{M}, \mathbf{K}, \mathbf{V}$). Assume we have a sequence of $T$ inputs $\mathbf{x}_1, \ldots, \mathbf{x}_T$ which have already transformed into a sequence of keys and values $(\mathbf{k}_1, \mathbf{v}_1), \ldots, (\mathbf{k}_T, \mathbf{v}_T)$. We will use these keys and values as the dataset for our test-time regression problem. Let $\mathbf{K}_t = \begin{bmatrix} \mathbf{k}_1 \ldots \mathbf{k}_t \end{bmatrix}^\top \in \mathbb{R}^{t\times {D_k}} $ and $ \mathbf{V}_t = \begin{bmatrix} \mathbf{v}_1 \ldots \mathbf{v}_t \end{bmatrix}^\top \in \mathbb{R}^{t\times {D_v}} $ be matrices that contain $ t $ keys and $ t $ values respectively, for $ t=1, \ldots, T$. Each key/value vector is a row of the key/value matrix.

Vignette 1: Linear attention is (suboptimal) linear least squares regression

We start by deriving an architecture from linear least squares, the most classic regression method. Linear least squares corresponds to the following choices in our recipe.

1. Choice of weights: assign equal weight $\gamma_i^{(t)} = 1$ to each association,

$ {m}{\textsc{linear}, , t} = \operatorname{\arg!\min}{{m}\in {\mathcal{M}{\textsc{linear}}}} \frac{1}{2}\sum{i=1}^t \lVert \mathbf{v}_i - {m}(\mathbf{k}_i)\rVert _2^2. $

2. Choice of function class: linear functions ${\mathcal{M}_{\textsc{linear}}} = {{m}\mid {m}(\mathbf{k}) = \mathbf{M} \mathbf{k}, \ \mathbf{M} \in \mathbb{R}^{{D_v}\times {D_k}}}$.
3. Choice of minimization algorithm: analytical solution equivalent to one step of Newton's method[^1].

[^1]: Since the objective function is quadratic, the analytical solution is equivalent to one iteration of Newton's method initializing from the origin.

Since ${m}_{\textsc{linear}, , t}$ is a linear function, it is equivalent to a linear map $\mathbf{M}_t\in \mathbb{R}^{{D_v}\times {D_k}}$. Minimizing the linear least squares objective results in the solution:

$ \begin{aligned} \mathbf{M}t = \operatorname{\arg!\min}{\mathbf{M}} \frac{1}{2}\sum_{i=1}^t \lVert \mathbf{v}_i - \mathbf{M} \mathbf{k}_i\rVert _2^2 =\begin{cases} \mathbf{V}_t ^\top \mathbf{K}_t(\mathbf{K}_t ^\top \mathbf{K}_t)^{-1}, \ \text{overdetermined, e.g. } t \geq {D_k}, \ \mathbf{V}_t ^\top (\mathbf{K}_t \mathbf{K}_t ^\top)^{-1}\mathbf{K}_t, \ \text{underdetermined, e.g. } t < {D_k}, \end{cases} \end{aligned} $

where we choose the min-norm solution in the underdetermined case. The overdetermined case corresponds to $\mathbf{K}_t$ having rank ${D_k}$ (linearly independent columns) and the underdetermined case corresponds to $\mathbf{K}_t$ having rank $t$ (linearly independent rows). In general, we can express the solution in terms of a Moore-Penrose pseudoinverse: $\mathbf{M}_t = \mathbf{V}_t ^\top(\mathbf{K}_t ^{\dagger})^\top$. We restrict our discussion to the case of $t \geq D_k$ with linearly independent columns, which is the case for typical applications.

Relationship to linear attention.

Although ${m}_{\textsc{linear}, , t}$ is the optimal linear memory with respect to our objective function, it can be expensive to compute, requiring us to invert a ${D_k}\times {D_k}$ matrix at every timestep. One crude but hardware-efficient simplification would be to assume that $\mathbf{K}_t ^\top \mathbf{K}_t = \mathbf{I}$, avoiding matrix inversion completely, at the cost of some recall ability. After making this rough approximation, we arrive exactly at the equations for linear attention:

$ \mathbf{y}t = {m}{\textsc{linear}}(\mathbf{q}_t) = \mathbf{V}_t ^\top \mathbf{K}_t(\mathbf{K}_t ^\top \mathbf{K}_t)^{-1} \mathbf{q}_t \approx \mathbf{V}_t ^\top \mathbf{K}_t \mathbf{q}t =\sum{i=1}^t \mathbf{v}_i \mathbf{k}_i ^\top \mathbf{q}_t. $

Our derivation clarifies exactly how linear attention underperforms as an associative memory. When the keys are zero mean, $\mathbf{K}_t ^\top \mathbf{K}_t$ is proportional to the empirical covariance matrix of the keys. Thus linear attention is a crude associative memory that ignores the covariance between the dimensions of the key vectors. In fact, it is an optimal linear memory only when $t \leq {D_k}$ and the keys are orthonormal $\mathbf{k}_i ^\top \mathbf{k}j = \delta{ij}$.

From a training stability perspective, the inverse covariance factor helps prevent the output $\mathbf{y}_t$ from becoming unbounded. We can bound the norm of the output of a least squares layer (the section) by

$ \lVert \mathbf{y}_t\rVert _2 = \lVert \mathbf{V}_t ^\top \mathbf{K}_t(\mathbf{K}_t ^\top \mathbf{K}_t)^{-1} \mathbf{q}_t\rVert _2 \leq \frac{\lVert \mathbf{q}_t\rVert 2\left(\sum{i=1}^t\lVert \mathbf{v}_i\rVert _2\lVert \mathbf{k}_i\rVert 2\right)}{\lambda{\text{min}}(\mathbf{K}_t ^\top \mathbf{K}_t)} $

where $\lambda_{\text{min}}(\mathbf{K}_t ^\top \mathbf{K}_t)$ is the minimum eigenvalue of $\mathbf{K}_t ^\top \mathbf{K}_t$. See Appendix A for the derivation. When we ignore the covariance matrix, approximating it with the identity, the denominator disappears and the output $\mathbf{y}_t$ can grow arbitrarily large with the sequence length, leading to training instability. This instability indeed shows up in practice for linear attention, and [42] proposed output normalization to stabilize training. Our derivation shows that output normalization works by approximately restoring the self-normalizing property of that linear attention loses.

Regardless of training stability, we will see in Section 6 that ignoring the covariance causes linear attention to perform worse at associative recall.

Efficient computation.

Conveniently, we can compute $\mathbf{M}_t$ in for $t=1, \ldots T$ in parallel or sequentially. To compute them in parallel, we solve a batch of $T$ independent linear systems. To compute them sequentially, we note that $(\mathbf{K}_t ^\top \mathbf{K}t)^{-1} = (\mathbf{K}{t-1}^\top \mathbf{K}_t + \mathbf{k}_t \mathbf{k}_t ^\top)^{-1}$, allowing us to autoregressively compute the inverses via the Woodbury matrix identity. The sequential form is known as online linear regression or recursive least squares (RLS) ([60]). Thus during training we can parallelize across the sequence axis by solving a batch of $T$ linear systems, and during testing we can unroll the recurrence into a recurrent neural network. In prior deep learning literature, this idea of linear-regression-as-a-layer has been explored by [61] and [62] under the names "intention" and "mesa-layer", respectively, motivated by heuristic similarities to softmax attention, rather than deriving them from first principles via enforcing associative recall as we do here.

Nonlinear associative memory.

So far we have only considered linear associative memory implemented by linear functions in ${\mathcal{M}{\textsc{linear}}}$, which have limited expressiveness. We can derive nonlinear memory by using a feature map $\phi: \mathbb{R}^{D_k} \to \mathbb{R}^{D\phi}$, replacing each key $\mathbf{k}i$ with $\phi(\mathbf{k}i)$. Then our function class is defined by $\mathcal{M}{\phi} = {{m} \mid {m}(\mathbf{k}) = \mathbf{M}\phi(\mathbf{k}), \ \mathbf{M} \in \mathbb{R}^{{D_v}\times {D\phi}}}$ and we solve the least-squares objective $\operatorname{\arg!\min}{\mathbf{M}} \sum{i=1}^t\lVert \mathbf{v}_i - \mathbf{M} \phi(\mathbf{k}_i)\rVert _2^2/2$. Solving this least squares objective for $\mathbf{M}$ results in an optimal associative memory that is nonlinear with respect to the input keys. Parameterizing associative memory with a nonlinear feature map has a long history, dating back to [63] which found polynomial feature maps to be highly effective, consistent with more recent works ([21, 44]).

If we make the same crude approximation as linear attention of disregarding the covariance among ${\phi(\mathbf{k}i)}{i=1}^t$, we arrive at a suboptimal nonlinear memory. Thus, recent works that apply linear attention with a nonlinear feature map, such as [34, 20, 21, 23], and others, can be understood within our framework as suboptimal nonlinear memory as well.

Vignette 2: Gated linear attention and state-space models are (suboptimal) weighted linear least squares regression

For many kinds of sequential data, such as time-series data and text data, recent tokens are more informative of predicting the next token. Thus a natural requirement is to encourage our memory ${m}_t$ to focus more on the recent history by downweighting the past. Fortunately, this can be straightforwardly enforced through classic weighted linear least squares regression together with geometrically decaying weights to forget older key-value pairs.

To develop an architecture from weighted linear least squares, we only need to add weights to the optimization objective relative to ordinary least squares:

1. Choice of weights: assign time-decaying weights to associations,

$ {m}{\textsc{linear}, , t} = \operatorname{\arg!\min}{{m}\in {\mathcal{M}{\textsc{linear}}}} \frac{1}{2}\sum{i=1}^t \gamma_{i}^{(t)}\lVert \mathbf{v}_i - {m}(\mathbf{k}_i)\rVert _2^2, $

Here, to connect to past literature, we focus on weights are geometrically decaying in time: $\gamma_{i}^{(t)} = \prod_{j=i+1}^t \gamma_j$ and $\gamma_i \in [0, 1]$ for $i=1, \ldots, t$. Each association is down-weighted based on how far it is from the current timestep $t$: $\gamma_1^{(t)} \leq \ldots \gamma_{t-1}^{(t)} \leq \gamma_t^{(t)} = 1$. 2. Choice of function class: linear functions ${\mathcal{M}_{\textsc{linear}}} = {{m}\mid {m}(\mathbf{k}) = \mathbf{M} \mathbf{k}, \ \mathbf{M} \in \mathbb{R}^{{D_v}\times {D_k}}}$. 3. Choice of minimization algorithm: analytical solution equivalent to one step of Newton's method.

We can solve weighted linear regression by instead solving an unweighted regression problem involving rescaled keys and values, $\mathbf{V}_t\mapsto \boldsymbol{\Gamma}_t^{1/2}\mathbf{V}_t$ and $\mathbf{K}_t\mapsto \boldsymbol{\Gamma}_t^{1/2}\mathbf{K}_t$, where we defined a diagonal matrix $\boldsymbol{\Gamma}t\in \mathbb{R}^{t\times t}$ with entries $(\boldsymbol{\Gamma}t){ii} = \gamma{i}^{(t)}$. Making the rescaled substitution and reusing the analytical min-norm solution that we had before, we see that

$ \begin{aligned} \mathbf{M}_t &= \begin{cases} \mathbf{V}_t ^\top \boldsymbol{\Gamma}_t \mathbf{K}_t(\mathbf{K}_t ^\top \boldsymbol{\Gamma}_t \mathbf{K}_t)^{-1}, \ \text{overdetermined, e.g. }, t \geq {D_k}\ \mathbf{V}_t ^\top (\mathbf{K}_t \mathbf{K}_t ^\top)^{-1}\mathbf{K}_t, \ \text{underdetermined, e.g. }\ t < {D_k}. \end{cases} \end{aligned} $

The underdetermined case has the same solution in both the weighted and unweighted case; intuitively this is because when the system is underdetermined, there are infinitely many solutions that interpolate the data points and attain zero regression loss, rendering the weights irrelevant.

Relationship to linear attention with gating and state-space models.

Once again, computing $\mathbf{M}_t$ requires matrix inversion which is not efficient on matrix-multiplication based hardware. Making a crude approximation similar to before of $\mathbf{K}_t ^\top \boldsymbol{\Gamma}_t \mathbf{K}_t\approx \mathbf{I}$ results in $\mathbf{M}_t \approx \mathbf{V}_t ^\top \boldsymbol{\Gamma}_t \mathbf{K}t = \sum{i=1}^t \gamma_i^{(t)} \mathbf{v}_i \mathbf{k}_i ^\top$. We can further unroll this equation into a recurrence due to the multiplicative structure of our weights:

$ \mathbf{M}t \approx \sum{i=1}^t \gamma_i^{(t)} \mathbf{v}i \mathbf{k}i ^\top = \sum{i=1}^t \left(\prod{j=i+1}^t \gamma_j\right) \mathbf{v}_i \mathbf{k}i ^\top = \gamma_t \sum{i=1}^{t-1} \gamma_i^{(t-1)} \mathbf{v}_i \mathbf{k}_i ^\top + \mathbf{v}_t \mathbf{k}t ^\top = \gamma_t \mathbf{M}{t-1} + \mathbf{v}_t \mathbf{k}_t ^\top. $

This is exactly the recurrence implemented by gated variants of linear attention, such as GLA ([30]), RetNet ([24]), Gateloop ([26]), RWKV-6 ([29]), mLSTM ([31]), LRU ([25]), as well as state-space models like Mamba-2 ([33]). Forgetting older tokens significantly improves performance, which was first discovered empirically ([46, 64, 65]). Our derivation of gating from weighted least squares provides a mathematical grounding for the forgetting heuristic upon which these sequence layers were based.

Efficient computation.

Like before, we can compute the memory of our architecture (without approximation) sequentially or in parallel as long as we are willing to invest FLOPs in matrix inversion. Parallel computation can be done by solving a batch of $T$ ${D_k}\times D_k$ linear systems. The sequential version can be derived by applying the Woodbury matrix inversion formula while taking into account our geometrically-decaying weights. Geometrically-weighted recursive least squares (RLS) is a classic method from the adaptive filtering literature ([66]).

Vignette 3: Fast weight programmers and online learners are first-order methods for solving streaming least-squares

Up until now we have derived linear associative memory architectures through direct analytical solution, equivalent to one step of Newton's method, a second-order optimization method. However, second-order iterative methods are computationally expensive, requiring matrix inversions at each step. On the other hand, disregarding the matrix inversion as done by linear attention leads to poor performance, since it ignores the covariance structure of the data.

A computationally cheaper alternative is to apply a first-order gradient-descent method. We consider this here with the following design choices:

1. Choice of weights: assign equal weight $\gamma_i^{(t)} = 1$ to each association,

$ {m}{\textsc{linear}, , t} = \operatorname{\arg!\min}{{m}\in {\mathcal{M}{\textsc{linear}}}} \frac{1}{2}\sum{i=1}^t \lVert \mathbf{v}_i - {m}(\mathbf{k}_i)\rVert _2^2. $

2. Choice of function class: linear functions ${\mathcal{M}_{\textsc{linear}}} = {{m}\mid {m}(\mathbf{k}) = \mathbf{M} \mathbf{k}, \ \mathbf{M} \in \mathbb{R}^{{D_v}\times {D_k}}}$.
3. Choice of minimization algorithm: gradient descent, possibly with adaptive step sizes, momentum, or weight decay (L2 regularization).

We can now derive the test-time regression layer corresponding to a first-order iterative method.

Since we consider linear functions, solving is equivalent to minimizing the objective

$ \mathcal{L}t(\mathbf{M}) = \sum{i=1}^t L_i(\mathbf{M}) = \sum_{i=1}^t \frac{1}{2}\lVert \mathbf{v}_i - \mathbf{M} \mathbf{k}_i\rVert _2^2, $

using its gradient

$ \begin{aligned} \nabla \mathcal{L}t(\mathbf{M}) &= \sum{i=1}^t \nabla L_i(\mathbf{M}) = \sum_{i=1}^t (\mathbf{M} \mathbf{k}_i - \mathbf{v}_i)\mathbf{k}_i ^\top = (\mathbf{M} \mathbf{K}_t ^\top - \mathbf{V}_t ^\top)\mathbf{K}_t. \end{aligned} $

to find $\mathbf{M}_t$. Let the initialization be at the origin $\mathbf{M}^{(0)}_t= \mathbf{0}$. Iteratively solving this objective with full-batch gradient descent and step size $\beta_t^{(i)}$ at each iteration results in

$ \begin{aligned} \mathbf{M}^{(i)}_t &= \mathbf{M}^{(i-1)}_t - \beta_t^{(i)}\nabla \mathcal{L}_t(\mathbf{M}^{(i-1)}_t) \ &= \mathbf{M}^{(i-1)}t - \beta_t^{(i)}\sum{j=1}^t \nabla L_j(\mathbf{M}^{(i-1)}_t) \ &= \mathbf{M}^{(i-1)}t - \beta_t^{(i)}\sum{j=1}^t (\mathbf{M}^{(i-1)}_t \mathbf{k}_j - \mathbf{v}_j)\mathbf{k}_j ^\top \ &= \mathbf{M}^{(i-1)}t(\mathbf{I} - \sum{j=1}^t\beta_t^{(i)} \mathbf{k}_j \mathbf{k}j ^\top) + \beta_t^{(i)} \sum{j=1}^t \mathbf{v}_j \mathbf{k}_j ^\top. \end{aligned} $

After $I$ iterations, we produce the linear memory represented by $\mathbf{M}_t = \mathbf{M}^{(I)}_t$.

Linear attention and variants as one-step of batch gradient descent

Consider the $I=1$ case with step size $\beta_t^{(i)} = 1$. Since we initialize at the origin, we obtain the equation for linear attention

$ \mathbf{M}^{(1)}_t = \mathbf{M}^{(0)}_t - \nabla \mathcal{L}_t(\mathbf{M}^{(0)}_t) = \mathbf{V}_t ^\top \mathbf{K}t = \sum{i=1}^t \mathbf{v}_i \mathbf{k}_i ^\top. $

This is exactly the result of performing one-step of full batch gradient descent on the least square objective $\mathcal{L}_t$. A similar result can be shown for layers that implement linear attention with forgetting, performing one step of full batch gradient descent on the weighted least square objective.

A standard technique to accelerate the convergence of gradient descent is to precondition the gradient ([67]). The optimal preconditioner in this case, due to the quadratic objective, is the constant Hessian of $\mathcal{L}t$, $\mathbf{H}{\mathcal{L}_t}(\mathbf{M})= \mathbf{K}_t ^\top \mathbf{K}_t$. This preconditioner allows gradient descent to converge in a single update

$ \mathbf{M}^{(1)}_t = \mathbf{M}^{(0)}_t - \nabla \mathcal{L}_t(\mathbf{M}^{(0)}t)\mathbf{H}{\mathcal{L}_t}(\mathbf{M}^{(0)}_t)^{-1} = \mathbf{V}_t ^\top \mathbf{K}_t(\mathbf{K}_t ^\top \mathbf{K}_t)^{-1}, $

reproducing the analytic solution from. The preconditioned recurrence is equivalent to one step of Newton's method, a second order iterative method.

By taking this iterative optimization perspective, we see that linear attention and the ideal linear associative memory lie on two extremes: one-step of batch gradient descent with no preconditioning (the section) versus one-step of batch gradient descent with perfect preconditioning (the section). The difference in performance comes from whether the optimizer accounts for the curvature of the objective, governed by the covariance of the keys ([67]). Exploring other preconditioners, e.g. quasi-Newton methods, that interpolates between the two extremes may be an interesting direction for future research.

Stochastic gradient descent as a layer.

An alternative to full-batch gradient descent is stochastic gradient descent (SGD) where we use a subset of the samples to compute the gradient at each step. This is particularly convenient in causal sequence modeling where keys and values typically are presented one at a time, rather than all at once. In this streaming setting, we can instead perform single-example SGD to iteratively minimize. In the case of linear memories ${\mathcal{M}_{\textsc{linear}}}$, the section with a single-example and variable step sizes $\beta_t^{(i)}$ becomes

$ \begin{aligned} \mathbf{M}t^{(i)} &= \mathbf{M}t^{(i-1)} - \beta_t^{(i)} \nabla L_t(\mathbf{M}{t}^{(i-1)}) \ &= \mathbf{M}{t}^{(i-1)} + \beta_t^{(i)} (\mathbf{v}t - \mathbf{M}{t}^{(i-1)}\mathbf{k}_t)\mathbf{k}t ^\top \ &= \mathbf{M}{t}^{(i-1)}(\mathbf{I} - \beta_t^{(i)}\mathbf{k}_t \mathbf{k}_t ^\top) + \beta_t^{(i)} \mathbf{v}_t \mathbf{k}_t ^\top, \end{aligned} $

where $\mathbf{M}_t$ is the matrix representation of ${m}_t$. When we take only one gradient update per timestep $t$ and initialize our iterate using the previous timestep's linear map, $\mathbf{M}t^{(0)} = \mathbf{M}{t-1}$, we end up with the recurrence of DeltaNet ([34, 68]) from:

$ \mathbf{M}_t = \mathbf{M}t^{(1)} = \mathbf{M}{t-1}(\mathbf{I} - \beta_t \mathbf{k}_t \mathbf{k}_t ^\top) + \beta_t \mathbf{v}_t \mathbf{k}_t ^\top. $

As we perform stochastic gradient descent, we produce an online sequence ${\mathbf{M}i}{i=1}^t$, each of which is an associative memory ${m}_i$ for key-value pairs up to $(\mathbf{k}_i, \mathbf{v}_i)$. The stochastic gradient approach to streaming inputs has the advantage that ${m}_t$ focuses more on recent key-value pairs, rather than uniformly over all timesteps as in the case of unweighted RLS and unweighted linear attention.

the section is known by many names, such as the delta rule or the Widrow-Hoff algorithm ([69]). Within the adaptive filtering literature this recurrence is known as Least-Mean Squares (LMS) ([60]), serving as one of the most common algorithms for streaming signals processing. From this derivation, we see that the prevalence of the update across disciplines comes from the fact that the update rule is actually an SGD update.

In our case of sequence layers, the SGD update occurs within the layer as part of the forward pass. One extension to is to consider SGD with nonlinear parametric regressors ${m}_t$. Indeed, this extension is equivalent to test-time training (TTT) layers proposed by [37], which apply SGD to a nonlinear memory ${m}_t$, e.g. multilayer perceptrons or linear maps with layer normalizations. In contrast to our associative memory derivation here, [37] motivated their method by viewing the squared loss as an self-supervised auto-encoding objective. Another extension is to take multiple gradient steps per example $(\mathbf{k}_t, \mathbf{v}_t)$, which was recently explored by DeltaProduct ([70]).

Including adaptive step sizes.

To further improve the performance of a gradient-descent-based sequence layer, we can consider adaptive step sizes. A common way to derive adaptive step sizes is to use the relationship between gradient descent and online learning ([71, 72]), a connection that has produced many successful algorithms in the past, such as Adagrad ([73]) and Adam ([74]).

Consider the Longhorn sequence layer derived by [36] by solving the online learning objective

$ \mathbf{M}{t} = \operatorname{\arg!\min}\mathbf{M}\ \frac{1}{2}\lVert \mathbf{M} - \mathbf{M}_{t-1}\rVert _F^2 + \frac{\delta_t}{2} \lVert \mathbf{v}_t - \mathbf{M} \mathbf{k}_t\rVert _2^2. $

This corresponds to $d(\mathbf{M}{t-1}, \mathbf{M}) = \lVert \mathbf{M} - \mathbf{M}{t-1}\rVert _F^2/2$ and $\ell_t(\mathbf{M}) = \delta \lVert \mathbf{v}_t - \mathbf{M} \mathbf{k}_t\rVert _2^2/$ in. The solution to is the recurrence

$ \begin{aligned} \mathbf{M}{t} &= \mathbf{M}{t-1}\left(\mathbf{I} - \frac{\delta_t}{1 + \delta_t\lVert \mathbf{k}_t\rVert _2^2} {\mathbf{k}_t \mathbf{k}_t ^\top}\right) + \frac{\delta_t}{1 + \delta_t \lVert \mathbf{k}_t\rVert _2^2}\mathbf{v}_t \mathbf{k}_t ^\top. \end{aligned} $

However, by comparing the Longhorn update in to the SGD update in, we see that Longhorn's recurrence implements SGD but with a conservative adaptive step size of $\beta_t = \delta_t / (1 + \delta_t\lVert \mathbf{k}_t\rVert _2^2)$. Thus although the Longhorn layer was derived from an online learning perspective, its global objective is still the least squares regression objective.

We can better understand the effects of this adaptive step size by looking at $\delta_t$, which controls how much we care about memorizing the latest association $(\mathbf{k}_t, \mathbf{v}_t)$. When $\delta_t = 0$, we see that $\mathbf{M}t = \mathbf{M}{t-1}$ and the latest association is excluded completely. When $\delta_t \to \infty$, the section becomes equivalent to a constrained optimization problem

$ \mathbf{M}{t} = \operatorname{\arg!\min}\mathbf{M}\ \frac{1}{2}\lVert \mathbf{M} - \mathbf{M}_{t-1}\rVert _F^2\quad \text{subject to}\ \mathbf{M}_t \mathbf{k}_t = \mathbf{v}_t. $

Solving this constrained optimization problem results in a hyperparameter-free step size $\beta_t=1/\lVert \mathbf{k}_t\rVert _2^2$ with update

$ \begin{aligned} \mathbf{M}{t} &= \mathbf{M}{t-1}\left(\mathbf{I} - \frac{1}{\lVert \mathbf{k}_t\rVert _2^2} {\mathbf{k}_t \mathbf{k}_t ^\top}\right) + \frac{1}{ \lVert \mathbf{k}_t\rVert _2^2}\mathbf{v}_t \mathbf{k}_t ^\top. \end{aligned} $

This special case is also known as the normalized LMS (NLMS) update in adaptive filtering ([75]) and known as the projection method ([76]) or Kaczmarz's method ([77]) in optimization theory. The advantage of is that it is hyperparameter-free, which may make it easier to tune. However, to our knowledge, the recurrence in has not been explored as a sequence layer.

Including L2 regularization.

In addition to adaptive step sizes, we can add regularization. Consider adding L2 regularizing to the least squares objective

$ \mathcal{L}{\text{reg}, t}(\mathbf{M}) = \sum{i=1}^t L_i(\mathbf{M}) + \frac{\lambda_i}{2}\lVert \mathbf{M}\rVert _2^2 $

where $\lambda_i\in [0, 1]$ is the regularization strength. When we minimize this regularized objective with single-example stochastic gradient descent like in, we obtain the recurrence

$ \begin{aligned} \mathbf{M}t &= \mathbf{M}{t-1} -\beta_t \left[\nabla L_t(\mathbf{M}{t-1}) + \lambda_t \mathbf{M}{t-1}\right] \ &= (1 - \beta_t\lambda_t)\mathbf{M}{t-1} -\beta_t \nabla L_t(\mathbf{M}{t-1}) \ &= (1 - \beta_t\lambda_t) \mathbf{M}_{t-1} + \beta_t(\mathbf{v}t - \mathbf{M}{t-1}\mathbf{k}_t)\mathbf{k}_t ^\top, \end{aligned} $

which is known as Leaky LMS in adaptive filtering ([60]). Notably, the first term of our new recurrence resembles but also includes a multiplicative factor similar to a forget gate.

Recently, [35] proposed Gated DeltaNet motivated by combining DeltaNet ([34]) with the forget gating of Mamba-2 ([33]), following the online learning perspective of [36]. Here we show that the Gated DeltaNet recurrence is actually an equivalent reparameterization of our recurrence in. Starting with, we have

$ \begin{aligned} \mathbf{M}t &= (1 - \beta_t\lambda_t) \mathbf{M}{t-1} + \beta_t(\mathbf{v}t - \mathbf{M}{t-1}\mathbf{k}t)\mathbf{k}t ^\top \ &= \mathbf{M}{t-1}\left[\underbrace{(1 - \beta_t \lambda_t)}{\alpha_t}\mathbf{I} - \beta_t \mathbf{k}_t \mathbf{k}_t ^\top \right] + \beta_t \mathbf{v}t \mathbf{k}t ^\top \ &= \mathbf{M}{t-1}\alpha_t\left[\mathbf{I} - \underbrace{\frac{\beta_t}{\alpha_t}}{\eta_t}\mathbf{k}_t \mathbf{k}_t ^\top \right] + \frac{\beta_t}{\alpha_t} \underbrace{(\alpha_t \mathbf{v}t)}{\mathbf{v}'_t}\mathbf{k}t ^\top \ &= \mathbf{M}{t-1}\alpha_t\left(\mathbf{I} - \eta_t \mathbf{k}_t \mathbf{k}_t ^\top \right) + \eta_t \mathbf{v}'_t \mathbf{k}_t ^\top \quad \text{(Gated DeltaNet recurrence)} \end{aligned} $

Thus Gated DeltaNet is equivalent to L2-regularized regression via SGD where the regularization strength $\lambda_t$ is coupled with the step size $\beta_t$ and $\mathbf{v}_t$ is rescaled by $\alpha_t$. We can transform into via the invertible mapping $(\alpha_t, \eta_t, \mathbf{v}_t') = (1 - \beta_t\lambda_t, \beta_t/(1 - \beta_t\lambda_t), (1 - \beta_t\lambda_t)\mathbf{v}_t)$. The inverse map is given by $(\lambda_t, \beta_t, \mathbf{v}_t) = ((1-\alpha_t)/(\alpha_t\eta_t), \alpha_t\eta_t, \allowbreak \mathbf{v}'_t/\alpha_t)$. Past works have shown that decoupling weight decay from the step size results in better generalization and easier-to-tune hyperparameters ([78]), suggesting that similar modifications to Gated DeltaNet may improve model performance.

4.1 Vignette 4: Softmax attention is nonparametric regression

Having exhaustively derived multiple associative memory architectures via linear and nonlinear parametric regression, we can now derive self-attention-like architectures using the tools of nonparametric regression.

Unnormalized softmax attention is (suboptimal) kernel regression.

We start first from kernel regression as the most common nonparametric regression method ([43]), and applying the kernel trick to our feature-mapped linear memory. Let $k: \mathbb{R}^{{D_k}\times {D_k}} \to \mathbb{R}$ be a Mercer kernel such that $k(\mathbf{k}_i, \mathbf{k}_j) = \phi(\mathbf{k}_i)^\top\phi(\mathbf{k}j) = (\Phi\Phi ^\top){ij}$ where $\Phi = \begin{bmatrix} \phi(\mathbf{k}1) \ldots \phi(\mathbf{k}t) \end{bmatrix}^\top \in \mathbb{R}^{t\times {D\phi}} $ and $ \phi: \mathbb{R}^{D_k}\to \mathbb{R}^{D\phi} $ is a possibly infinite-dimensional feature map. Following our recipe for model design, we make the following design choices:

1. Choice of weights: assign equal weight $ \gamma_i^{(t)} = 1$ to each association

$ {m}{k, , t} = \operatorname{\arg!\min}{{m}\in \mathcal{M}k} \frac{1}{2}\sum{i=1}^t \lVert \mathbf{v}_i - \mathbf{M} \phi(\mathbf{k}_i)\rVert _2^2, $

where $\phi$ is possibly an infinite-dimensional feature map. 2. Choice of function class: linear functions $\mathcal{M}{k}= {{m}\mid {m}(\mathbf{k}) = \mathbf{M} \phi(\mathbf{k}), \ \mathbf{M} \in \mathbb{R}^{{D_v}\times {D\phi}}}$. 3. Choice of minimization algorithm: analytical solution.

The output of the optimal kernelized nonlinear associative memory induced by the kernel $k$ at timestep $t$ is the standard kernel regression prediction equation ([43])

$ {m}_{k, , t}(\mathbf{q}) = \mathbf{V}_t ^\top k(\mathbf{K}_t, \mathbf{K}_t)^{-1} k(\mathbf{K}_t, \mathbf{q}). $

Here we follow standard shorthand notation from the kernel learning literature: $k(\mathbf{K}_t, \mathbf{K}_t)$ is a shorthand for the matrix with $k(\mathbf{k}_i, \mathbf{k}_j)$ in entry $(i, j)$; $k(\mathbf{K}_t, \mathbf{q})$ is a shorthand for the vector with $k(\mathbf{k}_i, \mathbf{q})$ in entry $i$. When we use an exponential kernel $k(\mathbf{k}_i, \mathbf{k}_j) = \exp(\mathbf{k}_i ^\top \mathbf{k}_j/\sqrt{{D_k}})$, corresponding to an infinite-dimensional feature map, and make the crude approximation $k(\mathbf{K}_t, \mathbf{K}_t) \approx \mathbf{I}$, we obtain

$ \begin{aligned} {m}_{k, t}(\mathbf{q}) &= \mathbf{V}_t ^\top k(\mathbf{K}_t, \mathbf{K}_t)^{-1} k(\mathbf{K}_t, \mathbf{q}) \ &\approx \mathbf{V}_t ^\top k(\mathbf{K}t, \mathbf{q}) \ &= \sum{i=1}^t \mathbf{v}_i, k(\mathbf{k}i, \mathbf{q}) \ &= \sum{i=1}^t \mathbf{v}_i \exp\left(\frac{\mathbf{k}_i ^\top \mathbf{q}}{\sqrt{{D_k}}}\right), \end{aligned} $

an associative memory map equivalent to an unnormalized softmax attention. This relationship with kernel regression has been used by many previous works to derive more efficient versions of self-attention ([18, 79, 45, 80, 23]), borrowing techniques from kernel approximation methods. However unnormalized softmax attention is empirically known to be unstable to train ([80]). Thus we do not stop our derivations here; instead, we will see how the standard normalized softmax attention comes from another kind of nonparametric regressor.

Softmax attention with QKNorm is a locally-constant nonparametric regressor.

Local polynomial regression is another classic class of nonparametric estimators ([81]), which have been historically used as nonlinear smoothing functions in data analysis ([82]). Here, we consider the function class $\mathcal{M}^p$, which is the set of all vector-valued functions that are locally approximated by a degree $p$ polynomial around each point $\mathbf{q} \in \mathbb{R}^{D_k}$. Here locality is measured by some similarity function $s$ that only depends on the distance to $\mathbf{q}$, i.e. $s(\mathbf{k}, \mathbf{q}) = s(\lVert \mathbf{k} - \mathbf{q}\rVert 2)$. A classic example is the exponential smoothing kernel $s{\text{exp}, , B}(\mathbf{k}, \mathbf{q}) = \exp(-\lVert \mathbf{k} - \mathbf{q}\rVert _2^2/B)$ where $B$ is a parameter controlling the local neighborhood size around $\mathbf{q}$. Local polynomial regression is most often discussed for one-dimensional data, but here we generalize it to the multivariate case.

Nonparametric regression with local polynomial estimators corresponds to the following design choices:

1. Choice of weights: assign weights to associations based on similarity to $\mathbf{q}$,

$ {m}{p, , t}(\mathbf{q}) = \operatorname{\arg!\min}{{m}\in \mathcal{M}^p} \frac{1}{2}\sum_{i=1}^t \gamma_i^{(t)} \lVert \mathbf{v}_i - {m}(\mathbf{k}_i)\rVert _2^2 $

with weights given by $\gamma_i^{(t)}=s(\mathbf{k}_i, \mathbf{q})=s(\lVert \mathbf{k}_i - \mathbf{q}\rVert _2)$. 2. Choice of function class: $\mathcal{M}^p$, the class of functions ${m}_p$ such that for all points $\mathbf{q}\in \mathbb{R}^{D_k}$, the function can be locally approximated by a polynomial

$ \hat{{m}}\mathbf{q}(\mathbf{k}) = \mathbf{M}^{[0]} + \mathbf{M}^{[1]}(\mathbf{k}- \mathbf{q}) + \mathbf{M}^{[2]}(\mathbf{k} - \mathbf{q}, \mathbf{k} - \mathbf{q}) + \ldots + \mathbf{M}^{[p]}(\underbrace{\mathbf{k} - \mathbf{q}, \ldots, \mathbf{k} - \mathbf{q}}{p\ \text{arguments}}) $

where $\mathbf{M}^{[j]}$ is an order $j$ multilinear map. Intuitively, $\mathcal{M}^p$ is the class of functions that can be well approximated by an order $p$ Taylor expansion around each point. Consequently, the value retrieved from ${m}{p, , t}$ at a point $\mathbf{q}$ is given by the constant term of the polynomial: ${m}{p}(\mathbf{q}) = \hat {m}\mathbf{q}(\mathbf{q}) = \mathbf{M}^{[0]}$. 3. Choice of minimization algorithm: analytical solution. Solving for the minimum is equivalent to solving for the set of tensors $(\mathbf{M}^{[0]}, \ldots, \mathbf{M}^{[p]})$ that minimize the least squares objective. Note that we need to jointly solve for all of $(\mathbf{M}^{[0]}, \ldots, \mathbf{M}^{[p]})$ to get the retrieved value $\mathbf{M}^{[0]} = {m}{p, , t}(\mathbf{q})$.

Let us consider the simplest case of $p=0$, corresponding to the class of locally constant functions around $\mathbf{q}$, also known as Nadaraya-Watson estimators ([82, 83, 84]). With this simplification, the output of our nonparametric associative memory is

$ \begin{aligned} m_{0, , t}(\mathbf{q}) &= \operatorname{\arg!\min}{\mathbf{M}^{[0]}\in \mathbb{R}^{D_k}} \frac{1}{2}\sum{i=1}^t s(\mathbf{k}_i, \mathbf{q}) \lVert \mathbf{v}_i - \mathbf{M}^{[0]}\rVert 2^2 = \sum{i=1}^t \frac{s(\mathbf{k}i, \mathbf{q})}{\sum{j=1}^t s(\mathbf{k}_j, \mathbf{q})} \mathbf{v}_i. \end{aligned} $

We may be disappointed at first that $s$ needs to be a monotonic function of the distance $\lVert \mathbf{k} - \mathbf{q}\rVert _2$ to enforce the local approximation property of the estimator in; in contrast, softmax attention relies on an exponentiated inner product. However, if we normalize keys and queries such that $\lVert \mathbf{k}\rVert _2=\lVert \mathbf{q}\rVert _2=1$, the exponential smoothing kernel with bandwidth $B=2\sqrt{{D_k}}$ is equivalent to the exponential scaled dot product:

$ s_{\text{exp}, 2\sqrt {D_k}}(\mathbf{k}, \mathbf{q}) =\exp\left(\frac{-\lVert \mathbf{k} - \mathbf{q}\rVert _2^2}{2\sqrt{{D_k}}}\right) = \exp\left(\frac{-\lVert \mathbf{k}\rVert _2^2 - \lVert \mathbf{q}\rVert _2^2 + 2 \mathbf{k} ^\top \mathbf{q}}{2\sqrt{{D_k}}}\right) \propto \exp\left(\frac{\mathbf{k} ^\top \mathbf{q}}{\sqrt{{D_k}}}\right), $

allowing us to exactly derive softmax attention in. Moreover, this preprocessing step lines up with an existing technique called query-key normalization (QKNorm), empirically known to stablize the training of large Transformers ([39, 40]) but lacks theoretical justification. Our derivation provides a simple explanation: softmax attention with QKNorm is a locally zeroth order non-parametric approximation of the true key-value map.

By deriving self-attention from the same framework as the previous recurrent sequence layers, we can more clearly compare them through the perspective of how they do associative recall. Previous sequence layers that we derived maintained an explicit matrix-valued recurrent state to compress the prefix key-value pairs, relying on this constant-sized matrix to perform retrieval. In contrast, self-attention stores the entire set of prefix key-value pairs as its associative memory system, using all of them for retrieval. As a result, its memory (and expressiveness) grows with the sequence length, a typical property of non-parametric methods ([85, 81]).

Higher-order attention.

the section suggests a path to higher-order generalizations of self-attention by considering $p > 0$ at the cost of more computation. By reducing to a weighted least squares problem, we see the $p=1$ case of higher-order softmax self-attention is

$ \begin{aligned} \begin{bmatrix} \mathbf{M}^{[0]} & \mathbf{M}^{[1]} \end{bmatrix} &= (\mathbf{V}_t ^\top \mathbf{W}_t \mathbf{X}_t)(\mathbf{X}_t ^\top \mathbf{W}_t \mathbf{X}_t)^{-1}, & \mathbf{W}_t &\triangleq \operatorname{diag}\left(s(\mathbf{k}_1, \mathbf{q}), \ldots, s(\mathbf{k}_t, \mathbf{q})\right) \ \mathbf{X}_t &= \begin{bmatrix} 1 & \ldots & 1\ \mathbf{k}_1 - \mathbf{q} & \ldots & \mathbf{k}_t - \mathbf{q} \end{bmatrix}^\top & \mathbf{V}_t &= \begin{bmatrix} \mathbf{v}_1 & \ldots & \mathbf{v}_t \end{bmatrix}^\top, \end{aligned} $

where $\mathbf{W}_t\in \mathbb{R}^{t\times t}, \ \mathbf{X}_t \in \mathbb{R}^{t\times (1 + {D_k})}$, $\mathbf{M}^{[0]} \in \mathbb{R}^{{D_v}}$, and $\mathbf{M}^{[1]} \in \mathbb{R}^{{D_v} \times {D_k}}$. These equations correspond to the standard prediction equations for local linear regression ([82]).

An alternate form for the prediction $\mathbf{M}^{[0]}$ that is more intuitive is

$ \begin{aligned} \mathbf{M}^{[0]} &= \bar{\mathbf{v}} - \mathbf{M}^{[1]}\bar{\mathbf{d}}, & \bar{\mathbf{v}} &\triangleq \sum_{i=1}^t \frac{s(\mathbf{k}i, \mathbf{q})}{\sum{j=1}^t s(\mathbf{k}_j, \mathbf{q})} \mathbf{v}i, & \bar{\mathbf{d}} &\triangleq \sum{i=1}^t \frac{s(\mathbf{k}i, \mathbf{q})}{\sum{j=1}^t s(\mathbf{k}_j, \mathbf{q})} \mathbf{d}_i, & \mathbf{d}_i &\triangleq \mathbf{k}_i - \mathbf{q}, \end{aligned} $

resulting in an offset softmax attention as the output: $m_{1, , t}(\mathbf{q}) = \mathbf{M}^{[0]} = \bar{\mathbf{v}} - \mathbf{M}^{[1]}\bar{\mathbf{d}}$. This alternate form comes from the fact that if we already know $\mathbf{M}^{[1]}$, then we can solve the $p=0$ case with an alternate regression problem with preprocessed data ${(\mathbf{k}_i, \mathbf{v}_i - \mathbf{M}^{[1]}\mathbf{d}i)}{i=1}^t$, resulting in $\mathbf{M}^{[0]}$ being a normalized weighted sum of $\mathbf{v}_i - \mathbf{M}^{[1]}\mathbf{d}_i$. Unlike standard softmax attention (the $p=0$ case), which only considers the interaction between keys and queries, the $p=1$ case in also accounts for the covariance between keys through the inverse term, similar to how linear regression improves upon linear attention from before.

Unfortunately, directly computing for $\mathbf{q}\in {\mathbf{q}_1, \ldots, \mathbf{q}_T}$ in parallel to fully leverage hardware-acceleration is impractical for all but the smallest sequence problems. Materializing all $T$ prefix set of input tokens ${\mathbf{X}t}{t=1}^T$ requires $\mathcal{O}\left(T^2 {D_k} \right)$ space, and materializing all $T$ prefix weighted sums of outer products ${\mathbf{X}_t ^\top \mathbf{W}_t \mathbf{X}t}{t=1}^T$ requires $\mathcal{O}\left(T {D_k}^2 \right)$ space. Computing ${\mathbf{X} ^\top_t \mathbf{W}_t \allowbreak \mathbf{X}t}{t=1}^T$ requires $\mathcal{O}\left(T^2 {D_k} \right)$ time, and solving a batch of $T$ linear systems requires $\mathcal{O}\left(T {D_k}^3 \right)$ time. These challenges makes the current form computationally impractical for practical sequence tasks. Developing hardware-aware optimizations to compute similar to FlashAttention ([86]) is a promising future direction.

5. Online regression with sequence layers

Section Summary: The section demonstrates that various sequence layers can implicitly solve an online regression task—predicting the next value vector from a stream of key-value pairs—by performing a single forward pass with no training. Using synthetic data from a non-stationary autoregressive process that switches behavior midway, the authors compare parametric methods such as linear attention and recursive least squares against nonparametric ones like softmax attention and its local-linear variant, measuring one-step prediction error on held-out queries. Results show that layers able to discount older or dissimilar examples adapt quickly to the change in dynamics, while unweighted or correlation-ignoring methods incur sustained higher error.

We now empirically illustrate that the sequence layers discussed so far do in fact perform regression implicitly via a single forward pass. We generate a sequence of keys according to a non-stationary, switching, autoregressive process:

$ \begin{aligned} \mathbf{k}{t+1} = \begin{cases} 0.9 \mathbf{k}{t} + 0.02 \bm{\varepsilon}t, , t < T/4 \ 0.999 \mathbf{k}{t} + 0.02 \bm{\varepsilon}_t, , t \geq T/4 \end{cases} \end{aligned} $

with ${D_k}=64$ and $T=256$. Each $\bm{\varepsilon}_t \sim \operatorname{Normal}\left(0, \boldsymbol{\Sigma}\right)$ where $\boldsymbol{\Sigma}$ has unit diagonals and non-zero off-diagonals. The choice of a non-identity $\Sigma$ creates instantaneous correlation between the dimensions of the keys so that $\mathbf{K}_t ^\top \mathbf{K}t \neq \mathbf{I}$. We set the value vectors to be $\mathbf{v}t = \mathbf{k}{t+1}/\lVert \mathbf{k}{t+1}\rVert _2$, such that the dataset ${(\mathbf{k}_i, \mathbf{v}i)}{i=1}^T$ imitates a next-token prediction task. We make one-step-ahead predictions by setting $\mathbf{q}t = \mathbf{k}{t+1}$ so that $\mathbf{y}_t = {m}_t(\mathbf{q}_t) = {m}t(\mathbf{k}{t+1})$ is the test-time regressor's out-of-sample prediction of the next timestep, conditioned on the prefix ${(\mathbf{k}_i, \mathbf{v}i)}{i=1}^t$. To solve this online regression problem, we simply compute a single forward pass through each sequence layer.

We evaluate layers derived both from parametric regression and non-parametric regression. For our parametric regression layers we choose linear attention (the section) and recursive least squares (the section) to show the effects of accounting for correlations between keys. For our non-parametric regression layers, we chose softmax attention (Section 4.1) and our $p=1$ generalization of softmax attention based on local linear regression (the section). We also show the naive one-step forecaster baseline where we simply use $\mathbf{v}_{t-1}$ as the prediction given $\mathbf{k}_t$.

We plot the one-step out-of-sample test loss $\lVert \mathbf{v}_{t+1} - {m}t(\mathbf{k}{t+1})\rVert _2^2$ in Figure 2. During the initial $T/4$ timesteps, the dynamical system has a faster decay time with more rapid changes, making it less predictable. In the latter $3T/4$ timesteps, the system slows down and is more predictable. However, unweighted recursive least squares is unable to adapt to this change since it cannot discount datapoints from before the transition; as a result, it maintains a high error throughout the latter timesteps. In contrast, the nonparametric regressors are able to adapt quickly to the transition since earlier keys are dissimilar to the later keys. Interestingly, increasing our nonparametric softmax regressor from $p=0$ to $p=1$ allows the model to better adapt to both the initial fast-changing period and the sudden regime change. As expected, linear attention performs the worst since $\mathbf{K}_t ^\top \mathbf{K}_t \neq \mathbf{I}$ in our dynamical system[^2].

[^2]: We used a standard output normalization for linear attention ([42]) to prevent the output from scaling linearly with the sequence length.

6. Constructing effective key-value pairs for next-token recall

Section Summary: The section stresses that even capable models for associative recall depend critically on how key-value pairs are constructed from the input at test time, following the principle that model quality is only as good as its data. It shows that adding a simple short convolution—essentially blending the current token with the immediately prior one—lets even basic regression-based layers like linear attention capture bigram-style pairings and solve multi-query recall tasks, mimicking the useful “look-back” behavior of induction heads. Performance is ultimately bounded by memory capacity to store these pairs rather than by sequence length, since the underlying associations remain fixed.

Having shown that models derived via regression are able to perform associative recall, we now discuss the importance of constructing test-time key-value pairs that are pertinent to the task at hand. This echoes an unchanging principle in machine learning: a well-designed model is only as effective as the data it processes.

Historically, query-key-value sequences were constructed via a linear projection on the corresponding input $\mathbf{x}_t$ at that timestep, e.g. $\mathbf{k}_t = \mathbf{W}_K \mathbf{x}_t$ ([12]). Recent work on recurrent models, starting with [87], identified the importance of performing a "short convolution" in addition to the per-timestep projection:

$ \mathbf{k}t = \sum{i=0}^{K-1} w_i \mathbf{W}K \mathbf{x}{t-i} $

where $K$ is the length of the short convolution filter. The short convolution imitates the induction heads behavior learned by self-attention layers near the start of the neural network, which allow Transformers to "look back" at previous tokens. Setting $w_0=i$ and $w_j=0$ for $j\neq 0$ in Section 6 enables other sequence layers to also "look back" at the token from timestep $i$.

The short convolution is a crucial component of purely recurrent language models; removing the short convolution results in severe performance drops ([30, 68, 35, 37]). Including a short convolution before each self-attention layer can even improve the performance of Transformers, by freeing the earlier self-attention layers from having to learn the induction head behavior ([88]). In this section, we provide a partial explanation for why this short convolution is so crucial, through our perspective of associative memory. Specifically, we show that this short convolution allows the sequence layer to memorize and retrieve bigram-like key-value pairs, which has been shown to be important for language modeling ([15]).

One short conv is all you need for next-token recall.

Consider a standard associative recall task introduced by [15], in which a model must be able to recall multiple key-value pairings, known as multi-query associative recall (MQAR). In this task, there are $P$ pairs of unique and consistent cue-response[^3] tokens ${(\mathbf{u}j, \mathbf{v}j)}{j\in [P]}$ with each cue $\mathbf{u}j$ mapping one-to-one to its response $\mathbf{v}j$. The model is given a contextual sequence of randomly drawn, possibly repeated, cue-response pairs followed by a previously seen cue: $(\mathbf{x}1, \mathbf{x}2, \ldots, \mathbf{x}{T-1}, \mathbf{x}T, \mathbf{x}{T+1}) = (\mathbf{u}{i_1}, \mathbf{v}{i_1}, \ldots, \mathbf{u}{i{T/2}}, \mathbf{v}{i{T/2}}, \mathbf{u}{i_j})$. $\mathbf{u}{i_j}$ is a cue that appeared earlier (the $j$ th pair) and the model is expected to output the paired response $\mathbf{v}_{i_j}$ at timestep $T+1$. This benchmark abstracts the "Hakuna Matata" example "Hakuna Matata, it means no worries for the rest of your days. Hakuna " from [15], which we also discussed in Section 1.

[^3]: We use the "cue" and "response" terminology to avoid confusion with "keys" and "values" in the context of test-time regression.

We now show that a single test-time regression layer (e.g. as simple as linear attention) with one short convolution is sufficient to solve MQAR, without any parameters other than the embeddings, as long as it is provided with the appropriate key-value pairs to regress.

To solve this task, it suffices to test-time memorize all the bigram pairing ${(\mathbf{u}_j, \mathbf{v}_j)}$. Given $\mathbf{x}1, \ldots, \mathbf{x}T$, we causally construct our test-time regression dataset as follows. Let the keys be constructed via a short convolution $\mathbf{k}t = w_0 \mathbf{x}t + w_1 \mathbf{x}{t-1} = \mathbf{x}{t-1}$ where $w_0 = 0, w_1 = 1$. Let the queries and values be the same as the inputs: $\mathbf{q}t = \mathbf{v}t = \mathbf{x}t$. In other words, our AR sequence layer will first memorize ${(\mathbf{k}t, \mathbf{v}t)}{t=1}^{T+1} = {(\mathbf{x}{t-1}, \mathbf{x}t)}{t=1}^{T+1}$ into its associative memory map ${m}{T+1}$. Then we retrieve from ${m}{T+1}$ using $\mathbf{q}{T+1} = \mathbf{x}{T+1} = \mathbf{u}{i_j}$, which has previously appeared in the sequence of cues as the $i_j$ th bigram. The output is then

$ \mathbf{y}{T+1} = {m}{T+1}(\mathbf{q}{T+1}) = {m}{T+1}(\mathbf{u}_{i_j}). $

For simplicity, suppose ${m}_{T+1}$ is implemented via linear attention, following. Then

$ \mathbf{y}{T+1} = {m}{T+1}(\mathbf{u}{i_j}) = \left(\sum{t=1}^{T+1}\mathbf{v}t \mathbf{k}t ^\top\right) \mathbf{u}{i_j} = \sum{t=1}^{T+1} \mathbf{x}t \mathbf{x}{t-1}^\top \mathbf{u}_{i_j}. $

When the embedding space is large enough, we can set the embeddings such that all tokens are orthonormal. Then $\mathbf{x}{t-1}^\top \mathbf{u}{i_j}$ is nonzero only when $\mathbf{x}{t-1} = \mathbf{u}{i_j}$, simplifying the output to a sum of only tokens that are followed by $\mathbf{u}{i_j}$. Since the cue-response map is one-to-one, the remaining tokens are all $\mathbf{v}{i_j}$, producing the output $\mathbf{y}{T+1} \propto \mathbf{v}{i_j}$, solving the MQAR task.

Memory capacity, not sequence length, limits model performance.

Typical evaluations of models on the MQAR task look at model performance with respect to the length of the sequence, possibly varying the model capacity. A few examples include Figure 2 of [15], Figure 8 of [33], Figure 2 of [68], and likely others. However, our above construction shows that the sequence length doesn't matter for memorizers with a large enough memory capacity; once the memory is large enough, the difficulty of MQAR is independent of the sequence length $T$, which may be counter to the intuitive notion of longer sequences being more difficult. Instead, models are limited by their ability to fully memorize all $P$ key-value pairs.

Although our construction in relied on the embedding space being large enough to allow orthonormality, trained models can still succeed without perfectly orthonormal embeddings, which we empirically show in Figure 3. We train on the MQAR task following the procedure of ([15]). However, we use only one test-time regression layer and one length-2 convolution (for constructing the keys), following our construction. We also remove the typical MLP block following the sequence layer, since our construction shows that additional nonlinearities are unnecessary for this synthetic associative recall task. We fix our key dimensions ${D_k}$ to be the same as the embedding size, $d_{\text{model}} = 64$. We compare two test-time regression architectures, linear attention and unweighted recursive least squares (RLS) of. We do not use gated variants of linear attention since this task does not require any forgetting; the cue-response mapping is unchanged throughout the sequence. As shown by our previous construction, even linear attention is able to solve the task perfectly when the appropriate keys are constructed. In the case where $d_{\text{model}} = P=64$, linear attention solves MQAR perfectly, independent of sequence length as predicted. Once the number of keys increases to $P=128$, linear attention can no longer perfectly solve MQAR, since it is impossible for 128 keys to be orthogonal in a $d_{\text{model}}=64$ dimensional space. In contrast, an unweighted RLS layer (the section) improves upon linear attention by accounting for the key covariance matrices ${\mathbf{K}_t ^\top \mathbf{K}t}$, allowing it to better test-time memorize the cue-response map, even when $P > d{\text{model}}$. As a result, RLS is able to solve MQAR regardless of the sequence length.

7. Related works

Section Summary: The authors trace their test-time regression approach to a long tradition of building neural networks around associative memory, an idea explored since the 1970s and later linked to ideas from neuroscience and dynamical systems. They note that the capacity to recall earlier tokens is central to in-context learning in large language models and review studies showing that transformers can implicitly run optimization steps across layers. Their own linear-attention layers are presented as an explicit, built-in version of this process that differs from recurrent models performing gradient updates inside each layer.

Our test-time regression framework was derived by formalizing the notion of associative recall, resulting in sequence layers that first memorize then retrieve. Designing neural networks motivated with associative memory has a long history, dating back as early as the works of [89], [90], [91], [55], and [57]. Many more contemporary works have also sought to understand and derive neural network layers from the perspective of memory, drawing inspiration from both neuroscience and dynamical systems ([92, 93, 94, 95, 96, 38]). Indeed, even the feedforward layers of an MLP can be seen as a persistent associative memory ([97, 98, 92, 99, 100]).

The ability to reference and retrieve past tokens for further computation is believed to be key to in-context learning and zero-shot/few-shot learning in large language models ([14, 101]). There have been many attempts at understanding how sequence models are able to learn in-context ([102, 103]). Some works have shown that Transformers, due to their self-attention layers, learn linear functions in-context by learning to perform gradient descent, where the iterates are updated per layer ([104, 105]). This phenomenon is sometimes known as mesa-optimization ([62]) where the model learns to solve an auxiliary optimization problem in order to minimize its training objective. However, sometimes these intermediate optimization processes are actually implicitly in the architecture itself, as we showed in where each linear attention layer (and its gated variant) explicitly performs one step of gradient descent. Note that this behavior is different from the class of recurrent models that perform stochastic gradient descent within the layer, as defined by.

8. Conclusion and discussion

Section Summary: The paper introduces a unifying framework showing how sequence models can carry out associative recall by performing regression during testing, with different architectures arising naturally from choices about weighting associations, function classes, and optimization methods. This perspective explains existing attention mechanisms, derives useful generalizations, and demonstrates that simple components like a short convolution plus one sequence layer can handle complex recall tasks. The authors note the framework's limited scope and outline promising directions for future work on nonlinear test-time regressors and hardware-efficient implementations that could enable more adaptive models.

We have presented a unifying framework that derives sequence models that can explicitly perform associative recall via test-time regression. Each architecture emerges naturally by specifying three aspects: the relative importance of each association, the function class, and the optimization algorithm. Through this lens, we derived linear attention and its variants as approximations to the linear least squares solution, online learning layers and fast-weight layers as first-order least-squares solvers, and softmax attention as a local constant regressor. This perspective provided principled explanations for architectural choices like QKNorm and derived higher-order generalizations of attention that go beyond pairwise interactions. We also demonstrated that for multi-query associative recall (MQAR), a single short convolution and a sequence layer suffices, in contrast to typical approaches that use multiple layers and unnecessary feedforward blocks ([15]).

Our paper focuses specifically on sequence architectures that use query-key-value patterns for associative recall. Several important architectural families lie outside this scope, including structured state-space models with more general structured masks ([33]), convolutional architectures ([106, 107, 108, 48, 109, 110, 46]), and more novel computational patterns ([111, 112]). Moreover, complementary research on model backbones and initialization continues to yield important advances and performance improvements ([32, 33, 25, 113, 114]).

Our framework opens rich directions for future research by connecting to the extensive literature on regression and optimization. While initial successes like TTT-MLP ([37]) hint at the potential of neural test-time regressors, the space of nonlinear neural regression models remains largely unexplored. Recent work integrating weight decay and momentum ([38, 35]) suggests opportunities for developing more effective, stable, and parallelizable test-time optimizers. Concurrent work by [38] also suggests opportunities to more creatively utilize associative memory. The practical success of these approaches will depend heavily on efficient hardware implementations ([114, 68, 37, 33, 86]), aligning with established compute-scaling principles ([115, 116, 117]).

In the words of [56], "associative memory is a very delicate and complex concept which often has been attributed to the higher cognitive processes, especially those taking place in the human brain". We speculate that test-time associative memory, along with test-time compute more broadly ([37, 118, 119]), will be fundamental to developing truly adaptive models that can update and learn in changing environments ([120]).

Acknowledgements.

This work was supported in part by AFOSR Grant FA9550-21-1-0397, ONR Grant N00014-22-1-2110, NSF Grant 2205084, the Stanford Institute for Human-Centered Artificial Intelligence (HAI). EBF is a Chan Zuckerberg Biohub – San Francisco Investigator.

Appendix

Section Summary: The appendix derives a mathematical bound on the size of the output vector produced by a linear regression layer expressed in terms of attention-like operations. It shows that this output is controlled by the input query magnitude, summed key and value magnitudes, and a minimum eigenvalue term in the denominator that provides a natural scaling effect. Approximations used in linear attention remove this built-in normalization, which helps explain why explicit output normalization techniques are sometimes added back in.

appendixpage

A. Bounding the norm of a linear regression layer

Let $\mathbf{y}_t = \mathbf{V}_t ^\top \mathbf{K}_t(\mathbf{K}_t ^\top \mathbf{K}_t)^{-1} \mathbf{q}_t$ and let $\lVert \mathbf{A}\rVert _2$ indicate the spectral norm of a matrix $\mathbf{A}$, a matrix norm induced by the L2 norm on vectors. Then

$ \begin{aligned} \lVert \mathbf{y}_t\rVert _2 &= \lVert \mathbf{V}_t ^\top \mathbf{K}_t(\mathbf{K}_t ^\top \mathbf{K}_t)^{-1} \mathbf{q}_t\rVert _2\ &\leq \lVert \mathbf{q}_t\rVert 2\sup{\lVert \mathbf{q}\rVert _2=1} \lVert \mathbf{V}_t ^\top \mathbf{K}_t (\mathbf{K}_t ^\top \mathbf{K}_t)^{-1}\mathbf{q}\rVert _2 \ & = \lVert \mathbf{q}_t\rVert _2\lVert \mathbf{V}_t ^\top \mathbf{K}_t (\mathbf{K}_t ^\top \mathbf{K}_t)^{-1}\rVert _2 \ &\leq \lVert \mathbf{q}_t\rVert _2\lVert \mathbf{V}_t ^\top \mathbf{K}_t\rVert _2\lVert (\mathbf{K}_t ^\top \mathbf{K}_t)^{-1}\rVert _2 \ & = \lVert \mathbf{q}_t\rVert 2\lVert \sum{i=1}^t \mathbf{v}_i \mathbf{k}_i ^\top\rVert _2\lVert (\mathbf{K}_t ^\top \mathbf{K}_t)^{-1}\rVert _2 \ &\leq \lVert \mathbf{q}_t\rVert 2\left(\sum{i=1}^t\lVert \mathbf{v}_i\rVert _2\lVert \mathbf{k}_i\rVert _2\right)\lVert (\mathbf{K}_t ^\top \mathbf{K}_t)^{-1}\rVert _2 \ & = \lVert \mathbf{q}_t\rVert 2\left(\sum{i=1}^t\lVert \mathbf{v}_i\rVert _2\lVert \mathbf{k}_i\rVert _2\right)\max_i \lambda_i((\mathbf{K}_t ^\top \mathbf{K}_t)^{-1}) \ & = \lVert \mathbf{q}_t\rVert 2\left(\sum{i=1}^t\lVert \mathbf{v}_i\rVert _2\lVert \mathbf{k}_i\rVert _2\right)\max_i \frac{1}{\lambda_i(\mathbf{K}_t ^\top \mathbf{K}_t)} \ & = \frac{\lVert \mathbf{q}_t\rVert 2\sum{i=1}^t\lVert \mathbf{v}_i\rVert _2\lVert \mathbf{k}_i\rVert 2}{\lambda{\text{min}}(\mathbf{K}_t ^\top \mathbf{K}_t)}. \end{aligned} $

Hence $\lVert \mathbf{y}_t\rVert _2 \leq \lVert \mathbf{q}_t\rVert 2\sum{i=1}^t\lVert \mathbf{v}_i\rVert _2\lVert \mathbf{k}_i\rVert 2/\lambda{\text{min}}(\mathbf{K}_t ^\top \mathbf{K}_t)$. When we approximate $\mathbf{K}_t ^\top \mathbf{K}t \approx \mathbf{I}$, as in linear attention, we lose the self-normalizing property of dividing by ${\lambda{\text{min}}(\mathbf{K}_t ^\top \mathbf{K}_t)}$ since the denominator becomes 1. This explains how output normalization by [42] is an attempt at restoring this intrinsic self-normalizing property of linear regression.

References

Section Summary: The references section compiles a list of influential academic papers on artificial intelligence and machine learning. It highlights foundational work on neural networks, transformer models for handling sequences like text or images, and large language models capable of few-shot learning. Later entries extend these ideas to practical uses such as time-series forecasting, video analysis, and biological predictions like protein structures or genetic sequences.

[1] I. Sutskever, O. Vinyals, and Q. V. Le. Sequence to Sequence Learning with Neural Networks. In Advances in Neural Information Processing Systems, volume 27. Curran Associates, Inc., 2014. URL https://proceedings.neurips.cc/paper/2014/hash/a14ac55a4f27472c5d894ec1c3c743d2-Abstract.html.

[2] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. In J. Burstein, C. Doran, and T. Solorio, editors, Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pages 4171–4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi:10.18653/v1/N19-1423. URL https://aclanthology.org/N19-1423/.

[3] T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert-Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. M. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei. Language Models are Few-Shot Learners, July 2020. URL http://arxiv.org/abs/2005.14165. arXiv:2005.14165 [cs].

[4] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, J. Uszkoreit, and N. Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International conference on learning representations, 2021. URL https://openreview.net/forum?id=YicbFdNTTy.

[5] G. Bertasius, H. Wang, and L. Torresani. Is Space-Time Attention All You Need for Video Understanding? In Proceedings of the 38th International Conference on Machine Learning, pages 813–824. PMLR, July 2021. URL https://proceedings.mlr.press/v139/bertasius21a.html. ISSN: 2640-3498.

[6] D. Salinas, V. Flunkert, J. Gasthaus, and T. Januschowski. DeepAR: Probabilistic forecasting with autoregressive recurrent networks. International Journal of Forecasting, 36(3):1181–1191, July 2020. ISSN 0169-2070. doi:10.1016/j.ijforecast.2019.07.001. URL https://www.sciencedirect.com/science/article/pii/S0169207019301888.

[7] N. Gruver, M. Finzi, S. Qiu, and A. G. Wilson. Large Language Models Are Zero-Shot Time Series Forecasters. Advances in Neural Information Processing Systems, 36:19622–19635, Dec. 2023. URL https://proceedings.neurips.cc/paper_files/paper/2023/hash/3eb7ca52e8207697361b2c0fb3926511-Abstract-Conference.html.

[8] A. F. Ansari, L. Stella, A. C. Turkmen, X. Zhang, P. Mercado, H. Shen, O. Shchur, S. S. Rangapuram, S. P. Arango, S. Kapoor, J. Zschiegner, D. C. Maddix, H. Wang, M. W. Mahoney, K. Torkkola, A. G. Wilson, M. Bohlke-Schneider, and B. Wang. Chronos: Learning the language of time series. Transactions on Machine Learning Research, 2024. ISSN 2835-8856. URL https://openreview.net/forum?id=gerNCVqqtR.

[9] J. Jumper, R. Evans, A. Pritzel, T. Green, M. Figurnov, O. Ronneberger, K. Tunyasuvunakool, R. Bates, A. Žídek, A. Potapenko, A. Bridgland, C. Meyer, S. A. A. Kohl, A. J. Ballard, A. Cowie, B. Romera-Paredes, S. Nikolov, R. Jain, J. Adler, T. Back, S. Petersen, D. Reiman, E. Clancy, M. Zielinski, M. Steinegger, M. Pacholska, T. Berghammer, S. Bodenstein, D. Silver, O. Vinyals, A. W. Senior, K. Kavukcuoglu, P. Kohli, and D. Hassabis. Highly accurate protein structure prediction with AlphaFold. Nature, 596(7873):583–589, Aug. 2021. ISSN 1476-4687. doi:10.1038/s41586-021-03819-2. URL https://www.nature.com/articles/s41586-021-03819-2. Publisher: Nature Publishing Group.

[10] J. Zhou and O. G. Troyanskaya. Predicting effects of noncoding variants with deep learning–based sequence model. Nature Methods, 12(10):931–934, Oct. 2015. ISSN 1548-7105. doi:10.1038/nmeth.3547. URL https://www.nature.com/articles/nmeth.3547. Publisher: Nature Publishing Group.

[11] E. Nguyen, M. Poli, M. G. Durrant, B. Kang, D. Katrekar, D. B. Li, L. J. Bartie, A. W. Thomas, S. H. King, G. Brixi, J. Sullivan, M. Y. Ng, A. Lewis, A. Lou, S. Ermon, S. A. Baccus, T. Hernandez-Boussard, C. Ré, P. D. Hsu, and B. L. Hie. Sequence modeling and design from molecular to genome scale with Evo. Science, 386(6723):eado9336, Nov. 2024. doi:10.1126/science.ado9336. URL https://www.science.org/doi/10.1126/science.ado9336. Publisher: American Association for the Advancement of Science.

[12] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is All you Need. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https://papers.nips.cc/paper_files/paper/2017/hash/3f5ee243547dee91fbd053c1c4a845aa-Abstract.html.

[13] S. Rush. There are like 4 more linear RNN papers out today, Apr. 2024. URL https://x.com/srush_nlp/status/1780231813820002771.

[14] C. Olsson, N. Elhage, N. Nanda, N. Joseph, N. DasSarma, T. Henighan, B. Mann, A. Askell, Y. Bai, A. Chen, T. Conerly, D. Drain, D. Ganguli, Z. Hatfield-Dodds, D. Hernandez, S. Johnston, A. Jones, J. Kernion, L. Lovitt, K. Ndousse, D. Amodei, T. Brown, J. Clark, J. Kaplan, S. McCandlish, and C. Olah. In-context Learning and Induction Heads, Sept. 2022. URL http://arxiv.org/abs/2209.11895. arXiv:2209.11895 [cs].

[15] S. Arora, S. Eyuboglu, A. Timalsina, I. Johnson, M. Poli, J. Zou, A. Rudra, and C. Re. Zoology: Measuring and Improving Recall in Efficient Language Models. In The Twelfth International Conference on Learning Representations, Oct. 2023a. URL https://openreview.net/forum?id=LY3ukUANko.

[16] S. Arora, B. Yang, S. Eyuboglu, A. Narayan, A. Hojel, I. Trummer, and C. Ré. Language Models Enable Simple Systems for Generating Structured Views of Heterogeneous Data Lakes, Apr. 2023b. URL http://arxiv.org/abs/2304.09433. arXiv:2304.09433 [cs].

[17] A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret. Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. In Proceedings of the 37th International Conference on Machine Learning, pages 5156–5165. PMLR, Nov. 2020. URL https://proceedings.mlr.press/v119/katharopoulos20a.html. ISSN: 2640-3498.

[18] H. Peng, N. Pappas, D. Yogatama, R. Schwartz, N. Smith, and L. Kong. Random Feature Attention. In International Conference on Learning Representations, Oct. 2020. URL https://openreview.net/forum?id=QtTKTdVrFBB.

[19] Z. Qin, W. Sun, H. Deng, D. Li, Y. Wei, B. Lv, J. Yan, L. Kong, and Y. Zhong. cosFormer: Rethinking Softmax In Attention. In International Conference on Learning Representations, Oct. 2021. URL https://openreview.net/forum?id=Bl8CQrx2Up4.

[20] J. Kasai, H. Peng, Y. Zhang, D. Yogatama, G. Ilharco, N. Pappas, Y. Mao, W. Chen, and N. A. Smith. Finetuning Pretrained Transformers into RNNs. In M.-F. Moens, X. Huang, L. Specia, and S. W.-t. Yih, editors, Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 10630–10643, Online and Punta Cana, Dominican Republic, Nov. 2021. Association for Computational Linguistics. doi:10.18653/v1/2021.emnlp-main.830. URL https://aclanthology.org/2021.emnlp-main.830.

[21] M. Zhang, K. Bhatia, H. Kumbong, and C. Re. The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry. In The Twelfth International Conference on Learning Representations, Oct. 2023a. URL https://openreview.net/forum?id=4g02l2N2Nx.

[22] Y. Aksenov, N. Balagansky, S. M. L. C. Vaina, B. Shaposhnikov, A. Gorbatovski, and D. Gavrilov. Linear Transformers with Learnable Kernel Functions are Better In-Context Models, Feb. 2024. URL http://arxiv.org/abs/2402.10644. arXiv:2402.10644 [cs] version: 1.

[23] H. Chen, Liuzhicheng, X. Wang, Y. Tian, and Y. Wang. DiJiang: Efficient Large Language Models through Compact Kernelization. In Forty-first International Conference on Machine Learning, June 2024. URL {https://openreview.net/forum?id=0uUHfhXdnH&referrer=

[24] Y. Sun, L. Dong, S. Huang, S. Ma, Y. Xia, J. Xue, J. Wang, and F. Wei. Retentive Network: A Successor to Transformer for Large Language Models, Aug. 2023. URL http://arxiv.org/abs/2307.08621. arXiv:2307.08621 [cs].

[25] A. Orvieto, S. L. Smith, A. Gu, A. Fernando, C. Gulcehre, R. Pascanu, and S. De. Resurrecting Recurrent Neural Networks for Long Sequences. In Proceedings of the 40th International Conference on Machine Learning, pages 26670–26698. PMLR, July 2023. URL https://proceedings.mlr.press/v202/orvieto23a.html. ISSN: 2640-3498.

[26] T. Katsch. GateLoop: Fully Data-Controlled Linear Recurrence for Sequence Modeling, Jan. 2024. URL http://arxiv.org/abs/2311.01927. arXiv:2311.01927 [cs].

[27] S. De, S. L. Smith, A. Fernando, A. Botev, G. Cristian-Muraru, A. Gu, R. Haroun, L. Berrada, Y. Chen, S. Srinivasan, G. Desjardins, A. Doucet, D. Budden, Y. W. Teh, R. Pascanu, N. D. Freitas, and C. Gulcehre. Griffin: Mixing Gated Linear Recurrences with Local Attention for Efficient Language Models, Feb. 2024. URL http://arxiv.org/abs/2402.19427. arXiv:2402.19427 [cs].

[28] Z. Qin, S. Yang, W. Sun, X. Shen, D. Li, W. Sun, and Y. Zhong. HGRN2: Gated Linear RNNs with State Expansion. In First Conference on Language Modeling, Aug. 2024. URL https://openreview.net/forum?id=y6SqbJfCSk#discussion.

[29] B. Peng, D. Goldstein, Q. Anthony, A. Albalak, E. Alcaide, S. Biderman, E. Cheah, X. Du, T. Ferdinan, H. Hou, P. Kazienko, K. K. GV, J. Kocoń, B. Koptyra, S. Krishna, R. M. Jr, J. Lin, N. Muennighoff, F. Obeid, A. Saito, G. Song, H. Tu, C. Wirawan, S. Woźniak, R. Zhang, B. Zhao, Q. Zhao, P. Zhou, J. Zhu, and R.-J. Zhu. Eagle and Finch: RWKV with Matrix-Valued States and Dynamic Recurrence, Sept. 2024. URL http://arxiv.org/abs/2404.05892. arXiv:2404.05892.

[30] S. Yang, B. Wang, Y. Shen, R. Panda, and Y. Kim. Gated Linear Attention Transformers with Hardware-Efficient Training. In Proceedings of the 41st International Conference on Machine Learning, pages 56501–56523. PMLR, July 2024b. URL https://proceedings.mlr.press/v235/yang24ab.html. ISSN: 2640-3498.

[31] M. Beck, K. Pöppel, M. Spanring, A. Auer, O. Prudnikova, M. Kopp, G. Klambauer, J. Brandstetter, and S. Hochreiter. xLSTM: Extended Long Short-Term Memory, May 2024. URL http://arxiv.org/abs/2405.04517. arXiv:2405.04517 [cs, stat].

[32] A. Gu and T. Dao. Mamba: Linear-Time Sequence Modeling with Selective State Spaces. In First Conference on Language Modeling, Aug. 2024. URL https://openreview.net/forum?id=tEYskw1VY2#discussion.

[33] T. Dao and A. Gu. Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. In Forty-first International Conference on Machine Learning, June 2024.

[34] I. Schlag, K. Irie, and J. Schmidhuber. Linear Transformers Are Secretly Fast Weight Programmers. In Proceedings of the 38th International Conference on Machine Learning, pages 9355–9366. PMLR, July 2021. URL https://proceedings.mlr.press/v139/schlag21a.html. ISSN: 2640-3498.

[35] S. Yang, J. Kautz, and A. Hatamizadeh. Gated Delta Networks: Improving Mamba2 with Delta Rule, Dec. 2024a. URL http://arxiv.org/abs/2412.06464. arXiv:2412.06464 [cs] version: 1.

[36] B. Liu, R. Wang, L. Wu, Y. Feng, P. Stone, and Q. Liu. Longhorn: State Space Models are Amortized Online Learners, Oct. 2024. URL http://arxiv.org/abs/2407.14207. arXiv:2407.14207.

[37] Y. Sun, X. Li, K. Dalal, J. Xu, A. Vikram, G. Zhang, Y. Dubois, X. Chen, X. Wang, S. Koyejo, T. Hashimoto, and C. Guestrin. Learning to (Learn at Test Time): RNNs with Expressive Hidden States, Aug. 2024. URL http://arxiv.org/abs/2407.04620. arXiv:2407.04620.

[38] A. Behrouz, P. Zhong, and V. Mirrokni. Titans: Learning to Memorize at Test Time, Dec. 2024. URL http://arxiv.org/abs/2501.00663. arXiv:2501.00663 [cs].

[39] M. Dehghani, J. Djolonga, B. Mustafa, P. Padlewski, J. Heek, J. Gilmer, A. P. Steiner, M. Caron, R. Geirhos, I. Alabdulmohsin, R. Jenatton, L. Beyer, M. Tschannen, A. Arnab, X. Wang, C. R. Ruiz, M. Minderer, J. Puigcerver, U. Evci, M. Kumar, S. V. Steenkiste, G. F. Elsayed, A. Mahendran, F. Yu, A. Oliver, F. Huot, J. Bastings, M. Collier, A. A. Gritsenko, V. Birodkar, C. N. Vasconcelos, Y. Tay, T. Mensink, A. Kolesnikov, F. Pavetic, D. Tran, T. Kipf, M. Lucic, X. Zhai, D. Keysers, J. J. Harmsen, and N. Houlsby. Scaling Vision Transformers to 22 Billion Parameters. In Proceedings of the 40th International Conference on Machine Learning, pages 7480–7512. PMLR, July 2023. URL https://proceedings.mlr.press/v202/dehghani23a.html. ISSN: 2640-3498.

[40] M. Wortsman, P. J. Liu, L. Xiao, K. E. Everett, A. A. Alemi, B. Adlam, J. D. Co-Reyes, I. Gur, A. Kumar, R. Novak, J. Pennington, J. Sohl-Dickstein, K. Xu, J. Lee, J. Gilmer, and S. Kornblith. Small-scale proxies for large-scale Transformer training instabilities. In The Twelfth International Conference on Learning Representations, Oct. 2023. URL https://openreview.net/forum?id=d8w0pmvXbZ.

[41] Y. Tay, M. Dehghani, D. Bahri, and D. Metzler. Efficient Transformers: A Survey. ACM Comput. Surv., 55(6):109:1–109:28, Dec. 2022. ISSN 0360-0300. doi:10.1145/3530811. URL https://dl.acm.org/doi/10.1145/3530811.

[42] Z. Qin, X. Han, W. Sun, D. Li, L. Kong, N. Barnes, and Y. Zhong. The Devil in Linear Transformer. In Y. Goldberg, Z. Kozareva, and Y. Zhang, editors, Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 7025–7041, Abu Dhabi, United Arab Emirates, Dec. 2022. Association for Computational Linguistics. doi:10.18653/v1/2022.emnlp-main.473. URL https://aclanthology.org/2022.emnlp-main.473.

[43] B. Schölkopf and A. J. Smola. Learning with Kernels: Support Vector Machines, Regularization, Optimization, and Beyond. The MIT Press, Dec. 2001. ISBN 978-0-262-25693-3. doi:10.7551/mitpress/4175.001.0001. URL https://direct.mit.edu/books/monograph/1821/Learning-with-KernelsSupport-Vector-Machines.

[44] S. Arora, S. Eyuboglu, M. Zhang, A. Timalsina, S. Alberti, J. Zou, A. Rudra, and C. Re. Simple linear attention language models balance the recall-throughput tradeoff. In Proceedings of the 41st International Conference on Machine Learning, pages 1763–1840. PMLR, July 2024. URL https://proceedings.mlr.press/v235/arora24a.html. ISSN: 2640-3498.

[45] K. M. Choromanski, V. Likhosherstov, D. Dohan, X. Song, A. Gane, T. Sarlos, P. Hawkins, J. Q. Davis, A. Mohiuddin, L. Kaiser, D. B. Belanger, L. J. Colwell, and A. Weller. Rethinking Attention with Performers. In International Conference on Learning Representations, Oct. 2020. URL https://openreview.net/forum?id=Ua6zuk0WRH.

[46] S. Hochreiter and J. Schmidhuber. Long Short-Term Memory. Neural Computation, 9(8):1735–1780, Nov. 1997. ISSN 0899-7667. doi:10.1162/neco.1997.9.8.1735. URL https://ieeexplore.ieee.org/abstract/document/6795963. Conference Name: Neural Computation.

[47] A. Gu, T. Dao, S. Ermon, A. Rudra, and C. Ré. HiPPO: Recurrent Memory with Optimal Polynomial Projections. In Advances in Neural Information Processing Systems, volume 33, pages 1474–1487. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/paper/2020/hash/102f0bb6efb3a6128a3c750dd16729be-Abstract.html.

[48] A. Gu, K. Goel, and C. Re. Efficiently Modeling Long Sequences with Structured State Spaces. In International Conference on Learning Representations, Oct. 2021. URL https://openreview.net/forum?id=uYLFoz1vlAC.

[49] A. Gu, A. Gupta, K. Goel, and C. Ré. On the parameterization and initialization of diagonal state space models. In Proceedings of the 36th International Conference on Neural Information Processing Systems, NIPS '22, pages 35971–35983, Red Hook, NY, USA, 2022. Curran Associates Inc. ISBN 978-1-7138-7108-8.

[50] D. Han, Z. Wang, Z. Xia, Y. Han, Y. Pu, C. Ge, J. Song, S. Song, B. Zheng, and G. Huang. Demystify Mamba in Vision: A Linear Attention Perspective, May 2024. URL https://arxiv.org/abs/2405.16605v1.

[51] A. Trockman, H. Harutyunyan, J. Z. Kolter, S. Kumar, and S. Bhojanapalli. Mimetic Initialization Helps State Space Models Learn to Recall, Oct. 2024. URL http://arxiv.org/abs/2410.11135. arXiv:2410.11135 [cs].

[52] R. Ren, Z. Li, and Y. Liu. Can Mamba Always Enjoy the "Free Lunch"?, Oct. 2024. URL http://arxiv.org/abs/2410.03810. arXiv:2410.03810 [cs].

[53] J. Schmidhuber. Learning to Control Fast-Weight Memories: An Alternative to Dynamic Recurrent Networks. Neural Computation, 4(1):131–139, Jan. 1992. ISSN 0899-7667. doi:10.1162/neco.1992.4.1.131. URL https://doi.org/10.1162/neco.1992.4.1.131.

[54] K. Clark, K. Guu, M.-W. Chang, P. Pasupat, G. Hinton, and M. Norouzi. Meta-Learning Fast Weight Language Models. In Y. Goldberg, Z. Kozareva, and Y. Zhang, editors, Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 9751–9757, Abu Dhabi, United Arab Emirates, Dec. 2022. Association for Computational Linguistics. doi:10.18653/v1/2022.emnlp-main.661. URL https://aclanthology.org/2022.emnlp-main.661.

[55] T. Kohonen. Correlation Matrix Memories. IEEE Transactions on Computers, C-21(4):353–359, Apr. 1972. ISSN 1557-9956. doi:10.1109/TC.1972.5008975. URL https://ieeexplore.ieee.org/abstract/document/5008975. Conference Name: IEEE Transactions on Computers.

[56] T. Kohonen. Self-Organization and Associative Memory, volume 8 of Springer Series in Information Sciences. Springer, Berlin, Heidelberg, 1989. ISBN 978-3-540-51387-2 978-3-642-88163-3. doi:10.1007/978-3-642-88163-3. URL http://link.springer.com/10.1007/978-3-642-88163-3.

[57] G. E. Hinton and J. A. Anderson. Parallel Models of Associative Memory. Psychology Press, 1989.

[58] N. Shazeer. Fast Transformer Decoding: One Write-Head is All You Need, Nov. 2019. URL http://arxiv.org/abs/1911.02150. arXiv:1911.02150 [cs].

[59] J. Ainslie, J. Lee-Thorp, M. d. Jong, Y. Zemlyanskiy, F. Lebron, and S. Sanghai. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. In The 2023 Conference on Empirical Methods in Natural Language Processing, Dec. 2023. URL https://openreview.net/forum?id=hmOwOZWzYE.

[60] S. S. Haykin. Adaptive Filter Theory. Pearson, 2014. ISBN 978-0-13-267145-3. Google-Books-ID: J4GRKQEACAAJ.

[61] M. Garnelo and W. M. Czarnecki. Exploring the Space of Key-Value-Query Models with Intention, May 2023. URL http://arxiv.org/abs/2305.10203. arXiv:2305.10203.

[62] J. von Oswald, E. Niklasson, M. Schlegel, S. Kobayashi, N. Zucchet, N. Scherrer, N. Miller, M. Sandler, B. A. y Arcas, M. Vladymyrov, R. Pascanu, and J. Sacramento. Uncovering mesa-optimization algorithms in Transformers. CoRR, abs/2309.05858, 2023. URL https://doi.org/10.48550/arXiv.2309.05858. tex.cdate: 1672531200000 tex.publtype: informal.

[63] T. Poggio. On optimal nonlinear associative recall. Biological Cybernetics, 19(4):201–209, Sept. 1975. ISSN 1432-0770. doi:10.1007/BF02281970. URL https://doi.org/10.1007/BF02281970.

[64] K. Greff, R. K. Srivastava, J. Koutník, B. R. Steunebrink, and J. Schmidhuber. LSTM: A Search Space Odyssey. IEEE Transactions on Neural Networks and Learning Systems, 28(10):2222–2232, Oct. 2017. ISSN 2162-237X, 2162-2388. doi:10.1109/TNNLS.2016.2582924. URL http://arxiv.org/abs/1503.04069. arXiv:1503.04069 [cs].

[65] J. van der Westhuizen and J. Lasenby. The unreasonable effectiveness of the forget gate, Sept. 2018. URL http://arxiv.org/abs/1804.04849. arXiv:1804.04849 [cs, stat].

[66] R. M. Johnstone, C. R. Johnson, R. R. Bitmead, and B. D. O. Anderson. Exponential convergence of recursive least squares with exponential forgetting factor. In 1982 21st IEEE Conference on Decision and Control, pages 994–997, Dec. 1982. doi:10.1109/CDC.1982.268295. URL https://ieeexplore.ieee.org/document/4047398.

[67] S. Boyd and L. Vandenberghe. Convex optimization. Cambridge University Press, Cambridge, 2004.

[68] S. Yang, B. Wang, Y. Zhang, Y. Shen, and Y. Kim. Parallelizing Linear Transformers with the Delta Rule over Sequence Length, Aug. 2024c. URL http://arxiv.org/abs/2406.06484. arXiv:2406.06484.

[69] B. Widrow and M. E. Hoff. Adaptive switching circuits. In Neurocomputing: foundations of research, pages 123–134. MIT Press, Cambridge, MA, USA, Jan. 1988. ISBN 978-0-262-01097-9.

[70] J. Siems, T. Carstensen, A. Zela, F. Hutter, M. Pontil, and R. Grazzi. DeltaProduct: Increasing the Expressivity of DeltaNet Through Products of Householders, Feb. 2025. URL http://arxiv.org/abs/2502.10297. arXiv:2502.10297 [cs].

[71] N. Cesa-Bianchi, A. Conconi, and C. Gentile. On the generalization ability of on-line learning algorithms. IEEE Transactions on Information Theory, 50(9):2050–2057, Sept. 2004. ISSN 1557-9654. doi:10.1109/TIT.2004.833339. URL https://ieeexplore.ieee.org/document/1327806/?arnumber=1327806. Conference Name: IEEE Transactions on Information Theory.

[72] M. Zinkevich. Online convex programming and generalized infinitesimal gradient ascent. In Proceedings of the Twentieth International Conference on International Conference on Machine Learning, ICML'03, pages 928–935, Washington, DC, USA, Aug. 2003. AAAI Press. ISBN 978-1-57735-189-4.

[73] J. Duchi, E. Hazan, and Y. Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(61):2121–2159, 2011. URL http://jmlr.org/papers/v12/duchi11a.html.

[74] D. P. Kingma and J. Ba. Adam: A Method for Stochastic Optimization. 2015. URL https://arxiv.org/abs/1412.6980.

[75] F. T. Castoldi and M. L. R. de Campos. Minimum-disturbance description for the development of adaptation algorithms and a new leakage least squares algorithm. In 2009 IEEE International Conference on Acoustics, Speech and Signal Processing, pages 3129–3132, Apr. 2009. doi:10.1109/ICASSP.2009.4960287. URL https://ieeexplore.ieee.org/document/4960287. ISSN: 2379-190X.

[76] K. Tanabe. Projection method for solving a singular system of linear equations and its applications. Numerische Mathematik, 17(3):203–214, June 1971. ISSN 0945-3245. doi:10.1007/BF01436376. URL https://doi.org/10.1007/BF01436376.

[77] S. Kaczmarz. Angenaherte auflosung von systemen linearer glei-chungen. Bull. Int. Acad. Pol. Sic. Let., Cl. Sci. Math. Nat., pages 355–357, 1937. URL https://cir.nii.ac.jp/crid/1570009749812709888.

[78] I. Loshchilov and F. Hutter. Decoupled weight decay regularization. In International conference on learning representations, 2019. URL https://openreview.net/forum?id=Bkg6RiCqY7.

[79] Y. Chen, Q. Zeng, H. Ji, and Y. Yang. Skyformer: remodel self-attention with Gaussian kernel and nyström method. In Proceedings of the 35th International Conference on Neural Information Processing Systems, NIPS '21, pages 2122–2135, Red Hook, NY, USA, 2021. Curran Associates Inc. ISBN 978-1-7138-4539-3.

[80] J. Lu, J. Yao, J. Zhang, X. Zhu, H. Xu, W. Gao, C. XU, T. Xiang, and L. Zhang. SOFT: Softmax-free Transformer with Linear Complexity. In Advances in Neural Information Processing Systems, volume 34, pages 21297–21309. Curran Associates, Inc., 2021. URL https://proceedings.neurips.cc/paper/2021/hash/b1d10e7bafa4421218a51b1e1f1b0ba2-Abstract.html.

[81] L. Wasserman. All of Nonparametric Statistics. Springer Texts in Statistics. Springer, New York, NY, 2006. ISBN 978-0-387-25145-5. doi:10.1007/0-387-30623-4. URL http://link.springer.com/10.1007/0-387-30623-4.

[82] J. Fan. Local Polynomial Modelling and Its Applications: Monographs on Statistics and Applied Probability 66. Routledge, New York, May 2018. ISBN 978-0-203-74872-5. doi:10.1201/9780203748725.

[83] Y. Zhang, B. Liu, Q. Cai, L. Wang, and Z. Wang. An Analysis of Attention via the Lens of Exchangeability and Latent Variable Models, Apr. 2024. URL http://arxiv.org/abs/2212.14852. arXiv:2212.14852 [cs].

[84] G. Goel and P. Bartlett. Can a transformer represent a Kalman filter? In A. Abate, M. Cannon, K. Margellos, and A. Papachristodoulou, editors, Proceedings of the 6th annual learning for dynamics & control conference, volume 242 of Proceedings of machine learning research, pages 1502–1512. PMLR, July 2024. URL https://proceedings.mlr.press/v242/goel24a.html.

[85] K. A. Wang, G. Pleiss, J. Gardner, S. Tyree, K. Q. Weinberger, and A. G. Wilson. Exact Gaussian Processes on a Million Data Points. Advances in Neural Information Processing Systems, 32, 2019. URL https://papers.nips.cc/paper/2019/hash/01ce84968c6969bdd5d51c5eeaa3946a-Abstract.html.

[86] T. Dao, D. Y. Fu, S. Ermon, A. Rudra, and C. Re. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. In Advances in Neural Information Processing Systems, Oct. 2022. URL https://openreview.net/forum?id=H4DqfPSibmx.

[87] D. Y. Fu, T. Dao, K. K. Saab, A. W. Thomas, A. Rudra, and C. Re. Hungry Hungry Hippos: Towards Language Modeling with State Space Models. In The Eleventh International Conference on Learning Representations, Sept. 2022. URL https://openreview.net/forum?id=COZDy0WYGg.

[88] M. Xu, W. Cheng, B. Wang, and W. Chen. KV Shifting Attention Enhances Language Modeling, Dec. 2024. URL http://arxiv.org/abs/2411.19574. arXiv:2411.19574 [cs].

[89] J. J. Hopfield. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the National Academy of Sciences, 79(8):2554–2558, Apr. 1982. doi:10.1073/pnas.79.8.2554. URL https://www.pnas.org/doi/abs/10.1073/pnas.79.8.2554. Publisher: Proceedings of the National Academy of Sciences.

[90] D. J. Willshaw, O. P. Buneman, and H. C. Longuet-Higgins. Non-Holographic Associative Memory. Nature, 222(5197):960–962, June 1969. ISSN 1476-4687. doi:10.1038/222960a0. URL https://www.nature.com/articles/222960a0. Publisher: Nature Publishing Group.

[91] D. Willshaw. Holography, Associative Memory, and Inductive Generalization. In Parallel Models of Associative Memory. Psychology Press, 1989. ISBN 978-1-315-80799-7. Num Pages: 22.

[92] A. Bietti, V. Cabannes, D. Bouchacourt, H. Jegou, and L. Bottou. Birth of a Transformer: A Memory Viewpoint. In Thirty-seventh Conference on Neural Information Processing Systems, Nov. 2023. URL https://openreview.net/forum?id=3X2EbBLNsk.

[93] D. Krotov and J. J. Hopfield. Dense Associative Memory for Pattern Recognition. In Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016. URL https://proceedings.neurips.cc/paper_files/paper/2016/hash/eaae339c4d89fc102edd9dbdb6a28915-Abstract.html.

[94] D. Krotov and J. Hopfield. Large Associative Memory Problem in Neurobiology and Machine Learning, Apr. 2021. URL http://arxiv.org/abs/2008.06996. arXiv:2008.06996.

[95] H. Ramsauer, B. Schäfl, J. Lehner, P. Seidl, M. Widrich, L. Gruber, M. Holzleitner, T. Adler, D. Kreil, M. K. Kopp, G. Klambauer, J. Brandstetter, and S. Hochreiter. Hopfield Networks is All You Need. In International Conference on Learning Representations, 2021. URL https://openreview.net/forum?id=tL89RnzIiCd.

[96] B. Millidge, T. Salvatori, Y. Song, T. Lukasiewicz, and R. Bogacz. Universal Hopfield Networks: A General Framework for Single-Shot Associative Memory Models. In Proceedings of the 39th International Conference on Machine Learning, pages 15561–15583. PMLR, June 2022. URL https://proceedings.mlr.press/v162/millidge22a.html. ISSN: 2640-3498.

[97] S. Sukhbaatar, E. Grave, G. Lample, H. Jegou, and A. Joulin. Augmenting Self-attention with Persistent Memory, July 2019. URL http://arxiv.org/abs/1907.01470. arXiv:1907.01470.

[98] J. Zhang, N. Nolte, R. Sadhukhan, B. Chen, and L. Bottou. Memory mosaics. In The thirteenth international conference on learning representations, 2025. URL https://openreview.net/forum?id=IiagjrJNwF.

[99] E. Nichani, J. D. Lee, and A. Bietti. Understanding Factual Recall in Transformers via Associative Memories, Dec. 2024. URL http://arxiv.org/abs/2412.06538. arXiv:2412.06538 [cs].

[100] V. Cabannes, E. Dohmatob, and A. Bietti. Scaling Laws for Associative Memories. In The Twelfth International Conference on Learning Representations, Oct. 2023. URL https://openreview.net/forum?id=Tzh6xAJSll.

[101] Q. Dong, L. Li, D. Dai, C. Zheng, J. Ma, R. Li, H. Xia, J. Xu, Z. Wu, T. Liu, B. Chang, X. Sun, L. Li, and Z. Sui. A Survey on In-context Learning, Oct. 2024. URL http://arxiv.org/abs/2301.00234. arXiv:2301.00234 [cs].

[102] S. Garg, D. Tsipras, P. Liang, and G. Valiant. What can transformers learn in-context? A case study of simple function classes. In A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, editors, Advances in neural information processing systems, 2022. URL https://openreview.net/forum?id=flNZJ2eOet.

[103] Y. Zhang, F. Zhang, Z. Yang, and Z. Wang. What and How does In-Context Learning Learn? Bayesian Model Averaging, Parameterization, and Generalization, Oct. 2023b. URL http://arxiv.org/abs/2305.19420. arXiv:2305.19420 [stat].

[104] J. V. Oswald, E. Niklasson, E. Randazzo, J. Sacramento, A. Mordvintsev, A. Zhmoginov, and M. Vladymyrov. Transformers Learn In-Context by Gradient Descent. In Proceedings of the 40th International Conference on Machine Learning, pages 35151–35174. PMLR, July 2023. URL https://proceedings.mlr.press/v202/von-oswald23a.html. ISSN: 2640-3498.

[105] K. Ahn, X. Cheng, H. Daneshmand, and S. Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36:45614–45650, Dec. 2023. URL https://proceedings.neurips.cc/paper_files/paper/2023/hash/8ed3d610ea4b68e7afb30ea7d01422c6-Abstract-Conference.html.

[106] J. Shi, K. A. Wang, and E. B. Fox. Sequence modeling with multiresolution convolutional memory. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of ICML'23, pages 31312–31327, Honolulu, Hawaii, USA, July 2023. JMLR.org.

[107] M. Poli, S. Massaroli, E. Nguyen, D. Y. Fu, T. Dao, S. Baccus, Y. Bengio, S. Ermon, and C. Ré. Hyena hierarchy: towards larger convolutional language models. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of ICML'23, pages 28043–28078, Honolulu, Hawaii, USA, July 2023. JMLR.org.

[108] J. Gehring, M. Auli, D. Grangier, D. Yarats, and Y. N. Dauphin. Convolutional Sequence to Sequence Learning. In Proceedings of the 34th International Conference on Machine Learning, pages 1243–1252. PMLR, July 2017. URL https://proceedings.mlr.press/v70/gehring17a.html. ISSN: 2640-3498.

[109] J. Bradbury, S. Merity, C. Xiong, and R. Socher. Quasi-recurrent neural networks. In International conference on learning representations, 2017. URL https://openreview.net/forum?id=H1zJ-v5xl.

[110] J. T. Smith, A. Warrington, and S. Linderman. Simplified state space layers for sequence modeling. In The eleventh international conference on learning representations, 2023. URL https://openreview.net/forum?id=Ai8Hw3AXqks.

[111] Z. Qin, X. Han, W. Sun, B. He, D. Li, D. Li, Y. Dai, L. Kong, and Y. Zhong. Toeplitz neural network for sequence modeling. In The eleventh international conference on learning representations, 2023. URL https://openreview.net/forum?id=IxmWsm4xrua.

[112] L. Ren, Y. Liu, S. Wang, Y. Xu, C. Zhu, and C. Zhai. Sparse modular activation for efficient sequence modeling. In Thirty-seventh conference on neural information processing systems, 2023. URL https://openreview.net/forum?id=TfbzX6I14i.

[113] H. Mehta, A. Gupta, A. Cutkosky, and B. Neyshabur. Long range language modeling via gated state spaces. In The eleventh international conference on learning representations, 2023. URL https://openreview.net/forum?id=5MkYIYCbva.

[114] W. Hua, Z. Dai, H. Liu, and Q. Le. Transformer Quality in Linear Time. In Proceedings of the 39th International Conference on Machine Learning, pages 9099–9117. PMLR, June 2022. URL https://proceedings.mlr.press/v162/hua22a.html. ISSN: 2640-3498.

[115] R. Sutton. The Bitter Lesson, 2019. URL http://www.incompleteideas.net/IncIdeas/BitterLesson.html.

[116] J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. d. L. Casas, L. A. Hendricks, J. Welbl, A. Clark, T. Hennigan, E. Noland, K. Millican, G. v. d. Driessche, B. Damoc, A. Guy, S. Osindero, K. Simonyan, E. Elsen, J. W. Rae, O. Vinyals, and L. Sifre. Training Compute-Optimal Large Language Models, Mar. 2022. URL http://arxiv.org/abs/2203.15556. arXiv:2203.15556 [cs].

[117] J. Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, S. Gray, A. Radford, J. Wu, and D. Amodei. Scaling Laws for Neural Language Models, Jan. 2020. URL http://arxiv.org/abs/2001.08361. arXiv:2001.08361 [cs].

[118] Y. Sun, X. Wang, Z. Liu, J. Miller, A. Efros, and M. Hardt. Test-Time Training with Self-Supervision for Generalization under Distribution Shifts. In Proceedings of the 37th International Conference on Machine Learning, pages 9229–9248. PMLR, Nov. 2020. URL https://proceedings.mlr.press/v119/sun20b.html. ISSN: 2640-3498.

[119] E. Akyürek, M. Damani, L. Qiu, H. Guo, Y. Kim, and J. Andreas. The Surprising Effectiveness of Test-Time Training for Abstract Reasoning, Nov. 2024. URL http://arxiv.org/abs/2411.07279. arXiv:2411.07279 [cs] version: 1.

[120] OpenAI. OpenAI o1 System Card, 2024. URL https://openai.com/index/openai-o1-system-card/.