ArXiv: 2309.01775
🎯 Pitch
Gated RNNs with multiplicative interactions can exactly implement linear self-attention using a finite number of neurons—and gradient descent naturally discovers this construction, revealing that RNNs and Transformers may not be as architecturally distinct as previously thought.
1. Executive Summary
This paper studies how gated recurrent neural networks (RNNs) with multiplicative interactions can exactly implement linear self-attention, and demonstrates through reverse-engineering experiments on synthetic in-context learning tasks that gradient descent discovers this construction in practice. The core contribution is a constructive proof that a simplified class of diagonal linear gated RNNs — featuring input gating (elementwise multiplication of linearly transformed input branches) and output gating (elementwise multiplication of recurrent state representations) — can match any linear self-attention layer using a finite number of neurons, requiring O(d²) hidden units and O(d⁴) total parameters for a d-dimensional input. The paper establishes that trained gated RNNs discover compressed versions of this construction — using only d(d+1)/2 neurons to store symmetric key-value matrices when value matrices are invertible, or scaling with the rank of the attention matrices rather than the full dimension — and that on in-context linear regression tasks, these RNNs learn to implement one step of gradient descent with optimal learning rate, the very same algorithm that linear self-attention layers typically discover, revealing that certain RNNs may be unexpectedly implementing attention under the hood only when both multiplicative gating and perfect-memory recurrence (λ = 1) are architecturally enabled.
2. Context and Motivation
The Core Question: Are RNNs and Attention Fundamentally Different Computational Mechanisms?
This paper tackles a question that sits at the heart of modern deep learning architecture design: can recurrent neural networks implement the same computations as attention-based models, and if so, does gradient descent naturally discover such implementations? This is not merely an academic curiosity — it cuts to whether the architectural dichotomy between RNNs and Transformers reflects genuine computational differences or is an artifact of our training procedures and inductive biases.
The practical stakes are high. Transformers dominate sequence modeling because softmax attention provides a powerful mechanism for dynamically routing information across time steps — each position can directly attend to any previous position through learned similarity scores. However, this power comes at a cost: the memory and computation required for standard softmax attention scale quadratically with sequence length, making Transformers expensive or intractable for very long sequences. RNNs, by contrast, process sequences with constant memory cost per step — their recurrent state serves as a fixed-size summary of all past information. This efficiency is the reason for the recent resurgence of interest in deep linear RNN architectures (Gu et al., 2022; Orvieto et al., 2023b; Gu & Dao, 2023), which have demonstrated surprising competitiveness with Transformers on long-sequence benchmarks while maintaining the inference efficiency that made RNNs attractive in the first place.
But this efficiency comes with a nagging theoretical concern: the constant-size state of an RNN imposes a bottleneck that attention mechanisms bypass entirely. If RNNs are fundamentally less capable than attention at certain kinds of sequence processing — particularly tasks requiring selective retrieval of past information or dynamic recombination of stored elements — then the efficiency gains come at a performance ceiling that no amount of scaling can overcome. Conversely, if RNNs can implement attention-like computations, then the efficiency-performance tradeoff might be resolvable through better architectural design rather than being an inherent limitation of recurrence.
The Gap: No Constructive Bridge Between RNNs and Attention
Prior to this paper, the relationship between RNNs and attention was understood primarily through two incomplete lenses, neither of which provided a satisfying answer to the capability question.
Universality arguments are too weak. The literature on the approximation capabilities of recurrent networks establishes that linear RNNs with nonlinear readouts are universal function approximators in the limit of infinite hidden state dimension (Boyd & Chua, 1985; Grigoryeva & Ortega, 2018; Orvieto et al., 2023a). From this, one can trivially conclude that RNNs could approximate attention — given enough neurons. But universality theorems tell us nothing about whether this approximation is practical. Does it require an astronomically large state? Does gradient descent reliably find the right parameters? Is there a finite construction that implements attention exactly, not just approximately? Universality theorems are existence proofs, not constructive recipes, and they leave entirely open the question that matters for practitioners: if I train an RNN on a task where attention would perform well, will the RNN actually learn an attention-like computation?
Linearized attention and fast weight memories provide conceptual connections, but no constructive equivalence. A parallel line of work has explored linearized attention mechanisms, where the softmax in standard attention is removed or approximated, enabling a reformulation where the forward pass can be executed in an RNN-like fashion (Katharopoulos et al., 2020; Choromanski et al., 2021; Schlag et al., 2021; Peng et al., 2021). In this view, a linear self-attention layer maintains an effective weight matrix that accumulates outer products of keys and values, and processes the current query through . This is structurally reminiscent of fast weight memories (Schmidhuber, 1992; Ba et al., 2016), where a "slow" neural network modifies "fast" associative weights through Hebbian-like updates. The connection is tantalizing: linear attention looks like an RNN with a matrix-valued state and a particular update rule.
But the existing literature treats linear attention as an alternative architectural choice — a Transformer variant, not an RNN. The question of whether standard RNN architectures, with their vector-valued states and gating mechanisms, can recover this matrix-valued computation is entirely different. Linear attention stores a matrix in its state, requiring numbers. A standard RNN with a hidden state of size stores only numbers. Can be finite (and reasonably small) while still exactly implementing the attention computation? This is a constructive question about architectural capacity, not an asymptotic one about universality.
The role of gating is empirically observed but theoretically unexplained. Modern high-performing RNNs — LSTMs, GRUs, and the recent deep linear RNN architectures like LRU (Orvieto et al., 2023b), S4 (Gu et al., 2022), and Mamba (Gu & Dao, 2023) — all prominently feature multiplicative interactions, typically in the form of elementwise gating. For instance, the Gated Linear Unit (GLU; Dauphin et al., 2017) computes , and appears as a standard component after recurrent layers in many architectures. The Mamba architecture uses input-dependent matrices that, when multiplied with the input, produce a multiplicative gating effect. These gating mechanisms are empirically important — removing them degrades performance — but their theoretical function has been unclear. Are they simply providing nonlinearities that improve optimization? Or do they serve a more specific computational role? The paper hypothesizes a precise answer: elementwise multiplication is the fundamental operation needed to compute outer products (and thus key-value accumulations) and to multiply accumulated matrices with current queries — the two core operations of linear attention.
Conflicting Narratives in the Literature
The paper enters a landscape where two partially contradictory narratives coexist, neither of which captures the full picture:
Narrative 1: RNNs are obsolete for sequence modeling. The empirical dominance of Transformers since Vaswani et al. (2017) has created a widespread perception that attention mechanisms are simply superior to recurrence for sequence processing tasks. This narrative is reinforced by the practical success of large language models, all of which use attention, and by the theoretical argument that attention provides direct, distance-independent access to all past tokens, whereas RNNs must route information through a bottlenecked state. From this perspective, recurrent architectures are a historical artifact, retained only for specialized applications where inference efficiency is paramount and performance is secondary.
Narrative 2: Modern RNNs are catching up, but the mechanism is mysterious. The counter-narrative comes from the recent empirical successes of deep linear RNNs. These models have matched or exceeded Transformers on challenging long-sequence benchmarks like the Long Range Arena (Tay et al., 2020), and in some cases on language modeling tasks (Gu & Dao, 2023). However, the explanation for why these RNNs work well has been unsatisfying. Some attribute it to better optimization properties of diagonal linear recurrence (Orvieto et al., 2023b; Zucchet et al., 2023b). Others point to the complex-valued parameterization enabling long-range memory. The possibility that these RNNs might be implementing attention internally has not been seriously explored, and the community lacks a mechanistic understanding of what computations the trained networks are actually performing.
Narrative 3 (the paper's contribution): RNNs can implement attention, and gradient descent finds this solution. The paper's thesis reconciles these narratives: RNNs with the right architectural elements (specifically, multiplicative input and output gating) have the capacity to exactly implement linear self-attention, and gradient descent discovers this capacity when the task demands it. This means that the performance of modern RNNs on tasks where attention excels may not be evidence that attention is unnecessary — it may be evidence that the RNN is secretly doing attention internally, hidden in its vector-valued state. The distinction matters profoundly: if RNNs succeed by implementing attention, then their performance ceiling is bounded by the quality of the attention they can implement, which depends on architectural details (state size, gating design). Alternatively, if RNNs succeed through a fundamentally different mechanism, they might have complementary strengths that attention lacks. The paper's findings lean toward the former interpretation, with important caveats about the parameter efficiency of the implementation.
Why This Matters: Bridging Two Research Communities
The significance of establishing a constructive equivalence between gated RNNs and linear attention extends beyond theoretical curiosity into several practical domains:
For architecture design, knowing that gated RNNs can implement attention provides a principled framework for guiding development. If a new RNN variant is being proposed, one can ask: does it contain the necessary multiplicative interactions to compute outer products and matrix-vector products? If not, it may lack the capacity to implement attention-like routing of information, which could explain performance deficits on tasks requiring selective memory access. The paper explicitly demonstrates this design principle by modifying the LRU architecture — adding an input GLU to complement the existing output GLU significantly improves its ability to mimic attention (Section 4.3, Figure 4.A), and replacing the GLU with an MLP of equal parameter count destroys this ability. These are not obvious architectural choices unless one understands the computational role of gating.
For mechanistic interpretability, the paper provides a concrete hypothesis about what trained RNNs might be computing. Rather than treating the hidden state as an opaque vector, the construction suggests that subsets of neurons may have interpretable roles: some neurons (with ) accumulate outer products of specific input features, while others (with ) transiently hold current query representations, and the output gating multiplies these representations together. This decomposition into key-value accumulators and query carriers is a testable prediction about the structure of trained weights, and the paper validates it through reverse-engineering experiments (Section 4.1, Figure 2).
For the training-inference efficiency tradeoff, understanding the parameter cost of implementing attention in RNNs is crucial. The paper's construction requires neurons and parameters — substantially more than the parameters of the linear self-attention layer itself. This immediately explains why RNNs are not a drop-in replacement for attention at scale: the cost is prohibitive. However, the paper also shows that trained networks discover compressed implementations when the attention matrices have low-rank structure (Section 4.2, Figure 3.B), reducing the required state size from to . This suggests that the practical viability of RNNs as attention-implementers depends crucially on whether real-world attention patterns admit such compression, and it opens the door to RNN designs that explicitly encode low-rank inductive biases.
For neuroscience, the connection between gated recurrence and attention has provocative implications. If multiplicative interactions in recurrent circuits can implement Hebbian-like fast weight updates (storing key-value associations in neural activity rather than in synaptic plasticity), this suggests a mechanism by which biological neural networks might perform attention-like computations. The paper notes that several single-neuron and circuit-level mechanisms have been experimentally identified that could support the required multiplication operations (Silver, 2010), and speculates that recurrent circuits with long integration time constants — such as those found in prefrontal cortex — might be encoding associative weights in their neural activity patterns.
Prior Work and Its Limitations
The paper positions itself against several specific lines of prior work:
Linear Transformers as fast weight programmers. Schlag et al. (2021) explicitly framed linear Transformers as "fast weight programmers" — systems where a neural network updates a separate set of weights through outer product accumulation. This connection is conceptually important but it works in the opposite direction from this paper's contribution: it shows that attention can be viewed as recurrence, not that RNNs can implement attention. The paper builds on this conceptual bridge but provides the constructive reverse direction, showing exactly what architectural elements RNNs need to replicate the fast weight programming computation.
Universality theorems for linear RNNs. Orvieto et al. (2023a) established that linear RNNs followed by nonlinear readouts are universal approximators. While this guarantees that RNNs can in principle approximate attention, it requires the hidden state dimension to grow to infinity for exact representation and provides no guidance on what finite state size is needed. The paper's contribution is stronger: it provides an explicit finite-dimensional construction (with state size in the symmetric case) and then demonstrates that gradient descent finds this construction or compressed variants of it.
In-context learning algorithms in Transformers. Von Oswald et al. (2023) showed that linear self-attention layers trained on in-context linear regression converge to implementing one step of gradient descent with a particular structure. Mahankali et al. (2023) proved that this is the globally optimal solution for single-layer linear attention on this task. The paper builds on this work by showing that gated RNNs trained on the same task discover the same algorithm — not an approximation, but the functionally identical computation (Section 5.2, Table 2). This is a strong mechanistic alignment result: two very different architectures converge to the same computation when trained on the same problem, suggesting the algorithm is determined by the task structure, not the architecture.
Architecture-specific capabilities. The paper provides a somewhat surprising analysis of how these results apply to specific RNN architectures. LSTMs can implement linear self-attention in a single layer (due to the structural similarity between the LSTM's input gate + cell state mechanism and outer product accumulation, and between its output gate and query multiplication), but GRUs cannot — even with stacking — because they lack a mechanism to compute the multiplicative interaction between keys and values that produces the outer product. However, when actually trained, LSTMs fail to discover this implementation, performing similarly to GRUs (Section 4.3). This disconnect between theoretical capacity and empirical convergence to that capacity is a crucial finding: architectural inductive bias matters not just for what is theoretically possible, but for what gradient descent actually finds. The paper shows that the simpler, more constrained architectures (diagonal linear recurrence with explicit linear gating) have the right inductive bias, while more complex architectures (LSTMs with their sigmoid and tanh nonlinearities) obscure the path to the attention solution.
How This Paper Positions Itself
The paper's core methodological move is to combine constructive theory with mechanistic reverse-engineering. Rather than establishing a loose conceptual similarity or an asymptotic equivalence, the paper provides explicit weight matrices (spelled out in Appendix A.1) that make a gated RNN compute the exact same function as any given linear self-attention layer. This construction is then used as a template for reverse-engineering trained networks: the authors look for networks where the learned weights closely match the structural pattern of the construction (block-sparse weight matrices, specific groupings of and neurons, interpretable roles for input gating, output gating, and readout weights).
By doing this, the paper shifts the question from "Can RNNs approximate attention?" (answer: yes, trivially, by universality) to "Do RNNs actually learn to implement attention when trained on tasks that benefit from it, and if so, how efficiently do they do so?" The answer to this more precise question reveals both the promise and the limitations: yes, gated RNNs do learn attention, often discovering clever compressions (using symmetry, exploiting low-rank structure), but the parameter cost is substantial — state neurons for a computation that attention implements with parameters total. This positions gated RNNs not as a replacement for attention, but as architectures that can recapitulate attention-like computations when needed, with efficiency characteristics that depend heavily on the structure of the attention patterns being implemented.
3. Technical Approach
3.1 Reader Orientation
This paper is primarily a constructive theoretical analysis combined with mechanistic reverse-engineering experiments — it provides explicit mathematical proofs that gated RNNs can exactly implement linear self-attention, then empirically examines whether gradient descent discovers these implementations when RNNs are trained on tasks where attention is the natural solution. The "system" being studied is not a novel practical method but rather a framework for understanding the computational relationship between two dominant neural architectures. The paper solves the problem of "can recurrent networks do what attention does?" not through asymptotic approximation theory but through finite constructive equivalence: it shows exactly which architectural elements are necessary (multiplicative gating for computing outer products, perfect-memory recurrence for accumulation, and output gating for query multiplication), exactly how many neurons are required (O(d²) in the general case, reducible via symmetry and low-rank structure), and exactly what pattern of weights implements the computation, then verifies that trained networks discover this pattern.
3.2 Big-Picture Architecture (Diagram in Words)
The paper's technical framework has four interconnected components that form a logical progression from theory to empirical validation:
1. The Teacher Architecture: Linear Self-Attention (Section 2.1) — A causally-masked linear attention layer that processes sequences by maintaining a matrix-valued state (accumulating outer products of value and key vectors) and using this matrix to transform the current query: . This is the computation that the RNN must learn to replicate.
2. The Student Architecture: Gated Diagonal Linear RNN (Section 2.2) — A simplified recurrent network where the hidden state evolves as , with implementing multiplicative input gating as , and the output is produced through output gating followed by a linear readout: where . Each neuron has a decay factor that determines whether it accumulates () or holds only the current value ().
3. The Constructive Proof (Section 3) — An explicit assignment of weight matrices that makes the student RNN exactly equal any given teacher linear self-attention layer. The construction maps: input gating computes outer products of keys and values (storing each of the d² entries of the key-value matrix) and passes through current queries; recurrent neurons with accumulate these outer products over time; recurrent neurons with hold only the current query; output gating multiplies each accumulated key-value entry by the corresponding query component; and the readout matrix sums these products to produce the final output. The construction requires hidden neurons in its basic form and total parameters.
4. The Reverse-Engineering Pipeline (Section 4) — After training student RNNs to match a teacher attention layer or to solve in-context learning tasks, the authors analyze the learned weights to determine whether they implement the construction. This involves: (a) identifying which neurons have versus , (b) pruning neurons whose input or output weights are all zeros, (c) verifying that key-values can be linearly reconstructed from the neurons and queries from the neurons, (d) comparing the polynomial coefficients of the RNN's instantaneous input-output mapping to those of the teacher attention layer, and (e) checking whether the weight matrices exhibit the block-sparse structure predicted by the construction.
3.3 Roadmap for the Deep Dive
The following detailed breakdown proceeds in the order that mirrors how the paper builds its argument:
- First, the simplified gated RNN model (Equation 2) — its precise mathematical definition, the role of each component (diagonal recurrence, input gating, output gating, readout), and how this simplified class relates to practical architectures like LRUs, LSTMs, GRUs, and Mamba. Understanding this model is essential because all constructive proofs and reverse-engineering analyses operate on it.
- Second, the constructive proof (Section 3.1, Appendix A.1) — the explicit mapping from any linear self-attention layer to an equivalent gated RNN. This covers: how input gating computes outer products and passes through queries, how values determine memory behavior, how output gating performs the matrix-vector multiplication, and how the readout sums the results. This is the theoretical heart of the paper.
- Third, the parameter count and compression strategies (Section 3.2, Appendices A.3-A.4) — why the naive construction requires neurons and parameters, how symmetry of the key-value matrix reduces this to neurons when value matrices are invertible, and how low-rank structure in the attention matrices enables further compression to neurons. This connects the theory to practical questions about efficiency.
- Fourth, the alternative construction with side gating (Appendix A.2) — a variant where output gating is replaced by side gating , which reduces the parameter cost to and provides a stronger inductive bias. This variant connects to the Mamba architecture.
- Fifth, the implications for existing RNN classes (Section 3.3, Appendix B) — a detailed comparison showing which practical architectures can implement the construction (LRU with stacking: yes; Mamba: yes in one layer; LSTM: yes in one layer theoretically; GRU: no, even with stacking). This bridges the simplified theory to real-world models.
- Sixth, the teacher-student experimental methodology (Section 4) — how the paper trains student RNNs to emulate teacher attention layers, the loss function, the optimization setup, and the analytical tools used to determine whether the student has learned the construction (linear reconstructability scores, polynomial coefficient comparison, weight matrix structural analysis).
- Seventh, the in-context learning experimental methodology (Section 5) — the linear regression task formulation, the connection to one-step gradient descent, and how the paper determines that trained RNNs implement the same algorithm as linear self-attention on this task.
3.4 Detailed, Sentence-Based Technical Breakdown
This is primarily a constructive theory and mechanistic analysis paper, whose core claim is that gated RNNs with multiplicative interactions can implement linear self-attention exactly, and that gradient descent in practice discovers this implementation under the right architectural conditions.
The Simplified Gated RNN Model
Why this model matters. The paper analyzes a simplified class of gated diagonal linear recurrent neural networks because it isolates the minimal computational mechanisms needed to establish the connection to attention while remaining closely related to practical architectures. This model is not proposed as a new architecture to be deployed; rather, it serves as a tractable analytical object — simple enough to provide explicit constructions, yet containing the essential elements (multiplicative gating, diagonal recurrence, linear readout) that appear in various forms across modern RNN designs. The simplifications (linear gating instead of sigmoid-gated, diagonal instead of dense recurrence, real-valued instead of complex-valued) are deliberately chosen to make the mathematical analysis transparent while preserving the computational phenomena of interest.
Recurrence with elementwise decay. The core recurrence equation defines the temporal dynamics:
where is the hidden state at time , is a vector of per-neuron decay factors, is the input vector (augmented with a constant 1 to enable bias-like terms), and is the input gating function.
What this equation computes: Each coordinate of the hidden state independently mixes its previous value (scaled by ) with a new contribution . When , the neuron perfectly integrates all past inputs — it acts as an accumulator with infinite memory. When , the neuron completely discards past information at each step, holding only the current input. When , the neuron implements an exponential moving average with decay rate , where older inputs are gradually forgotten.
Why this form: Diagonal linear recurrence is chosen for three reasons. First, it matches the recurrence structure used in modern deep linear RNN architectures (LRU, S4, S5), making the analysis directly relevant to practical models. Second, diagonal recurrence is analytically tractable — each coordinate evolves independently, so the behavior of the entire state can be understood by analyzing individual neurons. Third, diagonal recurrence has been shown to facilitate gradient-based learning compared to dense recurrence (Orvieto et al., 2023b; Zucchet et al., 2023b), meaning that the solutions the paper identifies analytically are also likely to be reachable by gradient descent. The elementwise product means that each neuron has its own independent forgetting rate, allowing different neurons to operate at different timescales simultaneously — this heterogeneity is crucial because the construction requires both perfect-memory neurons () and instantaneous neurons () in the same layer.
Input gating as bilinear computation. The input gating function is defined as:
where and are learnable weight matrices, denotes elementwise (Hadamard) multiplication, and the input has been augmented with a constant 1 as its -th entry to allow the gating to produce terms of degree 1 (linear) when one branch uses only the constant entry.
What this equation computes: For each neuron , the two matrices and produce two scalars and , which are then multiplied elementwise to yield the neuron's input . Since and are each linear functions of the input, their product is a quadratic (degree-2) polynomial in — specifically, a sum of terms of the form . When the constant-1 dimension is involved, one of the terms becomes linear, producing a degree-1 polynomial.
Why this form: The elementwise multiplication of two linear transformations is precisely what enables the computation of outer products between different linear projections of the input. If and are configured so that one branch produces values and the other produces keys , then their product is exactly entry of the outer product . Without this multiplicative interaction, a purely linear recurrent network could only accumulate degree-1 polynomials of past inputs — it could never represent the key-value outer products that are the core computational primitive of attention. Alternative forms of nonlinearity (e.g., applying a sigmoid or ReLU after a linear transformation) would produce different polynomial structures that cannot cleanly decompose into outer products. The specific bilinear form is chosen because it directly generalizes the Gated Linear Unit (GLU; Dauphin et al., 2017), which is the standard gating mechanism in modern RNNs, and because its degree-2 polynomial structure matches the degree-2 structure of the key-value outer product in linear attention.
Output gating as matrix-vector multiplication. The output gating function operates on the hidden state:
where and are learnable weight matrices, and is the number of output gating neurons (potentially different from ).
What this equation computes: Each output gating neuron computes the product of two linear combinations of the hidden state: . If the hidden state contains accumulated key-value entries (in the neurons) and the current query (in the neurons), and if the rows of and are configured to select specific key-value entries and query components respectively, then each output gating neuron computes the product of one key-value accumulation with one query component — the elementary operation needed for the attention output.
Why this form: The output gating multiplies state-derived quantities together, which is essential because the hidden state mixes accumulated key-values (which depend on all past inputs) with the current query (which depends only on the current input). A purely linear readout could only compute linear combinations of key-values and queries, not products of them. The attention output requires multiplying the accumulated key-value matrix by the current query vector — this is fundamentally a bilinear operation in the stored matrix and the current vector. Output gating provides this bilinearity by multiplying two different linear projections of the state. The alternative of storing the product directly in the recurrent state would require the recurrent computation itself to be multiplicative, which would violate the linear recurrence structure needed for stable training and long-range memory.
Final linear readout. The network output is produced by a linear readout of the output-gated representation:
where is a learnable weight matrix.
What this equation computes: It sums the output gating neuron activations into the output dimensions using learned weights. In the attention-matching construction, this summation corresponds to the matrix-vector multiplication that combines all key-value-query products into the final output — each output gating neuron produces one term , and the readout sums over to produce .
Why this form: After the output gating has computed all the pairwise products between key-value entries and query components, the readout matrix provides the summation step that completes the matrix-vector multiplication. This decomposition — output gating for multiplication, readout for summation — is not strictly necessary (the gating could in principle incorporate the summation as well), but it makes the construction modular and interpretable. It also mirrors the structure of practical architectures where a GLU is followed by a linear projection.
Relationship to practical architectures (Appendix B). The simplified model relates to several real RNN designs as follows. The LRU (Orvieto et al., 2023b) uses complex-valued diagonal recurrence with an output GLU, matching our recurrence and output gating but lacking input gating; stacking two LRU layers allows the first layer's output gating to serve as the second layer's input gating, requiring 2 layers to implement attention (Table 3). The LRU In+Out variant (introduced in this paper) adds an explicit input GLU, enabling one-layer attention implementation. The Mamba layer (Gu & Dao, 2023) uses input-dependent matrices that, when multiplied with the input, create multiplicative gating in the recurrence input, and its side-gating mechanism matches our side-gating variant — enabling one-layer attention implementation (Table 3). LSTMs have the theoretical capacity to implement attention in one layer because their input gate + cell state update mechanism can compute outer products and their output gate can multiply queries, but this requires linearizing the tanh and sigmoid nonlinearities. GRUs cannot implement attention even with stacking because they lack a mechanism to compute multiplicative key-value interactions — the reset gate modulates the previous state but does not create products of different projections of the input.
The Constructive Proof: Mapping Attention to Gated RNNs
What the proof establishes. The core theoretical contribution (Section 3.1, detailed in Appendix A.1) is a constructive proof that for any linear self-attention layer with parameters , there exists a finite-dimensional gated RNN of the form described above that computes the identical function on all input sequences. This is fundamentally stronger than a universality result — it provides explicit numerical values for all weight matrices, it uses a finite (and characterized) number of neurons, and the equivalence is exact rather than approximate.
The teacher computation. The linear self-attention layer to be matched processes input sequences as:
where are the value, key, and query matrices respectively. Define , , and as the value, key, and query vectors at time . The matrix is the accumulated key-value outer product, and the output is .
What this equation computes: At each timestep , the layer takes all previous inputs (through their extracted values and keys) and the current query, and produces a weighted combination of past values where the weights are determined by the similarity (inner product) between the corresponding key and the current query. Each output dimension is — it sums over all past timesteps and all key dimensions the product of three terms: the value component at that past step, the key component at that past step, and the query component now. This is a degree-3 polynomial in the input variables (degree 2 from the key-value accumulation, times degree 1 from the query).
Why this attention form matters: Linear self-attention represents the core computational primitive that distinguishes Transformers from standard RNNs: the ability to dynamically route information between arbitrary positions in a sequence based on learned similarity functions. The softmax-nonlinearity of standard attention changes the optimization landscape and enables sharper selectivity, but the fundamental computational structure — accumulate outer products, multiply by query — is preserved in the linearized version.
The construction: input gating for outer products and queries. The construction uses hidden neurons. The input vector is augmented with a constant 1 to become . The weight matrices and have dimensions and are defined as follows (Equations 9-10 in Appendix A.1):
For (the key-value accumulator neurons), let and identify which entry of the key-value matrix this neuron stores. Then for :
with all entries for (the constant dimension) and for outside these ranges set to zero.
For (the query pass-through neurons), let identify which query component this neuron carries. Then for :
with for these neurons (all other entries zero).
What these matrices compute: For a key-value neuron (first neurons), the input gating produces:
This is exactly one entry of the outer product between the current value and key vectors. Each of the neurons stores a different entry of this matrix.
For a query neuron (last neurons), the input gating produces:
because reads the constant-1 dimension (giving 1) and reads the query projection. This produces the current query component without any multiplication.
Why this structure: The separation into key-value neurons and query neurons reflects the two different temporal behaviors needed: key-values must be accumulated over time (requiring ), while queries must be instantaneous (requiring ). Using the constant-1 dimension for the query neurons ensures that the input gating produces a linear function (degree 1) of the input for queries, while producing a quadratic function (degree 2) for key-values — this matches the degree structure of the attention computation where key-values are products of two input projections and queries are single projections.
The construction: recurrence with selective memory. The decay vector is defined as:
What this equation computes: The first neurons have , meaning their state at time is the sum of all inputs they have ever received:
This is exactly the entry of the accumulated key-value matrix . The last neurons have , meaning they hold only the current input:
This is exactly the -th component of the current query vector.
Why this binary choice: The construction requires exactly two distinct temporal behaviors — perfect accumulation and zero memory — because the attention computation separates information into a cumulative matrix (all past key-values) and an instantaneous vector (current query). Intermediate values (exponential decay) would produce a weighted sum that does not match the attention computation, which treats all past positions equally (no recency bias). The ability to set for some neurons and for others in the same layer is crucial — it means the RNN can simultaneously maintain both types of temporal processing, which is what enables the attention mapping in a single recurrent layer rather than requiring separate layers for accumulation and pass-through.
The construction: output gating for matrix-vector multiplication. The output gating has neurons. The weight matrices and have dimensions and are defined as follows (Equations 12-13 in Appendix A.1):
For output neuron (where for row and column ):
What these equations compute: is a selector matrix — its -th row has a single 1 at position , so it extracts from the hidden state. For , this is the entry of the accumulated key-value matrix. is a query-selector matrix — its -th row has a single 1 at position (the -th query neuron), so it extracts from the hidden state. The output gating neuron then computes:
This is the product of one accumulated key-value entry with the corresponding query component.
Why this selector structure: The output gating needs to pair each key-value entry with the appropriate query component — specifically, entry must be multiplied by query component to contribute to output dimension . The selector matrices achieve this by having read from the key-value part of the state and read from the query part of the state, with the row index encoding the pair. The extreme sparsity (one non-zero entry per row) is not strictly necessary — any weight configuration that extracts the correct linear combinations would work — but this sparse form makes the construction transparent.
The construction: readout summation. The readout matrix sums the output gating neurons into the output dimensions:
What this equation computes: Output dimension sums all output gating neurons whose index satisfies — that is, all neurons storing entries in row of the key-value matrix. Each such neuron contributes , so the sum over yields , which is exactly the -th output of the linear self-attention layer.
Why this summation structure: The readout matrix completes the matrix-vector multiplication by summing over the column index . This three-stage decomposition — (1) recurrence accumulates key-values, (2) output gating computes elementwise products with queries, (3) readout sums over columns — mirrors the mathematical structure of the attention computation. Each stage uses a different architectural component optimized for its role: recurrence for temporal accumulation, gating for multiplicative interaction, and linear readout for summation.
Complete information flow. Putting all components together, the full computation at time proceeds as follows. The input arrives. For each key-value neuron : input gating produces , this is added to the existing state (since ), yielding the running sum . For each query neuron : input gating produces , this replaces the previous state (since ), yielding . Output gating neuron multiplies (accumulated key-value entry) by (current query component). Readout neuron sums the output gating neurons for column indices , producing . This matches the linear self-attention output exactly.
Parameter Count and Compression Strategies
Naive construction cost. The basic construction described above requires hidden neurons and output gating neurons, leading to total parameter counts of approximately: and each have parameters, and each have parameters, and has parameters. The total parameter count is , dominated by the output gating matrices. This is substantially larger than the parameters of the linear self-attention layer being emulated — a quadratic blowup in state size and a quartic blowup in parameter count.
Why this cost is necessary: The output gating needs to represent all possible connections from hidden state to output gating neurons, because in the general case any key-value entry might need to be multiplied by any query component. This combinatorial explosion reflects a fundamental tension: the RNN represents the key-value matrix as a -dimensional vector (by flattening), and the subsequent matrix-vector multiplication requires pairwise products ( entries times query components). The RNN's architectural constraint — that gating can only multiply two scalars per neuron — forces it to explicitly enumerate all these pairwise products.
Compression via symmetry (Appendix A.3). When the value matrix is invertible, the construction can be compressed to require only hidden neurons. The key insight leverages an invariance of linear self-attention: the layer's output depends on and only through the product , not on the individual matrices. Specifically, for any invertible matrix :
If we replace with and adjust to , the product remains unchanged:
This means we can assume without loss of generality that , i.e., the key and value matrices are equal. When , the key-value matrix is symmetric — entry equals entry . Therefore, we only need to store the entries in the upper triangle (including the diagonal) rather than all entries.
What this compression achieves: The number of key-value accumulator neurons reduces from to , cutting the storage requirement nearly in half for large . The output gating still requires neurons and parameters because the matrix-vector multiplication step remains, but the hidden state size is reduced. Importantly, the query-related neurons ( of them) remain unchanged.
When this compression applies: The symmetry argument requires to be invertible so that the transformation is well-defined. If is singular, the compression may not be possible because the equivalence between different pairs breaks down. In practice, randomly initialized square matrices are invertible with probability 1, so the symmetric compression is generally applicable. The paper confirms that trained RNNs discover this compressed representation — in the teacher-student experiments, the network uses exactly key-value neurons for , rather than the naive .
Compression via low-rank structure (Appendix A.4). When the value and query-key matrices have low rank, further compression is possible. Let the singular value decompositions be:
where has only non-zero entries, and has only non-zero entries. The attention output can then be written as:
The inner key-value matrix has non-zero entries only in the first rows and columns. Therefore, only key-value entries need to be stored, and only query components are active. The total number of required recurrent neurons becomes .
What this compression achieves: When the attention matrices have rank significantly smaller than the full dimension , the RNN can implement attention with a state size that scales with the ranks rather than with . For example, with and , the construction requires only neurons instead of . The parameter count of the output gating similarly reduces because only output neurons are needed.
Why this compression matters in practice: The quadratic scaling of attention parameters with rank instead of dimension means that if real-world attention patterns are low-rank — a common empirical observation (many attention heads can be pruned or factorized with minimal performance loss) — then RNNs can implement attention with substantially smaller state sizes than the worst-case . This provides a pathway for RNNs to be practical as attention-implementers despite the theoretical quartic parameter cost. The paper validates this empirically: when teacher attention matrices have rank 6 with , trained RNNs discover solutions using approximately 42 neurons (Figure 3.B), exactly matching the low-rank construction's prediction.
Alternative construction with side gating (Appendix A.2). If the output gating is replaced by a side gating mechanism defined as:
where directly multiplies the hidden state (rather than computing a product of two state-based projections), the parameter count improves significantly. The hidden state now only needs the key-value accumulator neurons (no separate query neurons needed, because queries are provided directly through the side gate). The side gating matrix has entries, and it copies queries to the appropriate key-value entries:
What this form achieves: Instead of storing queries in the hidden state and then extracting them via output gating, the side gating provides queries directly as a function of the current input. This eliminates the query pass-through neurons (reducing hidden state from to ) and replaces the output gating matrices with the side gating matrix. The total parameter count becomes rather than . The per-neuron computation is , which multiplies the query (from the side gate) by the key-value (from the hidden state) directly.
Why this connects to Mamba: The Mamba architecture (Gu & Dao, 2023) features a similar mechanism where the output is computed as modulated by a side gating transformation — this is structurally equivalent to our side gating construction. The Mamba layer can thus implement attention in a single layer, whereas the input-output gating variant requires two layers (or a modified single layer with both input and output GLUs).
Implications for Existing RNN Classes (Section 3.3, Appendix B)
LRU requires two layers for attention (input-output gating construction). The standard LRU layer (Orvieto et al., 2023b) contains diagonal linear recurrence and an output GLU of the form applied after the recurrent state readout. It lacks an input gating mechanism. To implement the construction described above, two stacked LRU layers are needed: the first layer uses its output GLU to compute the key-value outer products (acting as ), stores them in its hidden state with , and passes queries through with . The second layer's recurrence is bypassed (or used with ), and its output GLU acts as the output gating, multiplying the accumulated key-values from the first layer by the queries. The GLU's sigmoid nonlinearity must be linearized — this can be achieved by using small weights in and large compensating weights in , approximating . Half the GLU neurons provide the linear gating, and half (with ) provide a bias term that can be subtracted in a subsequent linear transformation.
LRU In+Out enables one-layer attention. Adding an explicit input GLU to the LRU layer — so the recurrence receives — allows a single layer to implement both the input gating (computing outer products for accumulation) and the output gating (multiplying accumulated key-values by queries). The paper introduces this variant and demonstrates empirically that it significantly improves the ability to mimic attention compared to the standard LRU (Figure 4.A).
Mamba implements attention in one layer (via side gating). The Mamba layer's core recurrence is , where is an input-dependent matrix. The term can be decomposed as:
When is diagonal, this becomes a standard multiplicative gating (one branch is , the other is ). When is preceded by a linear projection, it can implement the outer product computation needed for key-values. The output of Mamba is , which is structurally equivalent to our side gating construction. Thus, a single Mamba layer can implement linear self-attention (Table 3, Appendix B.2).
LSTM can implement attention theoretically but doesn't learn it. The LSTM cell has the computational components needed: the input gate multiplied by the cell update candidate can compute key-value outer products, the cell state can accumulate them (with forgetting controlled by ), and the output gate can multiply by queries when the tanh is linearized (Table 4, Appendix B.3). The LSTM requires only one layer because the output gate directly provides the query to the output without needing separate query-pass-through neurons. However, the construction requires linearization of both the tanh in and the sigmoid in and , which imposes specific constraints on weight magnitudes (small weights for linearization, compensated by large readout weights). When actually trained on the teacher-student task, LSTMs fail to achieve low loss — they perform similarly to GRUs (Figure 4.A) — indicating that gradient descent does not find the attention implementation within the LSTM's parameter space under standard training conditions. This negative result is important: theoretical capacity does not imply learnability.
GRU cannot implement attention at all. The GRU lacks a mechanism to compute the product of two different projections of the input that produces the outer product. The update gate modulates the mixing of old and new states but does not multiply different input projections. The reset gate gates the previous hidden state but only in service of computing the update candidate , not for producing multiplicative interactions of input features. Stacking multiple GRU layers does not remedy this because no layer can compute key-value outer products — the fundamental primitive is architecturally absent. This is formalized in Table 3 and Appendix B.4.
Teacher-Student Experimental Methodology (Section 4)
Task formulation. The teacher-student experiments train a gated RNN to emulate the output of a randomly initialized linear self-attention layer. The teacher has -dimensional inputs and outputs, with sampled from a normal distribution. Training sequences have length 32, with inputs sampled i.i.d. from a normal distribution at each position. The teacher processes the entire sequence causally, producing outputs at each position. The student RNN receives the same input sequence and produces outputs , trained to minimize the mean squared error across all timesteps and dimensions.
Why i.i.d. inputs: Using independent inputs at each timestep ensures that the teacher's output depends on the entire sequence history, not on any temporal structure in the inputs themselves. This forces the student to learn the teacher's temporal processing algorithm (accumulating key-values and multiplying by queries) rather than exploiting input correlations. Each sequence is presented only once (no epoch repetition of specific sequences) to prevent memorization of specific input-output patterns.
Student architecture and hyperparameters. The student is a gated RNN as defined in Equation 2, with input dimension , hidden state dimension (varied experimentally), and output dimension . For the main teacher identification experiment (Section 4.1), , hidden neurons, and output gating neurons. The parameters are parameterized as (similar to the LRU's exponential parameterization) to ensure they stay in and prevent state explosion. Training uses the AdamW optimizer (Loshchilov & Hutter, 2019) with initial learning rate , cosine annealing to over the training duration, weight decay applied to all parameters except (the parameters). Training runs for approximately 1 million iterations (50,000 sequences per epoch, 1000 epochs) with batch size 64 and sequence length 32.
Why this optimization setup: AdamW with cosine annealing provides stable optimization for the bilinear gating operations, which can have poorly conditioned loss landscapes. The exclusion of from weight decay is important because weight decay would bias toward zero (forgetting), which would prevent the network from learning the perfect-memory accumulators () needed for the attention implementation.
Reverse-engineering analysis: pruning. After training, many weights in , , , , and become exactly zero or very close to zero. The authors prune all hidden neurons whose input weights (rows of and combined) are entirely zero, and all output gating neurons whose output weights (columns of and , and corresponding rows of ) are entirely zero. This pruning is lossless — the pruned network computes exactly the same function as the original, because removed neurons contribute nothing at any stage. In the experiment, pruning removes 86 out of 100 hidden neurons and 87 out of 100 output gating neurons, leaving a compact network of 14 hidden and 13 output gating neurons (Figure 2.A summarizes the post-processing).
What this pruning reveals: The surviving neurons after pruning are precisely those that contribute to the network's function. Their values, weight patterns, and connectivity reveal what computation the network learned to perform. The fact that so many neurons can be pruned losslessly indicates that the network was substantially overparameterized relative to the solution it found — the effective architecture discovered by training is much smaller than the one provided.
Reverse-engineering analysis: value categorization. The authors categorize surviving hidden neurons into two groups based on their learned values: those with (perfect memory) and those with (no memory). In the experiment, the surviving 14 hidden neurons split into two clean groups (Figure 2.B): 10 neurons with (key-value accumulators) and 4 neurons with (query carriers). This matches the symmetric construction's prediction: neurons for key-values, neurons for queries.
Reverse-engineering analysis: linear reconstruction scores. To verify that the neurons encode key-values and the neurons encode queries, the authors compute Score KV and Score Q (Table 1). For Score KV: at each timestep , construct the true key-value matrix , flatten it into a vector of length , and fit a linear regression from the activations of the 10 neurons to this flattened vector. The score is of this regression — a value near 0 means the key-value vector is perfectly linearly reconstructable from the neuron activations. For Score Q: similarly regress the true query vector from the activations of the 4 neurons and report . The achieved scores are and (Table 1), indicating essentially perfect linear reconstruction.
What these scores establish: The neurons form an exact linear basis for the key-value matrix — there exists a linear transformation mapping their activations to the flattened key-value matrix. Combined with (perfect accumulation), this means the hidden state at time is exactly a linear encoding of . Similarly, the neurons are a linear encoding of the current query . This confirms that the temporal processing component of the RNN correctly implements the attention accumulation and query pass-through.
Reverse-engineering analysis: polynomial coefficient comparison. The output gating and readout together define the instantaneous mapping from hidden state to output. Since the hidden state linearly encodes key-values and queries (from the reconstruction scores), the full input-output mapping of the RNN at a single timestep can be expressed as a polynomial in the input dimensions. For the teacher linear self-attention, this mapping is the degree-3 polynomial . For the student RNN, since input gating produces degree-2 polynomials and output gating multiplies degree-2 by degree-2, the instantaneous mapping is a degree-4 polynomial in the input (with some coefficients potentially zero). The authors extract the coefficients of both polynomials for each output dimension, normalize them into vectors, and compute the L2 distance between the teacher's coefficient vector and the student's coefficient vector — the Polynomial distance (averaged over the 4 output dimensions).
What this metric captures: Even if the RNN's internal representations differ from the construction (due to the invariances discussed earlier — permutations, scalar rescalings, invertible transformations), the input-output polynomial must match that of the teacher if the RNN has truly learned the attention function. This is a functional equivalence test that transcends the specifics of weight representation. The achieved polynomial distance is (Table 1), confirming that the RNN implements the same degree-3 polynomial as the teacher (the degree-4 coefficients are essentially zero). This test covers the entire input domain, not just the training distribution — it establishes that the RNN has learned the attention function, not just memorized training examples.
Reverse-engineering analysis: weight matrix structure. After pruning and reordering neurons (grouping and neurons, permuting rows of gating matrices for visual clarity), the learned weight matrices exhibit the block-sparse structure predicted by the construction (Figure 2.B). Specifically: and have a clear separation — rows corresponding to key-value neurons have non-zero entries only in the value and key projection columns respectively, while rows for query neurons have non-zero entries in the query projection columns (and the constant column for ). and have the selector structure — reads from key-value neurons, reads from query neurons. The readout has the summation structure over column indices.
What the structure confirms: The learned network has not just any solution that achieves low loss — it has specifically discovered the structure of the constructive proof: separate neurons for key-values and queries, binary values, output gating that multiplies key-value entries by query components, and readout that sums over columns. This means gradient descent, when provided with the right architectural elements (multiplicative gating, diagonal recurrence, learnable per-neuron decay), reliably converges to the attention implementation rather than some alternative computation that happens to match the training outputs.
Handling non-construction weights: neuron merging (Appendix C.2). In the learned network of Figure 2.B, three output gating neurons (rows 11, 12, 13 of the output gating matrices) show a pattern that doesn't match the construction's clean selector structure. However, these three neurons can be merged into a single neuron with the construction structure. The contribution of these three neurons to output dimension is:
This is a quadratic form in the hidden state, with a kernel that is the weighted sum of three rank-1 matrices. The authors observe that for all four output dimensions, these kernels are proportional to each other and are themselves rank-1 (Figure 2.C). Therefore, the three neurons can be replaced by a single neuron whose and are the singular vectors of the combined kernel, and whose readout weight is the appropriate scaling factor. This merged neuron now exhibits the construction pattern: one vector reads from key-value neurons, the other reads from query neurons.
What this merging reveals about invariances: The existence of such redundant representations is a consequence of the invariances in the gated RNN parameterization. The output gating produces a sum of outer products of its weight rows, and any decomposition of the resulting kernel into rank-1 components is equivalent. The network can represent the same function with different numbers of output gating neurons, different weight matrices, and different patterns of readout weights — as long as the combined quadratic form matches. This makes exact weight retrieval nearly impossible (as the paper notes), but the structural signature (one vector from key-value neurons, one from query neurons) remains identifiable after appropriate merging.
Overparameterization requirement (Section 4.2, Figure 3.A). The paper systematically varies the number of hidden neurons and output gating neurons (keeping them equal) and measures the final training loss. For , the symmetric construction requires hidden neurons. The experiments show a sharp transition: when (roughly twice the construction minimum), the loss remains high — the network fails to learn the attention function. When , the loss drops to near-zero — the network successfully learns the construction. This "mild overparameterization" requirement (approximately the theoretical minimum) suggests that gradient descent needs extra capacity to navigate the loss landscape and find the sparse solution, after which the unused neurons can be pruned.
What this threshold implies for practical architectures: If an RNN is to learn attention on a task, its hidden state dimension must be at least twice the theoretical minimum for the construction given the input dimension and any applicable compression (symmetry, low-rank). This provides a concrete sizing guideline: for a -dimensional input with full-rank attention matrices, allocate at least hidden neurons.
Low-rank teacher experiment (Figure 3.B). To test whether trained RNNs discover the low-rank compression, the authors use a teacher attention layer with but with constrained to have rank 6 (by removing 6 singular values from randomly initialized matrices). The low-rank construction predicts that neurons suffice. The experiments show a loss transition at approximately 84 hidden neurons — again about the theoretical minimum — confirming that the network discovers and exploits the low-rank structure rather than naively allocating neurons.
Architectural comparisons (Section 4.3, Figure 4.A). The paper compares several architectures on the teacher-student task with : (1) the simplified gated RNN with diagonal recurrence, (2) a variant with dense linear recurrence (instead of diagonal), (3) one-layer LSTM, (4) one-layer GRU, (5) standard LRU (output GLU only, requiring 2 layers), (6) LRU In+Out (both input and output GLU, 1 layer), (7) LRU In+Out but with the GLU replaced by an MLP with equal parameter count. All architectures use a linear embedding to project the input to a latent dimension, the recurrent layer(s), and a linear readout to the output dimension. Recurrent state dimensions are set to 60 for LSTMs and GRUs and to the same value for the simplified RNNs for fair comparison (detailed parameter-count-controlled comparisons appear in Appendix Figure 6).
Key findings from the architecture comparison: (1) Diagonal versus dense recurrence makes little difference — the diagonal inductive bias is helpful but not essential for learning attention (both achieve similarly low loss). (2) LSTMs and GRUs achieve high loss and fail to learn attention — despite LSTMs having the theoretical capacity, gradient descent does not find the solution. The nonlinearities (sigmoid, tanh) appear to obstruct the optimization path. (3) The standard LRU (2 layers) achieves moderate loss but is outperformed by the LRU In+Out (1 layer with explicit input gating), confirming that adding input gating strengthens the inductive bias toward attention. (4) Replacing the GLU in LRU In+Out with an MLP causes a dramatic performance degradation, demonstrating that multiplicative interactions are fundamental — a feedforward network with the same number of parameters but without elementwise multiplication cannot implement the outer product computation central to attention. (5) The construction baseline (simplified gated RNN) achieves the lowest loss, consistent with it being explicitly designed to match attention.
In-Context Learning Experimental Methodology (Section 5)
Task formulation: in-context linear regression. The task is designed so that the optimal computation for solving it is one step of gradient descent, which can be implemented by linear self-attention with a specific weight structure. Each "task" (sequence) is characterized by a ground-truth linear mapping with entries sampled i.i.d. from . The sequence presents observations where is sampled i.i.d. from a uniform distribution , and . After observing pairs, the network is queried at position with an input — the concatenation of a new input and a zero placeholder — and must output to be compared against the true .
Why this task structure: The zero placeholder at the query position signals to the network that this is a prediction step, not another observation. Since is not provided, the network must use the previously observed pairs to infer and apply it to . The linearity of the underlying function and the Gaussian/uniform distributions make one-step gradient descent with zero initialization an analytically optimal solution (Mahankali et al., 2023).
Connection to gradient descent. One step of gradient descent on the mean squared error starting from with learning rate yields:
The prediction for is then:
What this equation represents: The prediction is a sum over all past observation pairs, where each past output is weighted by the inner product between the corresponding past input and the current query input. This is exactly the computation performed by linear self-attention if we identify (values are the outputs), (keys are the inputs), and (the query is the new input). Specifically, von Oswald et al. (2023) showed that when biases are zero, a linear self-attention layer with , , and implements precisely this gradient descent computation.
Why gradient descent is the optimal solution: Mahankali et al. (2023) proved that for the in-context linear regression task with these data distributions, one step of gradient descent with an optimal learning rate achieves the globally minimum expected loss among all functions implementable by a single layer of linear self-attention. The optimal learning rate can be derived analytically (Appendix D.2) and depends on the input dimension , the number of observations , and the variance of the input distribution:
where for the distribution (variance 1). For the paper's setting with and , this gives a specific optimal learning rate.
Experimental setup. The input dimension is (concatenating inputs and outputs). The sequence length is observations plus one query step. The gated RNN student uses 80 hidden neurons by default. Training uses the same optimization scheme as the teacher-student experiments (AdamW, cosine annealing from to , weight decay on non- parameters) but for only 300,000 iterations. During validation, the distribution of entries is changed to — double the training variance — to test whether the network has learned the algorithm (which should generalize to any with the right statistics) or has overfit to the specific training distribution of .
Reverse-engineering analysis: loss comparison. After training, the student RNN achieves a training loss of 0.0945, which is nearly identical to the loss of one optimal gradient descent step (0.0947). This near-equality is the first piece of evidence that the RNN has learned to implement gradient descent — no strategy that merely approximates gradient descent could achieve lower loss than the optimal step, and the closeness suggests exact implementation.
Reverse-engineering analysis: polynomial coefficient extraction (Table 2). The instantaneous function mapping the current concatenated input to the output prediction is a polynomial of degree at most 4 (since input gating produces degree-2 polynomials and output gating multiplies degree-2 by degree-2). For gradient descent, the prediction for output dimension at the query step is:
This is a sum of terms each proportional to — a degree-3 polynomial in the observation inputs and the query input. At any non-query step , the network's output is not directly trained but may be produced by the same mechanism.
The authors extract the coefficients of all degree-4 and lower terms of the polynomial implemented by the first output neuron of the trained RNN, averaged over 4 random seeds. The results (Table 2) show that only three terms have non-negligible coefficients: , , and , each with coefficient approximately . The optimal gradient descent prediction uses exactly these terms with coefficient (where and is computed from the formula above). The "residual norm" — the norm of all other polynomial coefficients — is , close to zero.
What this establishes: The trained RNN implements exactly the same polynomial as one optimal step of gradient descent. It has learned to extract and from the input, accumulate their products in the neurons, pass through in the neurons, and combine them through output gating to produce the prediction. This is a specific low-rank instance of the attention construction — the value matrix extracts (dimension ), the key matrix extracts (dimension ), and the accumulated key-value matrix has rank at most . The compressed construction would therefore require neurons.
Reverse-engineering analysis: parameter count scaling (Figure 3.C). The paper varies the number of hidden neurons in the gated RNN and measures the difference between the RNN's loss and the optimal gradient descent loss ("delta loss"). For (input dimension), the low-rank construction predicts that neurons suffice. The experiments show a sharp transition: when the RNN has fewer than approximately 24 neurons (again, roughly overparameterization), the delta loss remains high. When the RNN has 24 or more neurons, the delta loss drops to near zero. This confirms that the RNN discovers the rank-structured solution using the minimal number of neurons predicted by the low-rank compression.
What this means mechanistically: The RNN does not allocate neurons (the symmetric full-rank construction). Instead, it recognizes that the gradient descent attention matrices are low-rank and uses only the key-value accumulator neurons (storing the key-value matrix entries) and query neurons — a total of neurons, which with overparameterization requires about neurons to reliably learn. The observed transition at approximately 24 neurons is consistent with this analysis.
Architectural comparisons on in-context learning (Section 5.3, Figure 4.B). The paper compares LSTMs, GRUs, and several LRU variants on the in-context linear regression task with out-of-distribution generalization (validation variance doubled). The key findings contrast with the teacher-student results: (1) Nonlinearity helps significantly — LSTMs and GRUs, which failed at teacher-student attention mimicking, achieve the best in-context learning performance when given sufficient parameters. The GRU performs best overall, despite being the architecture furthest from attention. (2) Within LRU variants, there is a high correlation between attention-mimicking ability and in-context learning performance — the LRU In+Out (which most closely matches the construction) significantly outperforms the standard LRU, and replacing the GLU with an MLP causes a large performance drop. (3) All nonlinear architectures with sufficient capacity outperform one optimal gradient descent step, suggesting they implement more sophisticated in-context learning algorithms than simple gradient descent.
What this discrepancy reveals: The fact that LSTMs and GRUs perform well at in-context learning despite failing to implement the specific attention construction indicates that they succeed through different mechanisms. This is a crucial finding: while gated linear RNNs succeed at in-context learning by discovering attention, nonlinear RNNs like LSTMs and GRUs can achieve even better performance through alternative algorithms that exploit their nonlinearities. The paper thus demonstrates that attention is a sufficient but not necessary mechanism for in-context learning in RNNs. The correlation within the LRU family suggests that for architectures with weak nonlinearities (the LRU's sigmoid in the GLU can be approximately linearized), the attention mechanism is the primary path to in-context learning. For strongly nonlinear architectures (LSTM, GRU), other mechanisms become accessible and can outperform attention.
4. Key Insights and Innovations
Innovation 1: Finite Constructive Equivalence Between Gated RNNs and Attention — Not Just Universality, But an Explicit Blueprint
The prior theoretical landscape for understanding RNNs in relation to attention was dominated by two incomplete frameworks: asymptotic universality theorems (Boyd & Chua, 1985; Grigoryeva & Ortega, 2018) that guarantee RNNs could approximate attention given infinitely many neurons, and conceptual parallels between linearized attention and fast weight memories (Schmidhuber, 1992; Schlag et al., 2021) that work in the opposite direction — showing attention can be viewed as recurrence, not that recurrence can implement attention. Neither framework answers the practically decisive question: can a standard RNN with a finite state implement attention exactly, and if so, what architectural elements are necessary and how many neurons does it take?
This paper's distinctive intellectual move is to shift from asymptotic approximation to finite constructive equivalence. Rather than proving that an RNN's output can get arbitrarily close to attention's output as the state grows, the authors provide an explicit assignment of numerical weight values (Appendix A.1, Equations 9–14) that makes a gated RNN compute the identical function as any given linear self-attention layer, using hidden neurons and parameters for a -dimensional input. This is not an existence proof — it is a blueprint specifying exactly which neurons store which entries of the key-value matrix, which neurons pass through queries, how the input gating computes outer products, how the output gating multiplies accumulated key-values by queries, and how the readout sums the results.
What makes this contribution conceptually novel, rather than a routine constructive proof, is that it identifies the minimal architectural requirements for attention in a recurrent network. By working with a deliberately simplified model class — diagonal linear recurrence with multiplicative input and output gating — the construction isolates exactly which computational primitives are essential: (1) a bilinear operation to compute outer products of input projections (input gating as ), (2) perfect-memory recurrence () for accumulating these outer products over time, (3) instantaneous memory () for holding the current query, and (4) a second multiplicative operation to multiply the accumulated matrix by the query vector (output gating), followed by summation (readout). Remove any of these elements, and the construction fails — a fact the paper validates empirically by showing that replacing the GLU with an MLP of equal parameter count destroys the ability to mimic attention (Figure 4.A).
This decomposition into necessary primitives is a fundamental conceptual contribution because it explains why modern deep linear RNNs work well on tasks where attention excels: they contain these primitives, often in forms not recognized as attention-enabling by their designers. The LRU's output GLU and diagonal recurrence provide primitives (2) and (4); stacking two layers lets the first layer's output gating serve as the second layer's input gating, providing primitive (1). Mamba's input-dependent matrices and side gating provide primitives (1) and (4) in a single layer. The construction thus serves as a diagnostic tool: given any proposed RNN architecture, one can ask whether it contains the multiplicative interactions necessary for computing outer products and matrix-vector products, and predict its capacity for attention-like routing of information. The paper demonstrates this diagnostic value by showing that GRUs — which lack any mechanism for computing multiplicative outer products of input features — cannot implement attention even with stacking (Section 3.3, Appendix B.4), while LSTMs theoretically can but gradient descent doesn't find the solution (Section 4.3).
This is a fundamental shift from the field's prior framing. Before this paper, the question was "Can RNNs approximate what attention does?" — answered by universality theorems, which imply that with enough neurons, anything is possible. After this paper, the question becomes "Does this specific RNN architecture contain the primitives needed for attention, and if so, at what parameter cost?" — answered by checking against the construction template. This reframes architecture design from an empirical trial-and-error process to a principled one guided by computational requirements.
Innovation 2: Gradient Descent Discovers the Construction — Not Just Capacity, But Learnability
It is one thing to prove that a set of weights exists making an RNN equivalent to attention. It is entirely another to show that gradient descent actually finds these weights when the RNN is trained. Many theoretically possible configurations are unreachable by gradient-based optimization — the loss landscape may not connect the initialization to the solution, or the solution may be unstable under standard training dynamics. The paper's second major contribution is demonstrating, through careful reverse-engineering of trained networks, that standard gradient descent reliably discovers the attention implementation under the right architectural conditions (specifically, the simplified gated RNN with diagonal recurrence and linear gating).
The evidence for this claim comes from the teacher-student experiments (Section 4.1), where trained RNNs exhibit unmistakable signatures of the construction: (a) hidden neurons cleanly bifurcate into (perfect memory) and (no memory) groups (Figure 2.B), matching the construction's requirement for key-value accumulators and query pass-through; (b) the key-value matrix can be linearly reconstructed from the neuron activations with near-perfect accuracy (Score KV = , Table 1), and queries from the neurons (Score Q = ); (c) the polynomial coefficients of the trained RNN's instantaneous input-output mapping match those of the teacher attention layer with distance (Table 1), establishing functional equivalence across the entire input domain; and (d) the weight matrices exhibit the block-sparse structure predicted by the construction (Figure 2.B) — input gating rows for key-value neurons read from value and key projections respectively, output gating rows have the selector pattern multiplying key-value state entries by query state entries, and the readout sums over column indices.
What makes this finding intellectually distinctive is not just that learning succeeds, but how the network learns: it discovers the compressed version of the construction. Rather than naively allocating key-value neurons for , the trained network uses only , exploiting the symmetry of the key-value matrix when is invertible (Appendix A.3). Rather than using 16 key-value neurons when the teacher attention matrices have low rank, it uses only instead of 156 (Figure 3.B). The network doesn't just find any attention implementation — it finds the parameter-efficient one, suggesting that gradient descent in overparameterized gated RNNs has an inductive bias toward compressed representations that exploit algebraic structure in the target function.
This connects to a broader theme in deep learning theory — the implicit bias of gradient descent toward low-complexity solutions — but applies it in a novel and diagnostically precise way. The paper shows that the "complexity" being minimized is not a generic norm but specifically the rank and symmetry structure of the attention matrices being emulated, implemented through the minimal number of recurrent neurons.
The contrast with LSTMs is equally revealing as a negative result with theoretical implications: LSTMs can implement attention in one layer (Table 3, Appendix B.3), but when trained on the teacher-student task, they fail to achieve low loss and perform similarly to GRUs (Figure 4.A). This establishes a crucial distinction between theoretical capacity and empirical learnability that the field often elides. The nonlinearities in LSTMs (sigmoid gates, tanh activations) apparently create loss landscapes where the attention solution is not reachable by gradient descent from standard initializations, even though it exists in parameter space. This explains why prior work may have missed the connection between RNNs and attention: the architectures that were historically dominant (LSTMs, GRUs) either cannot implement attention at all (GRUs) or cannot learn to do so under standard training (LSTMs). The recent deep linear RNN architectures succeed precisely because their simplified nonlinearities (linearized GLUs, diagonal recurrence) make the attention solution more accessible to gradient descent.
This is a fundamental conceptual contribution, not an incremental empirical finding, because it reframes the question from "What can RNNs theoretically compute?" to "What do RNNs actually learn to compute, and why?" — and provides a rigorous template for answering that question through weight-level reverse-engineering.
Innovation 3: Multiplicative Interactions as the Computational Primitive for Attention — Redefining the Role of Gating
Gating mechanisms are ubiquitous in modern RNNs — the GLU in LRUs, the sigmoid gates in LSTMs and GRUs, the input-dependent matrices in Mamba — but their theoretical function has been framed primarily in terms of optimization: they mitigate vanishing gradients, they provide nonlinearities for representation learning, they enable selective memory updating. The paper's third innovation is to reconceptualize gating as the computational primitive that enables attention-like routing of information, shifting the understanding of gating from an optimization aid to a constructive computational mechanism.
The key insight is that elementwise multiplication — the operation at the heart of all gating mechanisms — is exactly the operation needed to compute outer products and matrix-vector products, the two core computations of linear self-attention. The input gating computes degree-2 polynomials that can represent entries of the outer product when the two linear projections extract values and keys respectively. The output gating pairs accumulated key-value entries with query components. Without multiplicative interactions, an RNN could only accumulate linear functions of past inputs — it could never represent the quadratic key-value accumulations that are attention's core state. With them, the RNN gains the capacity to dynamically route information between arbitrary past positions and the current input, exactly as attention does.
What makes this framing genuinely novel is that it inverts the standard narrative about gating. The standard view — inherited from the LSTM literature — is that gating controls information flow: the forget gate decides what to discard, the input gate decides what to add, the output gate decides what to expose. This is a selection or filtering role. The paper's view is that gating performs computation: it multiplies representations together to create new representations that are products of the originals, enabling the network to compute functions (specifically, bilinear functions of past and current information) that purely additive recurrence cannot express. This is not "controlling" information flow — it is creating information through multiplicative synthesis.
The empirical validation of this reconceptualization is striking and multi-pronged:
-
Architecture modification as causal test (Figure 4.A): Adding an input GLU to the LRU architecture (creating "LRU In+Out") significantly improves its ability to mimic attention compared to the standard LRU (which has only output gating), because a single layer now contains both multiplicative primitives (outer product computation and matrix-vector multiplication) rather than requiring two stacked layers where the first layer's output gating must double as the second layer's input gating.
-
Ablation of multiplicative interactions (Appendix B.1): Replacing the GLU in LRU In+Out with a 1-hidden-layer MLP of equal parameter count — which can represent any smooth function but lacks explicit elementwise multiplication — causes a dramatic performance degradation. This is a clean ablation: the MLP has equal representational capacity in terms of parameter count, but lacks the architectural inductive bias for computing outer products through multiplication. The failure demonstrates that multiplicative gating is not merely providing additional nonlinearity — it is providing a specific computational primitive that feedforward nonlinearities cannot efficiently replicate.
-
GRU impossibility result (Section 3.3, Appendix B.4): The GRU contains gating (reset gate, update gate) but these gates modulate the flow of information rather than multiplying different input projections together. The GRU's update candidate involves a tanh of a linear combination of input and (reset-gated) previous state — this is additive inside the nonlinearity, not multiplicative. As a result, stacked GRU layers cannot implement linear self-attention even in principle, because no layer can compute the key-value outer product. This negative result sharpens the positive claim: it's not "gating" in general that enables attention, but specifically multiplicative gating of the form where and are different linear projections of the same input.
-
In-context learning correlation (Section 5.3): Within the LRU family, there is a strong correlation between an architecture's ability to mimic attention (on the teacher-student task) and its in-context learning performance on the linear regression task. The LRU In+Out significantly outperforms the standard LRU and the MLP variant, suggesting that the attention mechanism is the bridge between these two capabilities — architectures that can implement attention do so, and this improves their in-context learning.
This is a fundamental reframing, not an incremental finding, because it gives the field a new vocabulary for thinking about what gating does. Before this paper, a practitioner choosing between an LSTM, a GRU, and an LRU for a sequence modeling task might think in terms of "vanishing gradients" or "long-range memory." After this paper, they can think in terms of "does this architecture have the multiplicative primitives to implement attention-like dynamic routing of information?" — a more precise and mechanistically grounded question that directly connects architectural choices to computational capabilities.
Innovation 4: The Parameter-Cost-of-Attention Diagnostic — Quantifying When and Why RNNs Are Inefficient Attention Implementers
The paper's constructive proof reveals something that was not obvious from universality arguments or conceptual analogies: implementing attention in an RNN is parameter-inefficient, and the inefficiency has a specific, quantifiable structure. The naive construction requires parameters to emulate a linear self-attention layer that itself has only parameters — a quadratic blowup in the total parameter count relative to the dimension (Section 3.2). This is not an artifact of a suboptimal construction; it reflects a fundamental mismatch between how attention represents its state (as a matrix) and how an RNN represents its state (as an -dimensional vector). The RNN must "flatten" the matrix into the vector, then — because its only multiplicative operation is scalar multiplication in the output gating — explicitly instantiate all pairwise products of matrix entries and query components through separate neurons and weights.
This insight is intellectually distinctive because it transforms a vague intuition ("RNNs are less efficient than Transformers for certain tasks") into a precise, quantifiable framework. The paper shows that the parameter cost of attention implementation depends on compressible structure in the attention matrices:
-
Full-rank case (Section 3.2, Appendix A.1): Requires hidden neurons and total parameters. The RNN is paying a quadratic premium in state size and a quartic premium in parameters compared to the attention layer it emulates.
-
Symmetric case (Appendix A.3): When is invertible, the key and query matrices can be transformed so that , making the key-value matrix symmetric. This reduces hidden neurons to — roughly half — but the output gating parameter count remains . The paper shows that trained networks discover this compression (Figure 2.B: 10 key-value neurons for instead of 16).
-
Low-rank case (Appendix A.4): When the value matrix has rank and the query-key product has rank , only neurons are needed. The output gating similarly compresses because only output neurons are active. The paper validates this with rank-6 teachers for : trained networks use approximately 42 neurons instead of the full-rank 156 (Figure 3.B).
-
Gradient descent case (Section 5.2): When trained on in-context linear regression, the network discovers that the attention matrices implementing gradient descent have rank , requiring only 12 key-value neurons instead of the full symmetric construction's 21. The loss transition occurs at approximately 24 hidden neurons (Figure 3.C), matching the predicted overparameterization threshold.
This framework provides a diagnostic tool for the field: given the dimension and estimated rank structure of attention patterns in a real-world task, one can predict how many RNN hidden neurons would be needed to implement attention. If the required state size exceeds what's practical, RNNs cannot efficiently replicate attention on that task. If the required state size is manageable, RNNs become a viable alternative.
More importantly, this contribution explains the performance landscape of modern RNN architectures in a way that was previously unavailable. The Mamba architecture (Gu & Dao, 2023) uses input-dependent state projections and side gating, which provide a stronger inductive bias toward low-rank attention implementation. Its empirical success on tasks where Transformers excel may be partly attributable to its ability to implement compressed attention efficiently. Conversely, the LRU's need for two stacked layers to implement full input-output gating may limit its attention capacity relative to Mamba. The paper's framework provides the vocabulary for making these comparisons precise and for guiding architecture design toward more efficient attention implementation.
This is a fundamental conceptual contribution rather than an incremental empirical finding, because it provides a quantitative bridge between two architectural paradigms that were previously compared only through benchmark scores. The vs. parameter cost comparison is a complexity-theoretic diagnostic that will remain valid regardless of future architectural improvements — any RNN implementing attention must pay some version of this cost, and understanding where the cost comes from (the matrix-vector multiplication step requiring explicit enumeration of pairwise products) points to which architectural innovations could reduce it (e.g., side gating reduces to , as shown in Appendix A.2).
Innovation 5: Attention as a Convergent Algorithm Across Architectures — Not an Architectural Property, But a Task-Determined Solution
The paper's most subtle but potentially deepest contribution emerges from the in-context learning experiments (Section 5): gated RNNs and linear self-attention converge to the identical algorithm when trained on the same task, despite having fundamentally different architectural parameterizations. Specifically, both architectures learn to implement one step of gradient descent with the optimal learning rate on the in-context linear regression task (Table 2, loss values), and the RNN does so using the minimal number of neurons permitted by the low-rank structure of the gradient descent attention matrices (Figure 3.C).
This is not obvious a priori. The RNN has a different parameterization, different optimization landscape, and different inductive biases than the linear self-attention layer. It could have discovered a completely different algorithm that achieves the same loss — perhaps some form of recursive least squares implemented through the recurrent dynamics, or a momentum-based estimator that leverages the parameters to perform exponential moving averages. The fact that it converges to gradient descent specifically — the same algorithm that linear self-attention provably converges to (von Oswald et al., 2023; Mahankali et al., 2023) — suggests that the algorithm is determined by the task structure, not the architecture.
This is a profound reframing of how to think about neural architecture. The standard view is that architecture determines algorithm: Transformers attend, RNNs accumulate, CNNs convolve. The paper's finding suggests a more nuanced picture: for certain tasks, there exists an optimal algorithm (gradient descent for linear regression), and architectures with sufficient capacity and the right computational primitives will discover that algorithm regardless of their parameterization. The architecture's role is not to determine what algorithm is learned, but to determine whether the algorithm can be represented and how efficiently it can be parameterized.
The contrast between the gated linear RNN and the LSTM/GRU on in-context learning (Figure 4.B) adds crucial texture to this picture. LSTMs and GRUs, which could not learn the attention implementation in the teacher-student setting, actually outperform one gradient descent step on in-context learning when given sufficient parameters — the GRU achieves the best performance overall. This means that while the linear architectures (gated RNN, LRU family) converge to the attention-based gradient descent algorithm, the nonlinear architectures discover different, more powerful algorithms that exploit their nonlinearities. The architecture doesn't determine whether a task can be solved, but it does determine which family of algorithms is accessible.
This insight reconciles conflicting findings in the literature in a way that goes beyond the specific results of this paper. Prior work has alternated between "RNNs can't do X" and "actually, new RNNs can do X" without a framework for understanding why the new ones succeed. This paper's framework provides that understanding: the new RNNs succeed on attention-like tasks not because they're "better RNNs" in some generic sense, but because they contain the specific multiplicative primitives that enable attention implementation. The LSTM and GRU failed not because they're "worse" architectures, but because their gating mechanisms serve a different computational role (information flow control rather than multiplicative synthesis) that doesn't enable attention — yet their nonlinearities enable other algorithms that can outperform attention on certain tasks.
This is a fundamental conceptual shift from an architectural-determinism view ("attention is what Transformers do") to a task-determinism view ("gradient descent is what this task demands, and multiple architectures converge to it if they can represent it"). It implies that the relevant question for architecture comparison is not "Which architecture is better?" but "Which algorithms can this architecture represent, and at what cost?" — a question that the paper's constructive framework provides tools for answering.
The evidence for this convergent-algorithm claim is unusually rigorous: the polynomial coefficient extraction in Table 2 shows that the trained RNN's input-output function is identical to gradient descent term-by-term, not just in aggregate loss. The residual norm of for non-gradient-descent polynomial terms confirms that the network has not learned a slightly different algorithm that happens to have similar loss — it has learned exactly gradient descent. This level of mechanistic precision in establishing algorithmic convergence across architectures is rare in the deep learning literature and represents a methodological contribution in its own right.
5. Experimental Analysis
Evaluation Methodology
-
Dataset. The teacher-student experiments use synthetic data: input sequences of length 32 with sampled i.i.d. from a normal distribution, processed by a randomly initialized linear self-attention layer to produce targets. The in-context linear regression task uses sequences of 13 steps ( observations plus one query), with entries sampled i.i.d. from during training and during validation, and inputs sampled from to achieve unit variance. The associative recall task (Appendix D.3) uses one-hot encoded inputs of dimension with , where each and appears exactly once per sequence.
-
Base model(s). The primary model is the simplified gated diagonal linear RNN defined in Equation 2, with input dimension , hidden state dimension (varied experimentally: 100 for the main identification experiment with , 80 by default for in-context learning), and output dimension . For architectural comparisons (Section 4.3 and Section 5.3), the paper additionally trains one-layer LSTMs, one-layer GRUs, standard LRUs (with output GLU only, requiring 2 stacked layers), LRU In+Out (with both input and output GLU, 1 layer), and LRU In+Out where the GLU is replaced by a 1-hidden-layer MLP with equal parameter count. Recurrent state dimensions for LSTMs and GRUs are set to 60 for the teacher-student comparison; parameter-count-controlled comparisons appear in Appendix Figure 6. All architectures use a linear embedding layer to project the input to a latent representation followed by the recurrent layer(s) and a linear readout.
-
Metrics. The teacher-student experiments use mean squared error (MSE) loss averaged over all sequence positions and output dimensions: . For in-context linear regression, the metric is the MSE between the predicted and the true . For reverse-engineering analysis, specialized metrics are used: Score KV is of a linear regression predicting the flattened true key-value matrix from neuron activations; Score Q is for predicting the true query from neuron activations; Polynomial distance is the normalized L2 distance between the coefficient vectors of the degree-4 polynomials representing the teacher attention layer's and the student RNN's instantaneous input-output mappings; delta loss is the difference between the RNN's training loss and the loss achieved by one optimal gradient descent step. For associative recall, cross-entropy loss is used.
-
Baselines. In the teacher-student experiments, the baseline is the randomly initialized linear self-attention layer itself (the "teacher"), and the question is whether the student RNN can match its output perfectly, not whether it can outperform some alternative method. In the in-context linear regression experiments, the primary baseline is one step of gradient descent with optimal learning rate (Mahankali et al., 2023; von Oswald et al., 2023), which provably achieves the minimum expected loss for single-layer linear self-attention on this task. The optimal learning rate is analytically derived (Appendix D.2) as , giving a value of approximately for the paper's setting. For architectural comparisons, the baseline implicit in the figures is the loss achieved by different architectures relative to each other and to the gradient descent optimum.
-
Generation budget / compute accounting. The paper does not use a "generation budget" concept akin to LLM sampling. Instead, all methods process the entire input sequence causally, and the comparison is on functional equivalence of the learned input-output mapping. The relevant cost measure is number of parameters (for comparing architectures, as in Appendix Figure 6) and number of hidden neurons (for studying the transition from inability to ability to learn the attention function, as in Figures 3.A-C). For the teacher-student experiments, training runs for approximately 1 million iterations (50,000 sequences per epoch, 1,000 epochs) with batch size 64 and sequence length 32. For in-context linear regression, training uses 300,000 iterations with the same batch size and sequence length 13.
-
Cross-validation / statistical protocol. There is no explicit cross-validation for strategy selection (the paper does not optimize hyperparameters over a validation set in the conventional sense — the reverse-engineering analysis is descriptive, not model selection). Statistical reliability is addressed through: (a) multiple random seeds: Table 2 reports polynomial coefficients averaged over 4 seeds with standard deviations; Figure 3.C shows delta loss transitions, and while seed variability is not explicitly quantified there, the sharp transition structure suggests robustness; (b) out-of-distribution generalization for in-context learning: the validation distribution uses double the variance of the training distribution (Section 5.3, Figure 4.B), testing whether the learned algorithm generalizes beyond the training statistics; (c) the polynomial coefficient extraction (Table 2) provides a global functional equivalence test — it compares the learned function to gradient descent over the entire input space, not just on training samples; (d) the pruning analysis (Section 4.1) is lossless — removing zero-weight neurons does not change the network output, so the post-pruning structure is not a statistical estimate but an exact property of the trained network. A limitation is that the number of seeds is small (4 for Table 2) and confidence intervals are not reported for the transition thresholds in Figures 3.A-C — the precise location of the "twice the theoretical minimum" overparameterization threshold may have variance that the paper does not quantify.
Main Quantitative Results
Teacher Identification: Gated RNNs Exactly Match Linear Self-Attention
The headline result of Section 4.1 is that a gated RNN with 100 hidden neurons and 100 output gating neurons, trained to emulate a linear self-attention layer, achieves a training loss of (Table 1), representing near-perfect functional replication of the teacher. After pruning zero-weight neurons, the effective network uses only 14 hidden neurons and 13 output gating neurons — precisely matching the symmetric construction's prediction of recurrent neurons (Figure 2.A-B).
The evidence that the learned solution is the construction, not an alternative function that matches the training outputs, comes from three complementary analyses:
Linear reconstruction scores (Table 1). The key-value matrix can be linearly reconstructed from the 10 neurons with with Score KV = , and the query vector from the 4 neurons with with Score Q = . These scores indicate that the hidden state is an exact linear encoding of the accumulated key-values and the current query, with essentially zero reconstruction error. This establishes that the temporal processing component of the RNN has learned the correct accumulation and pass-through dynamics.
Polynomial coefficient comparison (Table 1). The instantaneous input-output mapping of the RNN can be expressed as a degree-4 polynomial; the teacher's is a degree-3 polynomial. The normalized L2 distance between the coefficient vectors of these two polynomials, averaged over the 4 output dimensions, is Polynomial distance = . This small value means the RNN's degree-4 terms (which would represent computations beyond attention) are essentially zero, and the degree-3 terms match those of the teacher. Because polynomial equivalence holds over the entire input space, this establishes that the RNN has learned the attention function globally, not just on the training distribution.
Weight matrix structure (Figure 2.B). After pruning and re-ordering, the learned weight matrices exhibit the block-sparse structure predicted by the construction. The input gating matrices show a clean separation: rows for key-value neurons () have non-zero entries only in the portions reading value and key projections respectively, while rows for query neurons () read query projections (with using the constant-1 dimension). The output gating matrices have the selector pattern: reads from key-value state entries, reads from query state entries. The readout matrix shows the summation-over-columns structure. Three output gating neurons initially deviate from this pattern but can be merged into a single neuron with the correct structure (Figure 2.C, Appendix C.2), as their combined contribution forms a rank-1 kernel whose singular vectors align with the construction's selector directions.
Effective network size matches compressed construction. The pruned network retains exactly 10 key-value accumulator neurons (not the naive ), confirming that gradient descent discovers the symmetric compression described in Appendix A.3: when is invertible, can be set equal to (adjusting to preserve ), making the key-value matrix symmetric and requiring only entries. The learned values are binary: the 10 key-value neurons have (perfect accumulation), and the 4 query neurons have (no memory), with no intermediate values observed (Figure 2.B).
Overparameterization Threshold: Mild Redundancy Required for Discovery
The paper systematically varies the number of hidden neurons (keeping output gating neurons) and measures the final training loss, producing the transition curves in Figure 3.A. For inputs, where the symmetric construction requires neurons, the loss remains high when and drops to near-zero when , establishing a roughly overparameterization threshold for reliable identification.
What this demonstrates: gradient descent in gated RNNs does not find the attention solution when the network is minimally parameterized. The loss landscape appears to require extra capacity — likely because the optimization path involves exploring configurations where many neurons are active before pruning to the sparse construction solution. The threshold is an empirical regularity that the paper demonstrates but does not theoretically explain; it may depend on optimization hyperparameters (learning rate, initialization, number of iterations) and the specific structure of the teacher attention matrices.
Low-Rank Teacher Compression: Learned Networks Exploit Algebraic Structure
For a teacher attention layer with and constrained to rank 6 (by removing 6 singular values from randomly initialized full-rank matrices), the low-rank construction (Appendix A.4) predicts that only neurons are needed (where and , with almost surely given random sampling). The experiments (Figure 3.B) show a loss transition at approximately 84 hidden neurons — once again roughly the theoretical minimum — confirming that the trained RNN discovers and exploits the low-rank structure rather than naively allocating the symmetric full-rank requirement of neurons.
The fact that the transition occurs near 84 rather than 90 (the symmetric full-rank minimum) indicates that the network is indeed using the low-rank compression, not the symmetric compression. If the network were using the symmetric full-rank construction, it would need at least 90 neurons and would show a transition around 180 neurons under the overparameterization rule. Instead, the transition at approximately 84 () precisely matches the low-rank prediction. This demonstrates that gradient descent discovers the most parameter-efficient representation consistent with the algebraic structure of the target function.
Architectural Comparison on Teacher-Student Task: Multiplicative Interactions Are Essential
Figure 4.A compares seven architectures on the teacher-student task with inputs, using MSE loss on the test set. The headline findings are:
-
Simplified gated RNN (construction baseline): Achieves the lowest loss, serving as the upper bound for what is achievable with this model class. The diagonal and dense recurrence variants perform similarly (curves largely overlap in the figure), indicating that the diagonal inductive bias is helpful but not strictly necessary — the key enabler is the multiplicative gating, not the recurrence structure.
-
LSTM and GRU (1 layer each): Achieve substantially higher loss and fail to match the teacher. Despite the LSTM's theoretical capacity to implement attention in one layer (Table 3, Appendix B.3), gradient descent does not find this solution under standard training conditions. The GRU cannot implement attention even in principle (Appendix B.4), and its loss is similar to the LSTM's, suggesting that neither architecture converges to an attention-like solution.
-
Standard LRU (output GLU only, 2 layers): Achieves moderate loss — better than LSTM/GRU but substantially worse than the construction baseline. This is consistent with the theoretical requirement of two stacked layers: the first layer's output GLU serves as the second layer's input gating, enabling the computation but with additional complexity that makes the solution harder to find.
-
LRU In+Out (input and output GLU, 1 layer): Significantly outperforms the standard LRU, confirming that adding explicit input gating strengthens the inductive bias toward attention. A single layer now contains both multiplicative primitives needed (outer product computation via input gating, matrix-vector multiplication via output gating), making the attention solution more accessible to gradient descent.
-
LRU In+Out with GLU replaced by MLP (1 layer): Dramatically worse performance — the loss is closer to LSTM/GRU levels than to the GLU variants. This is a clean ablation demonstrating that multiplicative interactions are fundamental: a feedforward network with equal parameter count but without explicit elementwise multiplication cannot efficiently compute the outer products and matrix-vector products that attention requires. The MLP has universal approximation capacity in principle, but gradient descent does not find an attention-like solution within it.
The gap between the construction baseline and the LRU In+Out (which both have input and output gating) can be attributed to the LRU's additional complexities: complex-valued recurrence (only the real part is used), sigmoid nonlinearities in the GLU (requiring linearization via small weights), and the specific exponential parameterization of and . These elements add optimization challenges that the simplified architecture avoids.
In-Context Linear Regression: RNNs Converge to Gradient Descent
When gated RNNs (with 80 hidden neurons by default) are trained on the in-context linear regression task with observations and (input dimension ), the training loss converges to 0.0945, which is nearly identical to the loss of one optimal gradient descent step (0.0947; delta loss ). This near-equality, visible in Figure 3.C at the right side of the transition curve, is the first piece of evidence that the RNN has learned gradient descent rather than an alternative algorithm.
The polynomial coefficient analysis (Table 2) provides the mechanistic confirmation. The instantaneous function of the first output neuron is decomposed into coefficients of all degree-4 and lower polynomial terms. Only three terms have non-negligible coefficients:
- :
- :
- :
The optimal gradient descent step predicts coefficients of for exactly these three terms, where is the analytically derived optimal learning rate and . The empirical coefficients () match the theoretical value () to within approximately . The residual norm — the norm of all other polynomial coefficients — is , two orders of magnitude smaller than the dominant coefficients, confirming that the network implements exactly gradient descent and nothing else.
What this means structurally: The gradient descent computation for this task is a specific low-rank instance of linear self-attention. The value matrix extracts (the output portion of the concatenated input), the key matrix extracts (the input portion), and the query matrix extracts . The accumulated key-value matrix has rank at most . The compressed construction would require key-value neurons and 3 query neurons. With overparameterization, the transition to zero delta loss should occur at approximately 30 hidden neurons. The paper's systematic variation of hidden neuron count (Figure 3.C) shows a loss transition at approximately 24 neurons — consistent with the compressed construction given the approximate nature of the threshold.
The out-of-distribution generalization (validation variance doubled to ) is used in the architectural comparison (Figure 4.B) to distinguish between algorithms that truly learn gradient descent (which should generalize to any with finite variance) and those that overfit to the training distribution statistics. The gated RNNs, LRU variants, LSTMs, and GRUs all generalize to the doubled-variance setting (Figure 4.B shows validation loss, not training loss), though the paper does not separately report in-distribution vs. out-of-distribution loss for each architecture.
Architectural Comparison on In-Context Learning: Nonlinearity Enables Superior Algorithms
Figure 4.B presents the test loss (with variance doubled relative to training) for different architectures on the in-context linear regression task. The optimal gradient descent baseline is shown as a horizontal reference line. The findings reveal a stark contrast with the teacher-student results:
-
Gated RNNs and LRU variants achieve loss close to but slightly above the gradient descent optimum when given moderate parameters, consistent with them learning the gradient descent algorithm (as established by Table 2 for the gated RNN). The LRU In+Out achieves lower loss than the standard LRU, and replacing the GLU with an MLP causes a large drop — the same pattern as the teacher-student task, consistent with the attention-mimicking ability correlating with in-context learning performance within the linear/low-nonlinearity architectural family.
-
LSTMs, and particularly GRUs, achieve substantially lower loss than one gradient descent step when given sufficient parameters. The GRU achieves the best overall performance on this task. This is surprising in light of the teacher-student results: these architectures could not learn to implement a specific attention layer (Figure 4.A), yet they outperform the attention-based gradient descent algorithm on in-context learning (Figure 4.B).
Interpretation: LSTMs and GRUs, with their strong sigmoid/tanh nonlinearities, discover algorithms for in-context linear regression that are more sophisticated than one-step gradient descent — possibly implementing multi-step optimization, learning-rate adaptation, or leveraging the recurrent state for second-order information. Their failure at teacher-student attention mimicking (Figure 4.A) indicates that these superior algorithms are not attention-based. This is a crucial dissociation: attention is sufficient but not necessary for in-context learning in RNNs, and strong nonlinearities enable algorithms that can outperform attention on this task.
The correlation within the LRU family (between attention-mimicking ability in Figure 4.A and in-context learning performance in Figure 4.B) is high: architectures that can implement attention (LRU In+Out) perform better than those that cannot (standard LRU, MLP variant). This suggests that for architectures with weak nonlinearities (where sigmoid can be approximately linearized), the primary path to in-context learning is through implementing attention. For strongly nonlinear architectures, alternative paths open up that can be more effective.
Associative Recall: Different Algorithms for Different Gating Mechanisms
Appendix D.3 provides an illustrative case study on a simpler in-context task: associative recall, where the network must memorize pairs and retrieve when queried with . Two variants of the gated RNN are trained:
-
Input-output gating (the standard model, Equation 2): The trained network discovers a solution that does not match the linear self-attention construction. Instead, it allocates one neuron per possible value, stores a value associated with in that neuron when the pair appears, and uses a neuron to hold the current . The output gating then computes the negative squared difference between the stored and current values, with the zero-difference position (via argmax) indicating the correct . This solution exploits the one-hot encoding of inputs, which is specific to this task formulation.
-
Side gating (Appendix A.2): The trained network discovers a solution that is an instance of the attention construction, but with binary patterns in the matrices that differ from the canonical outer product. Specifically, the learned matrices have rank-1 components with an additional negative peak of equal amplitude (Figure 7), which does not affect classification accuracy given one-hot encodings but reveals that the network compresses the representation using neurons rather than the full that naive attention would require.
This comparison demonstrates that architectural differences (input-output gating vs. side gating) lead to different learned algorithms even on the same task, and that the specific algorithm discovered depends on available computational primitives and task structure (one-hot encodings enabling compression). It also confirms that the gated RNN can discover attention-like solutions (with side gating) or alternative solutions (with input-output gating) depending on which is more parameter-efficient for the specific task encoding.
Ablation Studies and Robustness Checks
Multiplicative gating vs. MLP (Figures 4.A and 4.B): Replacing the GLU in the LRU In+Out architecture with a 1-hidden-layer MLP of equal parameter count causes a dramatic performance degradation on both the teacher-student task (Figure 4.A: loss increases to LSTM/GRU levels) and the in-context learning task (Figure 4.B: substantial loss increase, falling well below the LRU In+Out and even below the standard LRU). This establishes that elementwise multiplication is the essential operation, not merely additional nonlinearity or increased representational capacity. An MLP with equivalent parameters has universal approximation power but lacks the architectural inductive bias for computing outer products and matrix-vector products efficiently, and gradient descent fails to discover alternative implementations that match attention.
Diagonal vs. dense recurrence (Figure 4.A): The diagonal and dense recurrence variants of the simplified gated RNN achieve nearly identical loss on the teacher-student task (curves overlap in the figure, though exact values are not tabulated). This indicates that the diagonal inductive bias — while useful for stable training and long-range memory (Orvieto et al., 2023b; Zucchet et al., 2023b) — is not essential for implementing attention. The critical enablers are the multiplicative gating mechanisms; the recurrence structure (diagonal vs. dense) primarily affects optimization dynamics rather than representational capacity for this computation.
Input gating addition to LRU (Figure 4.A, comparison of "LRU" vs. "LRU In+Out"): Adding an explicit input GLU to the standard LRU architecture (which has only output GLU) significantly improves attention-mimicking ability, reducing test loss substantially. This validates the theoretical analysis that two stacked standard LRU layers are needed to implement the construction (first layer's output GLU as second layer's input gating), whereas a single LRU In+Out layer provides both primitives directly, making the attention solution more accessible to gradient descent. The same pattern holds for in-context learning (Figure 4.B): LRU In+Out outperforms standard LRU.
Number of stacked layers for LRU (Table 3, theoretical): The paper provides a theoretical tally of how many layers each architecture needs to implement the simplified gated RNN class and linear self-attention: standard LRU requires 2 layers, LRU In+Out requires 1, LSTM requires 1 (theoretically), GRU cannot implement either with any number of layers. These theoretical layer requirements are validated empirically for LRU (the 2-layer standard LRU achieves lower loss than 1-layer would, though the 1-layer baseline is not shown separately) and implicitly for LSTMs (which have the theoretical capacity but don't achieve low loss in practice).
Overparameterization scaling (Figures 3.A-C): The systematic variation of hidden neuron count across three settings establishes a consistent empirical regularity: gradient descent requires approximately twice the theoretical minimum number of neurons to reliably discover the attention (or gradient descent) solution. This is observed for the full-rank teacher with (threshold at vs. minimum 14, Figure 3.A), the low-rank teacher with and rank 6 (threshold at vs. minimum 42, Figure 3.B), and the in-context linear regression task with (threshold at vs. minimum , Figure 3.C). The consistency of the factor across different , different ranks, and different tasks (teacher-student vs. in-context) suggests a general property of gradient descent dynamics in this architecture class, though the paper provides no theoretical explanation for the factor.
Output gating neuron merging (Figure 2.C, Appendix C.2): Three output gating neurons in the trained network initially exhibit weight patterns that deviate from the construction's block-sparse selector structure. The paper demonstrates that these three neurons can be functionally merged into a single neuron with the construction structure, without changing the network's input-output function. The merged neuron's kernel is rank-1 and its singular vectors align with the key-value and query reading directions, exactly matching the construction's selector pattern. This serves as a robustness check that apparent deviations from the construction in learned weights are often due to the invariances of the gating parameterization (permutations, scalar rescalings, decomposition of quadratic forms into different numbers of rank-1 components) rather than genuine functional differences. The fact that merging recovers the construction structure supports the claim that the network has functionally learned attention, even if the raw weights look different due to these invariances.
Lossless pruning (Figure 2.A, Section 4.1): After training the network, 86 out of 100 hidden neurons and 87 out of 100 output gating neurons have input or output weights that are entirely zero (or near-zero). Pruning them is lossless — the pruned network computes exactly the same function. This demonstrates that the network was substantially overparameterized relative to the solution found, and that gradient descent implicitly performs neuron-level sparsification (pushing unused neurons' weights to exactly zero). The pruned architecture size (14 hidden, 13 output gating) matches the theoretical minimum for the symmetric construction, confirming that the effective computation uses the minimal resources needed.
Out-of-distribution generalization for in-context learning (Figure 4.B): The in-context learning experiments use a validation set where entries are drawn from — double the training variance of . The fact that all architectures show meaningful performance on this distribution (rather than complete collapse) indicates that they have learned algorithms for linear regression that generalize to different weight distributions, not memorization of specific training tasks. The paper does not separately report in-distribution vs. out-of-distribution loss for each architecture, making it difficult to quantify the generalization gap; this would have been informative for distinguishing architectures that truly learn the algorithm vs. those that partially overfit.
LSTM negative result (Figure 4.A): The LSTM fails to achieve low loss on the teacher-student task despite having the theoretical capacity to implement attention in one layer (Appendix B.3). This negative result is itself an important ablation: it shows that architectural capacity does not imply learnability, and that the nonlinearities (sigmoid, tanh) in practical RNNs create optimization barriers that prevent gradient descent from finding theoretically possible solutions. The paper does not investigate whether modified training procedures (different initializations, learning rate schedules, or curriculum learning) could enable LSTMs to discover the attention solution, leaving open the question of whether the failure is fundamental or can be overcome with better optimization.
GRU impossibility (Appendix B.4, theoretical): The GRU cannot implement linear self-attention with any number of stacked layers because it lacks a mechanism to compute multiplicative outer products of different input projections. This is a theoretical ablation rather than an empirical one: it identifies exactly which architectural element (multiplicative gating of the form with different linear projections) is the necessary condition for attention implementation. The GRU's reset and update gates modulate information flow but do not create products of input features. This sharpens the paper's central claim: not all gating enables attention — only multiplicative gating of multiple input projections.
Side gating vs. input-output gating (Appendix A.2 and D.3): The associative recall experiment (Appendix D.3) compares the standard input-output gating model with a side gating variant (), which requires rather than parameters and provides a stronger inductive bias toward attention (matching Mamba's architecture). The side gating variant discovers an attention-like solution, while the input-output gating variant discovers a non-attention algorithm that exploits the one-hot encoding. This ablation demonstrates that different gating configurations bias the network toward different algorithmic strategies, even when both have the theoretical capacity to implement attention.
Critical Assessment
Does the paper demonstrate that gated RNNs can exactly implement linear self-attention?
Yes, this is the most solidly supported claim in the paper. The constructive proof (Section 3, Appendix A.1) provides explicit weight values, the construction requires a finite number of neurons (not asymptotic), and the teacher-student experiments (Section 4.1) demonstrate that trained networks achieve loss of with weight structures matching the construction (Figure 2.B, Table 1). The polynomial distance of establishes functional equivalence across the entire input domain, not just the training distribution.
However, there are important scope limitations that the paper acknowledges but that constrain the strength of this claim:
-
Only linear self-attention, not softmax attention: The construction maps to linearized attention (Equation 1), which removes the softmax nonlinearity. The paper does not claim that gated RNNs can implement standard softmax attention, and the construction does not easily extend to it because softmax requires computing exponentials and normalizing over all positions — operations that are not naturally expressed as elementwise products of linear projections. The claimed equivalence is therefore to a specific, simplified attention variant, not to the attention mechanism used in Transformers.
-
Causally-masked, single-head attention: The construction covers one attention head with causal masking. Multi-head attention would require parallel copies (increasing the state size proportionally) or interleaving. Cross-attention between different sequences is not considered.
-
Square weight matrices: The paper assumes for simplicity. The claims about neuron count ( for full-rank, for symmetric) depend on this assumption. For rectangular matrices (different key/query/value dimensions), the counts would generalize but the specific numbers would change.
-
The construction is parameter-inefficient: The paper is explicit that the parameter requirement is a weakness, not a strength. The claim is about representational capacity, not practical efficiency. No one would use an RNN with parameters to replace a -parameter attention layer; the construction's value is conceptual — it reveals the computational primitives needed and explains when/why RNNs might succeed at attention-like tasks.
Does the paper demonstrate that gradient descent discovers the construction?
Partially, with important caveats. The evidence is strong for the simplified gated RNN with diagonal linear recurrence and linear gating (Equation 2): the teacher-student experiments show clean convergence to the construction with the predicted weight structure, binary values, and exact functional equivalence. The in-context learning experiments (Section 5) additionally show convergence to the gradient descent algorithm, which is a specific instance of the low-rank construction.
However, the generality of this convergence claim is more limited than the paper's narrative might suggest:
-
Extremely high overparameterization in the identification experiment: The teacher-student experiment uses 100 hidden neurons when the compressed construction needs only 14 — over 7× overparameterization. The paper's own analysis (Figure 3.A) shows that at least 2× overparameterization is required for convergence, and the main experiment uses substantially more. This raises the question of whether convergence would occur with more modest overparameterization (say, 20-25 neurons) and whether the clean binary structure depends on having abundant spare capacity.
-
Only tested on small input dimensions: for the full identification experiment, for the architectural comparison, for the low-rank experiment. Real-world attention layers have in the hundreds or thousands. The parameter scaling means the experiments operate in a regime far from practical sizes, and it is unknown whether the convergence behavior scales — optimization difficulties may compound at larger , or the overparameterization factor might grow with dimension.
-
Single initialization and optimization setup: The paper uses one optimization protocol (AdamW with cosine annealing, specific learning rate, weight decay excluding ). The dependence of convergence on these choices is not ablated. Different optimizers, learning rate schedules, or initialization schemes might change the overparameterization threshold or prevent convergence entirely.
-
LSTM negative result highlights the fragility: The LSTM, which theoretically can implement attention in one layer, fails to do so under the same training protocol (Figure 4.A). This means the "gradient descent discovers the construction" claim is highly architecture-dependent. It holds for the specific simplified model class (linear gating, diagonal recurrence, binary ) but not for architectures with similar theoretical capacity but different nonlinearities (LSTM, and implicitly any architecture with sigmoid/tanh gates rather than linearized GLUs). The paper's narrative sometimes elides this distinction, speaking broadly about "gated RNNs" when the positive results are specific to the simplified linear-gated class.
Does the paper demonstrate that trained RNNs discover attention on in-context learning tasks?
Yes for the specific case of in-context linear regression, with the gated linear RNN architecture. The evidence in Table 2 (polynomial coefficients matching gradient descent to within 0.7%) and the delta loss near zero (0.0945 vs. 0.0947) are strong. The convergence to the low-rank construction size (Figure 3.C) further supports that the mechanism is the attention-based gradient descent algorithm.
However, the paper's own architectural comparison (Figure 4.B) demonstrates that nonlinear RNNs outperform this attention-based algorithm. The GRU achieves the best in-context learning performance, despite being theoretically incapable of implementing attention. This complicates the narrative that RNNs succeed at in-context learning "because they discover attention." For the gated linear RNN and LRU family, the attention mechanism does appear to be the primary path to in-context learning (as evidenced by the correlation between Figure 4.A and 4.B within that family). For LSTMs and GRUs, other mechanisms — possibly involving the nonlinear gating dynamics — achieve even better performance.
The paper does not investigate what algorithms the LSTMs and GRUs are actually implementing. The reverse-engineering analysis (polynomial coefficient extraction, linear reconstruction scores, weight structure visualization) is applied only to the simplified gated RNN, not to the LSTM or GRU. Understanding what superior algorithms the GRU discovers would significantly strengthen the paper's contribution and is a notable gap.
Additionally, the in-context experiments use only one task (linear regression) with specific dimensionalities (, ). The associative recall experiment (Appendix D.3) uses a different task but only analyzes the simplified gated RNN, not the LSTM/GRU. The claim that RNNs "discover attention-based in-context learning algorithms" would be stronger with evidence across a broader range of in-context learning tasks (e.g., nonlinear function classes, classification, sequence completion) and with mechanistic analysis of what the LSTM/GRU actually learn.
Does the paper validate that the identified computational primitives (multiplicative gating, binary ) are necessary?
Yes, through a combination of theoretical argument and empirical ablation. The GRU impossibility proof (Appendix B.4) establishes theoretically that architectures without multiplicative gating of multiple input projections cannot implement attention. The MLP ablation (Figure 4.A: replacing GLU with MLP destroys performance) provides the empirical complement: even a universal function approximator with equal parameters fails if it lacks elementwise multiplication. The LRU In+Out vs. standard LRU comparison (Figure 4.A) shows that adding input gating (the second multiplicative primitive) significantly improves attention-mimicking. The value analysis (Figure 2.B) shows that trained networks use exactly binary , consistent with the construction's requirement for perfect accumulation and zero memory — no intermediate decay rates are used.
A missing experiment that would strengthen the necessity claim: an ablation where is forced to be intermediate (e.g., constrained to rather than ) to test whether perfect accumulation () is truly necessary or whether exponential moving averages with long timescales could approximate attention adequately. The paper's construction requires exactly, but in practice, a neuron might approximate the behavior well enough for the loss to be low, while being reachable by gradient descent more easily.
Does the paper demonstrate that the connection between RNNs and attention has practical implications for architecture design?
Yes, but the evidence is primarily through controlled synthetic experiments, not real-world tasks. The paper's clearest practical contribution is the architectural diagnostic framework: given an RNN architecture, one can check whether it has the multiplicative primitives needed for attention and predict its capacity for attention-like computations. The paper validates this diagnostically for several architectures (Table 3, Figure 4.A): LRU In+Out has the primitives and performs well, standard LRU requires stacking and performs moderately, LSTM has them theoretically but fails empirically, GRU lacks them and cannot implement attention, MLP variants lack them and fail.
The step from this diagnostic framework to guiding architecture development (as claimed in Section 6 and Appendix B.5) is plausible but not directly demonstrated. The paper does not, for instance, propose a novel architecture based on the construction and show it outperforms existing ones on a practical benchmark. The "LRU In+Out" variant is introduced and shown to improve attention-mimicking, but this is a minor modification to an existing architecture tested only on synthetic tasks. The paper's suggestion that the Mamba architecture's success is "partly explained" by its attention-implementing capacity (Section 3.3) is post-hoc and not experimentally validated — the paper does not reverse-engineer trained Mamba models to show they actually implement attention on real tasks.
Missing experiments and analyses:
-
Reverse-engineering of LSTM/GRU solutions on in-context learning. The finding that GRUs outperform gradient descent on in-context linear regression is striking and suggestive, but the paper does not analyze what algorithm the GRU actually implements. Understanding this would significantly advance the field's understanding of in-context learning mechanisms beyond the attention=gradient-descent framework.
-
Larger-scale experiments. All experiments use tiny dimensions by modern standards (, state sizes ). While the constructive theory scales, the empirical learnability may not — demonstrating convergence at or would substantially strengthen the practical relevance, even if the experiment requires proportionally larger state sizes.
-
Softmax attention as a teacher. The paper only studies linear attention. An experiment training gated RNNs to emulate standard softmax attention (even at small scales) would test whether the construction's insights extend to the attention variant used in practice.
-
Real-world sequence modeling tasks. The paper's claims about architecture design implications would be bolstered by showing that adding input gating to the LRU (or removing multiplicative gating) impacts performance on standard benchmarks like Long Range Arena or language modeling, not just synthetic teacher-student tasks.
-
Sensitivity to optimization hyperparameters. The overparameterization threshold and the binary convergence might depend on the specific optimizer, learning rate schedule, and weight decay configuration. Ablating these would establish robustness.
-
The cost of learning vs. the cost of the solution. The experiments show that gradient descent finds the compressed construction, but they don't analyze the computational cost of finding it vs. the cost of the final solution. The overparameterization threshold is measured in terms of architecture size at convergence, not in terms of training iterations or wall-clock time to convergence. An architecture that requires more neurons for optimization but prunes to the minimum afterwards is still paying a training cost that a more carefully designed architecture might avoid.
Summary of experimental support for central claims:
The claim that gated RNNs can implement linear self-attention is conclusively supported by the constructive proof and the teacher-student experiments on the simplified architecture. The claim that gradient descent discovers this implementation is supported for the simplified architecture under specific conditions (sufficient overparameterization, appropriate optimization), but the failure on LSTMs and the lack of larger-scale or real-world experiments qualify the generality. The claim that trained RNNs discover attention on in-context learning tasks is supported for the gated linear RNN on linear regression, but the superior performance of GRUs (which cannot implement attention) on the same task complicates the narrative. The claim that multiplicative interactions are the essential primitive is strongly supported by the MLP ablation and the GRU impossibility proof. The claim that the framework can guide architecture design is plausible and demonstrated diagnostically but not yet validated through novel architecture proposals or real-world benchmarks.
6. Limitations and Trade-offs
The Quartic Parameter Cost Means This Construction Is Not a Practical Recipe for Replacing Attention
The assumption or constraint. The constructive proof mapping linear self-attention to gated RNNs requires O(d²) hidden neurons and O(d⁴) total parameters for a d-dimensional input, compared to the 3d² parameters of the attention layer being emulated (Section 3.2). The paper states this explicitly:
"Overall, the output gating requires O(d²) input and output entries for the gated RNN to match a linear self-attention layer. The RNN thus requires O(d⁴) parameters in total, with a lot of redundancy, significantly more than the 3d² parameters of the linear self-attention layer."
Even with the side-gating variant (Appendix A.2), the parameter cost remains O(d³) — still a cubic-vs-quadratic gap. The symmetric compression (Appendix A.3) halves the hidden neuron count but leaves the output gating blowup untouched.
The consequence. This cost structure means that for realistic input dimensions (d = 64 for a typical Transformer head, or d = 128–1024 for model dimensions in production LLMs), the RNN would require astronomically many parameters to implement even one linear attention head. For d = 64, the full-rank construction demands 64² + 64 = 4,160 hidden neurons and output gating matrices of size 4,096 × 4,160 — over 17 million parameters in the output gating alone, compared to 3 × 64² = 12,288 for the attention layer. At d = 128, the output gating alone would exceed 268 million parameters. No practitioner would use a million-parameter RNN component to replace a twelve-thousand-parameter attention layer. The construction therefore has zero practical deployment value as a direct replacement for attention. It is a theoretical existence proof, not an engineering recipe.
What evidence exists in the paper. All experiments operate at tiny scales: d = 4 for the main identification experiment (Section 4.1), d = 6 for the architecture comparison (Section 4.3), d = 12 for the low-rank experiment (Figure 3.B). Even at d = 12 with rank-6 compression, the trained network requires approximately 84 hidden neurons (Figure 3.B) and the output gating matrices have 42 × 84 dimensions — a regime where the quartic scaling has not yet manifested. The paper provides no experiments or extrapolations showing whether the construction scales to d = 64 or beyond, nor does it test whether gradient descent can find the solution when the parameter count becomes enormous. The paper acknowledges this is a limitation when it frames the result as partly a "negative result":
"Our results partly serve as a negative result: implementation of attention is possible but requires squaring the number of parameters attention has."
Mitigation status. The paper identifies two compression strategies (symmetry in Appendix A.3, low-rank structure in Appendix A.4) that reduce the hidden neuron count, but neither addresses the output gating bottleneck. The symmetry argument halves key-value neurons but output gating still requires O(d²) neurons and O(d⁴) weights. The low-rank compression reduces the effective dimension to the ranks r_V and r_KQ, making the parameter cost O(r_V · r_KQ · d) in the side-gating variant — but this still involves the full input dimension d, not just the ranks. The paper shows trained networks discover these compressions (Figures 3.B, 3.C), validating the theoretical analysis, but does not propose architectural modifications to reduce the output gating cost further. The side-gating variant (Appendix A.2) reduces output gating to O(d³) but this remains cubic, and the variant is architecturally different from the input-output gating model used in most experiments. The paper does not address whether training becomes practically infeasible at larger scales, whether the 2× overparameterization threshold grows with d, or whether optimization dynamics change when quartic parameter counts interact with limited training budgets.
Linear Self-Attention Only — No Extension to Softmax Attention, Multi-Head Attention, or Cross-Attention
The assumption or constraint. The entire theoretical construction and all experiments target causally-masked linear self-attention specifically (Equation 1), which removes the softmax nonlinearity:
This is explicitly stated in Section 2.1: "Such linear self-attention layers can be understood as a linearized version of the softmax attention mechanism." The paper studies only single-head attention with square weight matrices, does not consider cross-attention between different sequences, and does not address other Transformer components (feedforward layers, layer normalization, residual connections).
The consequence. Softmax attention — the mechanism actually used in Transformers — differs from linear attention in a computationally fundamental way: it computes a weighted sum of values where the weights are normalized exponentials of key-query inner products, not raw products. This normalization (the softmax) creates a nonlinear, non-separable interaction among all positions that cannot be expressed as a simple outer product accumulation in a fixed-size recurrent state. Specifically, the softmax attention state update depends on all past keys to compute the normalizing denominator ∑_{t'} exp(k_{t'}^\top q_t), which requires either storing all past keys (unbounded memory) or recomputing them — neither of which is compatible with the constant-memory recurrence studied in this paper. The paper's construction therefore does not demonstrate that RNNs can implement the attention mechanism used in all production Transformers. A practitioner cannot use this construction to understand whether an LSTM-based language model is secretly doing something analogous to GPT's attention — it only applies to the linearized variant found in efficient Transformer architectures (Performer, Linear Transformer, etc.), which trade expressivity for efficiency and generally underperform softmax attention on quality benchmarks.
Furthermore, multi-head attention — another universal feature of practical Transformers — is not addressed. A Transformer typically uses 8–64 parallel attention heads. Implementing one head in an RNN using the paper's construction costs O(d⁴) parameters; k heads would naively require k-fold replication of this cost, making the gap even wider.
What evidence exists in the paper. The paper does not provide any experiments with softmax attention teachers, nor does it discuss whether the construction could be extended to approximate softmax attention (e.g., via kernel approximations like random Fourier features, or via learned normalizations). The architectural comparisons (Section 4.3) use linear self-attention teachers exclusively. The in-context learning experiments (Section 5) use linear self-attention as the comparison mechanism because gradient descent can be implemented via linear attention (von Oswald et al., 2023), not because linear attention is the optimal mechanism for the task. The paper does not discuss cross-attention (where keys and values come from one sequence and queries from another), which is central to encoder-decoder Transformers.
Mitigation status. The paper is transparent about the scope: Section 2.1 defines the linear self-attention layer explicitly, and the title states "Gated recurrent neural networks discover attention" without qualification. However, Section 1 frames the contribution in terms of "attention-based neural networks, most notably Transformers," creating an impression of broader applicability that the technical scope does not support. The paper does not propose extensions or future work toward softmax attention or multi-head attention. The discussion (Section 6) references "bridging the gap between Transformers' computational power and RNNs' inference efficiency" as a motivation but does not suggest how to bridge the softmax-linear gap specifically. A reader unfamiliar with the distinction between linear and softmax attention may overinterpret the results as showing that RNNs can implement Transformer attention, when the paper only addresses a simplified variant.
The Difficulty Estimation Cost (Overparameterization Requirement) Is Not Accounted for in Any Practical Efficiency Claim
The assumption or constraint. The paper demonstrates that gradient descent in gated RNNs discovers the attention construction only when the network is sufficiently overparameterized — approximately 2× the theoretical minimum number of neurons (Section 4.2, Figure 3). This threshold is observed consistently across the full-rank teacher experiment (Figure 3.A: ~28 neurons needed vs. 14 minimum for d = 4), the low-rank teacher experiment (Figure 3.B: ~84 vs. 42 for d = 12, rank 6), and the in-context learning experiment (Figure 3.C: ~24 vs. ~15 for d = 6). Below this threshold, the loss remains high — the network fails to learn attention at all.
The consequence. The 2× overparameterization means the RNN must be provisioned with at least twice as many neurons as the construction actually needs, and training must discover the sparse solution within this redundant capacity. This has several practical implications that the paper does not quantify:
Training cost is not amortized in any efficiency claim. The paper reports effective network sizes after pruning (e.g., 14 hidden neurons for d = 4 in Section 4.1) and the 4× efficiency metaphor from the reference summary does not apply here, but the training compute required to train at 100 neurons (the experiment's setting) to converge to a 14-neuron solution is not compared to the cost of directly training a 14-neuron network. Since a 14-neuron network does not converge (Figure 3.A), the 2× overparameterization is not optional — it is a necessary condition for the optimization to succeed. The cost of training the larger network (more parameters, more memory, more gradient computation) is the price of reaching the sparse solution.
The 2× factor may grow with problem dimension. The paper tests only three (d, rank) configurations: (4, full), (12, rank-6), and (6, rank-3). Within this limited sample, the factor is consistently ~2. However, it is unknown whether the factor remains constant as d increases, as rank structure changes, or as the teacher function becomes more complex. If the factor grows (e.g., to 4× or 10×), the overparameterization cost would make the approach infeasible at scale even for tasks where the compressed construction is theoretically compact.
The cost of determining the "right" amount of overparameterization is not addressed. In practice, the minimum construction size is unknown a priori — it depends on the rank structure of the (unknown) target function. A practitioner would need to sweep hidden neuron counts, train multiple networks, and observe the loss transition to determine when sufficient capacity has been reached. This hyperparameter search multiplies the already-significant training cost.
What evidence exists in the paper. The loss transition curves in Figures 3.A-C are presented to demonstrate that networks discover compressed solutions, not to analyze the cost of overparameterization. The paper does not report training wall-clock time, FLOP counts, or memory usage for any experiment. The 100-neuron main experiment (Section 4.1) uses >7× overparameterization relative to the 14-neuron solution — far more than the 2× threshold — and the paper does not discuss whether this extreme overparameterization was necessary for convergence or was merely a conservative choice.
Mitigation status. The paper does not address the training cost of overparameterization as a limitation. The focus is on the representational capacity (how few neurons suffice to implement attention once learned) rather than on the optimization cost of reaching that representation. Section 8 acknowledges that "mild overparameterization" is needed for identification but does not quantify or problematize this cost. No experiments vary overparameterization systematically above the 2× threshold to determine whether additional capacity accelerates or decelerates convergence, or whether the final solution quality depends on initial overparameterization amount. This is a significant gap for any practitioner considering whether the construction is learnable in practice.
Single Architecture, Single Task Domain, Tiny Scale — No Evidence of Generality to Practical Settings
The assumption or constraint. All reverse-engineering experiments — where the paper demonstrates that trained networks actually implement the construction — use a single simplified architecture (diagonal linear gated RNN with linear GLU-style gating, Equation 2), on synthetic tasks (teacher-student with random attention weights, in-context linear regression, associative recall), at tiny scales (d ≤ 12, state sizes ≤ 100). The paper does not train or reverse-engineer on any real-world dataset (language modeling, machine translation, reinforcement learning) or at any scale approaching practical deep learning (d of hundreds, state sizes of thousands).
The paper acknowledges this scope explicitly in Section 4: "We now demonstrate that gated RNNs learn to mimic attention and comprehend how they do so. In this section, a student RNN is tasked to reproduce the output of a linear self-attention layer." The "in this section" scoping is clear, but the implications for generality are not discussed.
The consequence. Several critical uncertainties follow from this narrow scope:
Other architectures may not exhibit the same behavior, and the paper provides mixed evidence on this. The architecture comparison in Figure 4.A shows that LSTMs and GRUs fail to learn the attention construction on the teacher-student task. Even the standard LRU, which is architecturally closer to the simplified model, performs substantially worse than the construction baseline. This means the positive results (gradient descent discovers the construction) are specific to the simplified linear-gated architecture, not to "gated RNNs" in general. Yet the paper's title, abstract, and discussion use the broader term, potentially misleading readers about the generality to practical LSTM/GRU/LRU/Mamba implementations.
The learned algorithms on real tasks may differ. The paper shows that on a simple synthetic task (in-context linear regression), the gated linear RNN converges to gradient descent. But on a simpler task (associative recall), the same architecture discovers a non-attention solution when using input-output gating (Appendix D.3) — the network exploits one-hot encoding structure instead of implementing the construction. This suggests that the attention solution is one of several that gradient descent might find, and which one wins depends on task structure, data encoding, and the interaction between architecture and data. On real tasks (language modeling, where inputs are not one-hot and relationships are complex), there is no evidence about whether the network would discover attention, some other mechanism, or a hybrid.
Scale matters in ways the paper cannot assess. At d = 4, the construction requires 14–16 neurons and 2× overparameterization means ~30 neurons. At d = 128, the full-rank construction requires 8,256 neurons and ~16,500 with 2× overparameterization. Gradient descent dynamics in a 16,500-neuron recurrent layer with quartic output gating parameter counts are not tested. The loss landscape may change qualitatively at larger scales — the paper's observed binary λ convergence, clean block-sparse weight structure, and exact functional equivalence may be small-scale phenomena that degrade with dimension. Optimization challenges (vanishing/exploding gradients in the output gating, poor conditioning of the quartic parameter space) may prevent convergence at practical scales.
The "representative model" claim is unverified. The paper states in Section 4 that PaLM 2-S* is "representative of the capabilities of many contemporary LLMs" — this is about the reference paper's model, not the current paper, but the current paper uses gated RNNs with 100 neurons and claims broader relevance. The paper does not test whether its representative model (gated diagonal linear RNN) actually matches the behavior of practical large-scale RNNs like Mamba or RWKV when trained on real data. The reverse-engineering methodology (pruning, λ analysis, polynomial coefficient extraction) is feasible at d = 4 with 100 neurons but would be computationally prohibitive at d = 128 with 16,500 neurons, making it impossible to verify the construction on large-scale trained models even if resources were available.
What evidence exists in the paper. The evidence for generality is limited to the architecture comparison in Figure 4.A (showing LSTM and GRU failure) and the associative recall experiment in Appendix D.3 (showing task-dependent algorithm selection). These results cut against the generality of the claim: they demonstrate that the positive findings are architecture-specific and task-dependent. The paper does not provide any experiments on real data, at practical scales, or with deployed RNN architectures (Mamba, RWKV). The correlation between attention-mimicking ability and in-context learning performance within the LRU family (Figure 4.A vs. 4.B) is suggestive, but it comes from synthetic tasks at d = 6 and does not establish that this correlation holds on real benchmarks.
Mitigation status. The paper does not claim to have demonstrated generality to practical settings. The experiments are explicitly scoped as synthetic, small-scale, and architecture-controlled. The discussion (Section 6) acknowledges that "understanding whether real-world attention mechanisms lie in this regime remains an open question" regarding compression. However, the paper's framing — particularly the title "Gated recurrent neural networks discover attention" and the abstract's "suggest that certain RNNs might be unexpectedly implementing attention under the hood" — implies broader applicability than the experiments support. The phrase "under the hood" suggests a finding about what happens inside RNNs in general, when the evidence is specific to one simplified architecture on synthetic tasks. The paper suggests future work on "further investigations on the differences between RNNs and Transformers in controlled settings" (Section 6) but does not propose experiments that would test generality to practical RNNs or real tasks.
The LSTM Learns Nothing — Theoretical Capacity vs. Practical Learnability Gap Is Unexplained
The assumption or constraint. The paper proves that LSTMs can theoretically implement linear self-attention in a single layer (Appendix B.3, Table 3, Table 4), by configuring the input gate and cell state update to compute key-value outer products and the output gate to multiply by queries. However, when actually trained on the teacher-student task, LSTMs achieve high loss and fail to match the teacher — performing similarly to GRUs, which cannot implement attention even in principle (Figure 4.A).
The paper acknowledges this in Section 3.3: "we find somewhat surprisingly that our results extend only to LSTMs... the LSTM construction we provide requires a very specific configuration, which hints that the inductive bias towards attention-compatible configurations might be weaker for this architecture than for deep gated linear RNNs." And again in Section 4.3: "gradient descent does not find such a solution and the performance of LSTMs is close to that of GRUs that cannot implement attention."
The consequence. This negative result has profound implications that the paper does not fully explore. It demonstrates a capacity-learnability gap: an architecture can have the representational capacity to implement a function, yet gradient descent under standard training conditions cannot find the corresponding parameters. This gap is not a minor caveat — it affects the LSTM, the most historically important and widely-deployed RNN architecture. If a practitioner reads the paper's title and abstract and concludes "RNNs can implement attention," they might reasonably assume this applies to LSTMs (since LSTMs are the canonical gated RNN). The experiments show it does not — or at least, gradient descent cannot find it.
The reasons for this failure are not diagnosed by the paper, leaving several critical questions unanswered:
-
Is the failure fundamental or optimization-related? The LSTM construction (Appendix B.3) requires linearizing the tanh and sigmoid nonlinearities by using small weights in specific matrices and large compensating weights elsewhere. This creates a narrow parameter regime — perhaps too narrow for gradient descent to locate from standard random initializations. If alternative initializations, learning rate schedules, or curriculum strategies could enable LSTMs to find the solution, then the negative result is an optimization challenge, not a fundamental limitation. If not, then the LSTM's architecture has a representational capacity for attention that is practically inaccessible, which would be a significant finding about architectural inductive biases.
-
What do LSTMs learn instead? The LSTM achieves low-but-not-zero loss on the teacher-student task — it is not randomly guessing. What function is it implementing? Is it approximating attention poorly? Learning a different computation that partially matches the teacher's outputs on the training distribution? The paper does not reverse-engineer the LSTM solution as it does for the gated linear RNN (no polynomial coefficient extraction, no linear reconstruction scores, no weight structure visualization for the LSTM).
-
Does this failure extend to other architectures with similar nonlinearities? The Mamba architecture uses sigmoid-gated side modulation and input-dependent state projections, which are closer to LSTM-style nonlinearities than the paper's linear-gated model. If LSTM's sigmoid/tanh prevent gradient descent from finding attention, does Mamba face the same challenge? The paper speculates that Mamba has a "stronger inductive bias toward attention" (Section 3.3) but provides no teacher-student experiments with Mamba to test whether it actually learns attention or merely has the capacity for it.
What evidence exists in the paper. The evidence is the architecture comparison in Figure 4.A, where LSTM test loss is substantially higher than the gated RNN construction baseline and similar to the GRU. The paper does not provide training curves, loss trajectories, or any analysis of what the LSTM learned instead. No experiments attempt to close the gap (e.g., by initializing LSTM weights near the construction, by using different optimizers, or by annealing the sigmoid/tanh steepness during training). The paper does not test whether deeper LSTMs (2–3 layers) could implement attention when single-layer LSTMs cannot find the solution.
Mitigation status. The paper treats this as a secondary finding rather than a primary limitation. The abstract does not mention LSTMs; Section 3.3 notes the theoretical possibility and the empirical failure in passing; Section 4.3 reports the result briefly; Section 6 does not return to it. The gap between LSTM capacity and learnability is arguably the most important negative result in the paper — it constrains the generality of the central claim and raises questions about whether the connection to attention applies to any practical RNN beyond the simplified linear-gated model. By not investigating this gap further (no reverse-engineering of LSTM, no optimization ablations, no discussion of why the gap exists), the paper leaves a major uncertainty unresolved. The discussion in Section 6 states that the work is "of current practical relevance as it provides a framework that can guide future algorithmic developments," but this framework seems not to guide the understanding of why LSTMs, the most widely-studied gated RNN, fail to implement the construction. A practitioner considering whether to use an LSTM or a Mamba layer for a task where attention is known to excel receives no actionable guidance from this paper about whether the architecture will actually learn to attend, as opposed to merely having the theoretical capacity to do so under some configuration that gradient descent cannot find.
The Difficulty of Verifying the Construction in Large-Scale or Production Models Is Not Addressed
The assumption or constraint. The paper's reverse-engineering methodology — pruning zero-weight neurons, analyzing λ distributions, computing linear reconstruction scores for key-values and queries, extracting polynomial coefficients, and visualizing weight matrix structure — is applied successfully to small networks (d = 4, 100 neurons in Section 4.1). The methodology works because: (a) the optimized network is small enough that all weights can be inspected visually (Figure 2.B); (b) the teacher attention weights are known, enabling direct comparison of the RNN's internal representations to ground-truth key-values and queries; (c) the polynomial coefficient extraction requires computing and comparing coefficients of all degree-4 monomials in d variables — the number of such coefficients grows as O(d⁴), which is feasible at d = 4 (few hundred coefficients) but becomes computationally prohibitive at d = 64 (millions of coefficients) or d = 128 (hundreds of millions); (d) the trained network cleanly separates into λ = 0 and λ = 1 neurons with no intermediates, which may not hold in larger networks or different architectures.
The consequence. The paper's central empirical claim — that gradient descent discovers the construction — is verified only in a regime where verification is trivial. At the scales where the question becomes interesting (practical RNNs with hundreds of dimensions and thousands of neurons, trained on real data where the "teacher" is the data-generating process, not a known attention layer), the paper's verification methodology does not scale.
Several specific failure modes for verification at scale:
-
Weight matrix structure becomes uninterpretable. At
d = 64, the output gating matrices haved² × (d²+d)entries — over 17 million parameters. Visual inspection, block-sparse structure identification, and neuron merging (Appendix C.2) are impossible. Even automated structure detection would need to identify whether 4,096 key-value neurons and 64 query neurons can be linearly decoded, and whether output gating forms selector patterns — a substantial unsupervised learning problem in itself. -
The teacher is unknown in real tasks. On the in-context linear regression task, the paper knows the teacher is gradient descent (from von Oswald et al., 2023) and can verify polynomial coefficients against it. On a real language modeling task, the "teacher" is the corpus — there is no closed-form expression for the optimal function, and polynomial coefficient extraction in the high-dimensional token embedding space is meaningless. The paper provides no methodology for verifying the construction when the ground-truth teacher is unknown.
-
λvalues may not be binary in larger networks. The trainedd = 4network shows exactly binary λ values (Figure 2.B), but this may be an artifact of small scale and the specific optimization setup. At larger scales, with more complex tasks, λ values might spread continuously between 0 and 1, making it ambiguous which neurons are "key-value accumulators" versus "query pass-through" versus something else. The construction requires exactlyλ ∈ {0, 1}— a finding of intermediate λ values would not necessarily mean the construction is absent (approximate accumulation withλ ≈ 0.99might work), but would complicate verification. -
Compression structures may be harder to detect. The paper shows that trained networks discover symmetric and low-rank compressions (Figures 3.A-C). In a large network trained on real data, the target function's effective rank structure is unknown, and determining whether the network has discovered the "right" compressed representation requires knowing what compression is theoretically possible — which is exactly the challenge the paper does not solve at scale.
What evidence exists in the paper. The paper does not claim to have solved verification at scale, nor does it discuss this as a limitation. The reverse-engineering methodology is presented as a contribution (Section 4.1, Appendix C.2), not as a tool that needs to scale. The paper does not report computational costs for the polynomial coefficient extraction, linear reconstruction score computation, or weight structure analysis. There is no discussion of how these analyses would be conducted on a d = 64 or d = 128 network, or on a practical architecture like Mamba or RWKV.
Mitigation status. This limitation is not addressed in the paper. Section 6 suggests that "understanding whether real-world attention mechanisms lie in this regime remains an open question" but does not propose verification methodologies for answering this question. The discussion of future work focuses on "further investigations on the differences between RNNs and Transformers in controlled settings" — implying that the verification methodology stays in the controlled setting, not that it needs to be scaled. A practitioner or researcher who wants to test whether a large trained RNN is implementing attention (a question the paper's title and abstract motivate) will find that the paper provides no scalable tools for doing so. The connection between theory and practice — showing that the construction is not just possible but actually occurs in useful models — remains a promissory note uncashed by the paper's methodology.
7. Implications and Future Directions
How This Work Changes the Landscape
This paper instigates a fundamental reframing of the relationship between recurrent neural networks and attention, but it does so in a way that is more nuanced — and more useful — than a simple "RNNs can implement attention" headline suggests. The contribution is not a paradigm shift that renders one architecture obsolete in favor of another, nor is it an incremental benchmark improvement. Rather, it is a constructive diagnostic framework that gives the field a precise vocabulary and set of analytical tools for asking and answering the question: does this specific RNN architecture contain the computational primitives needed for attention-like dynamic routing of information, and if so, at what parameter cost?
Before this work, the discourse around RNNs and attention operated at two unhelpful extremes. On one side was the asymptotic universality argument — RNNs can approximate anything, so of course they can approximate attention given enough neurons — which is theoretically true but practically vacuous, providing no guidance on architectural design, parameter sizing, or whether gradient descent would actually find the attention solution. On the other side was a diffuse architectural intuition that multiplicative gating is "important" for modern RNN performance, but without a mechanistic theory of why or what computation gating enables. The paper replaces both of these with a precise account: multiplicative gating of the form (W₁x) ⊙ (W₂x) computes outer products of different input projections, which is exactly the primitive needed for key-value accumulation in linear attention; output gating of the form (W₁h) ⊙ (W₂h) multiplies accumulated key-values by current queries, which is exactly the primitive needed for the attention readout; perfect-memory recurrence (λ = 1) accumulates these outer products over time; and no-memory recurrence (λ = 0) passes through the instantaneous query. The consequence of this reframing is that architecture design shifts from an empirical trial-and-error process to a principled one: given a proposed RNN architecture, one can check whether it contains the necessary multiplicative primitives and predict its capacity for attention-like computations.
The paper's most important methodological contribution is demonstrating that this diagnostic framework can be empirically validated through reverse-engineering trained networks. The combination of λ value analysis, linear reconstruction scores for key-values and queries, polynomial coefficient comparison, and weight matrix structural analysis (Section 4.1, Table 1, Figure 2) provides a template for mechanistic interpretability that goes beyond black-box loss comparisons. The fact that trained gated linear RNNs exhibit clean binary λ values, block-sparse weight matrices matching the construction's selector structure, and exact functional equivalence (polynomial distance of 3.73 × 10⁻⁴) is strong evidence that the construction is not just a theoretical possibility but an attractor in the loss landscape — at least for the simplified architecture class studied. This methodology, while applied here at small scales, establishes a standard of evidence for mechanistic claims about what neural networks learn that the field would benefit from adopting more broadly.
However, the paper's most destabilizing finding for how the field thinks about this connection is the dissociation between theoretical capacity and empirical learnability. The LSTM, which the paper proves can implement linear self-attention in a single layer (Appendix B.3, Table 4), fails to do so under standard training — achieving the same loss as GRUs, which cannot implement attention at all (Figure 4.A). This negative result is arguably as important as the positive construction. It means that architectural inductive bias operates not just at the level of what functions can be represented, but at the level of what functions gradient descent can find. The simplified linear-gated architecture has the right inductive bias — its linearized GLU and diagonal recurrence create a loss landscape where the attention solution is reachable. The LSTM, with its sigmoid and tanh nonlinearities, apparently does not — the narrow parameter regime required for the construction (linearizing the nonlinearities via small weights in specific matrices with large compensation elsewhere) lies behind optimization barriers that standard training cannot cross. This reframes the question from "Can RNNs implement attention?" to "Which RNNs provide optimization pathways to the attention solution, and under what training conditions?"
The paper also reconciles conflicting narratives in the literature in a productive way. Prior work alternated between "RNNs are obsolete" (because Transformers outperform them on most benchmarks) and "modern RNNs are catching up" (because architectures like Mamba and LRU show competitive performance on long-sequence tasks). The paper's framework suggests a synthesis: modern RNNs catch up when and because they implement attention internally, and they do so precisely because their architectural innovations (multiplicative gating, diagonal recurrence, side gating) provide the computational primitives that older RNNs (simple RNNs, GRUs) lacked. This is not a story of RNNs defeating attention — it is a story of RNNs covertly becoming attention, with architectural modifications that make the attention computation accessible to gradient descent. The practical implication is that further architectural improvements should focus on making attention implementation more parameter-efficient and more optimization-friendly within the recurrent framework, rather than searching for entirely different mechanisms.
The in-context learning results (Section 5) add a crucial layer to this picture. The fact that gated linear RNNs converge to exactly the same gradient descent algorithm as linear self-attention on the in-context linear regression task (Table 2: polynomial coefficients match to within 0.7%, loss difference < 0.0002) suggests that the algorithm is determined by the task structure, not the architecture — architectures with sufficient capacity and the right primitives discover the task-optimal computation regardless of their parameterization. Yet the GRU, which cannot implement attention, outperforms both the attention-using architectures and the optimal gradient descent baseline on this task (Figure 4.B). This means that attention is sufficient but not necessary for in-context learning in RNNs, and that strong nonlinearities enable algorithms that can outperform attention-based gradient descent. The landscape this reveals is one where different architectural families — linear-gated RNNs, nonlinear-gated RNNs, and Transformers — converge to different points in algorithm space on the same task, with the linear-gated family converging to the same point as linear Transformers. This is a richer and more interesting picture than "architecture X is better than architecture Y."
Finally, the paper shifts the research agenda around verifier over-optimization from the LLM inference setting (where process reward models over-optimize under beam search) to a more general framing about the reliability of learned mechanisms. The parameter inefficiency of the construction — O(d⁴) parameters to implement what attention does with 3d² — is both a limitation and a diagnostic: it quantifies exactly what the RNN is paying for (the explicit enumeration of all pairwise products between key-value entries and query components) and thus identifies the bottleneck that future architectural innovations must address (side gating reduces this to O(d³), but further reduction requires fundamentally different approaches to the matrix-vector multiplication step).
Follow-Up Research This Work Enables
Closed-form reverse-engineering of LSTM and GRU solutions on in-context linear regression. The finding that GRUs outperform one-step gradient descent on in-context linear regression (Figure 4.B) is the paper's most tantalizing unexplained result. What algorithm does the GRU implement that beats gradient descent? The paper's reverse-engineering methodology — polynomial coefficient extraction, linear reconstruction of accumulated quantities from hidden state, analysis of gate activations over time — could be applied to trained GRUs and LSTMs on this task to determine whether they implement multi-step gradient descent, preconditioned gradient descent (using running estimates of input covariance for adaptive learning rates), recursive least squares, or something else entirely. A strong follow-up would train GRUs with varying numbers of hidden units on the linear regression task with d_x = d_y = 3 and T = 12 (matching the paper's setup), extract the polynomial coefficients of the learned input-output mapping as in Table 2, and compare against the polynomial forms of known optimization algorithms (2-step gradient descent, gradient descent with momentum, recursive least squares with forgetting factor). The result would either identify the algorithm or establish that the GRU has discovered a novel in-context learning algorithm, which would be a significant finding in its own right. Additionally, testing whether the GRU's advantage over gradient descent grows with sequence length T (does the algorithm compound gains over more observations?) or with input dimensionality would characterize the scaling properties of whatever mechanism the GRU employs.
Demonstrating or refuting that practical-scale architectures (Mamba, RWKV) converge to the construction when trained on real data. The paper's positive results are confined to the simplified linear-gated architecture at d ≤ 12 and n ≤ 100. A critical follow-up would train a Mamba layer (which has the side-gating variant of the construction, Appendix A.2) on a synthetic teacher-student task at moderate scale — say, d = 64 with a rank-8 teacher attention layer — and apply adapted versions of the paper's verification tools to determine whether the trained Mamba layer implements attention. The verification would need to handle the architectural differences (input-dependent B matrices instead of fixed input gating, sigmoid in the side gate instead of linear gating) and the scale (4,096 key-value neurons, output gating matrices with millions of entries). Specific checks could include: (a) whether the learned B matrices produce key-value outer products that can be linearly decoded from the hidden state (extending Score KV to Mamba's recurrence), (b) whether the side gating matrix aligns with the query projection (extending Score Q), (c) whether Mamba's λ parameters (the A matrix) converge to binary values or approximate them with exponential decay rates near 0 or 1. A negative result — Mamba failing to learn the attention solution at moderate scale despite having the theoretical capacity — would be highly informative: it would indicate that the paper's positive findings are specific to the linear-gated regime and that sigmoid-gated architectures require different optimization strategies or architectural modifications to realize their attention capacity. A positive result — Mamba converging to a compressed attention implementation at d = 64 — would substantially strengthen the paper's practical relevance and suggest that the connection to attention partly explains Mamba's empirical success.
Compression-aware RNN architectures with explicit low-rank inductive bias. The paper demonstrates that trained gated RNNs discover compressed representations that exploit symmetry and low-rank structure in the teacher attention matrices (Section 4.2, Figures 3.A-C), but the compression is discovered implicitly through overparameterization and pruning — the architecture itself imposes no low-rank structure. A natural follow-up is to design an RNN architecture that explicitly encodes the compressed construction: instead of having n general-purpose hidden neurons with learnable λ, the architecture would have r(r+1)/2 key-value accumulator neurons and r query neurons, where r is a hyperparameter controlling the effective rank of the emulated attention. The input gating would be structured so that one branch computes r values and the other r keys, the recurrence would separate into λ = 1 and λ = 0 groups by construction, and the output gating would have r²/d times fewer parameters than the naive construction. Training this structured architecture on the in-context linear regression task with d_x = d_y = 3 and T = 12 (matching Section 5) should achieve the gradient descent solution with r = 3 and far fewer parameters than the overparameterized general architecture. A controlled comparison — structured architecture with r = 3 vs. general architecture with equivalent parameter count — would test whether the explicit inductive bias accelerates convergence, improves sample efficiency, or enables learning at smaller scales where the general architecture fails. If successful at toy scale, the architecture could be tested on larger synthetic tasks (d = 64 with known low-rank teacher) and eventually on real benchmarks where attention patterns are hypothesized to be low-rank.
The overparameterization threshold: theoretical characterization and scaling behavior. The paper observes a consistent ~2× overparameterization requirement for gradient descent to discover the construction across three different settings (Figures 3.A-C), but provides no explanation for why the factor is 2, whether it is constant across scales, or what determines it. A theoretical follow-up could analyze the loss landscape of the gated RNN around the construction solution. Specifically: (a) characterize the null space of the network around the construction — how many redundant parameters can be perturbed without affecting the output, and does gradient descent use these directions during optimization? (b) Determine whether the 2× factor is a consequence of the specific parameterization (each effective neuron requiring a "twin" with opposite-sign weights to cancel bias terms from the gating, as suggested by Figure 2.C where three neurons merge into one) or a more general property of overparameterized bilinear optimization. (c) Empirically measure the threshold at larger scales — d = 16, d = 32, d = 64 with full-rank and low-rank teachers — to determine whether the factor grows, remains constant, or perhaps decreases with dimension. If the factor grows with d (e.g., as log d or √d), the practical feasibility of learning attention in RNNs at scale would be substantially worse than the paper's small-scale experiments suggest. If it remains constant at ~2×, the 2× rule becomes a practical sizing guideline for practitioners.
Extending the construction to approximate softmax attention via kernel methods. The paper's construction applies only to linearized attention (Equation 1), which removes the softmax nonlinearity. Softmax attention differs fundamentally because the normalization ∑_{t'} exp(k_{t'}^\top q_t) creates a non-separable interaction among all positions that cannot be expressed as a simple outer product accumulation. However, recent work on efficient Transformers has shown that softmax attention can be approximated by linear attention in a transformed feature space using random Fourier features (Rahimi & Recht, 2007) or positive random features (Choromanski et al., 2021). A follow-up could combine this kernel approximation literature with the paper's construction: use the gated RNN's input gating to compute not (W_V x)(W_K x)^\top but φ(W_V x) φ(W_K x)^\top where φ is a finite-dimensional random feature map approximating the Gaussian kernel, which in turn approximates the softmax. The question is whether the gated RNN can implement φ using its existing components (perhaps through a GLU followed by a nonlinear activation) or needs additional architectural modifications. A concrete experiment: train a gated RNN to emulate a softmax attention layer at d = 8 by providing the RNN with random Fourier feature projections as additional input dimensions, and test whether the RNN learns to use these features to approximate the softmax normalization. If successful, this would extend the paper's framework from the linearized attention special case to the attention variant used in all production Transformers, dramatically increasing practical relevance. The parameter cost would increase (because φ typically requires 2–3× the input dimension in random features), but the existence of a finite construction would establish that RNNs can approximate softmax attention with bounded memory, which is a stronger claim than the paper currently makes.
Mechanistic interpretability of in-context learning in Mamba and other deployed RNNs. The paper's methodology for verifying that a trained network implements attention — linear reconstruction of key-values and queries from hidden state, polynomial coefficient comparison, λ value analysis — provides a recipe for analyzing what large-scale deployed RNNs actually compute. A ambitious follow-up would apply these tools to a pre-trained Mamba language model (available open-source at various scales from Gu & Dao, 2023) on a controlled in-context learning probe. The probe could be the associative recall task from Appendix D.3, embedded into the token vocabulary of the pre-trained model, or a linear regression task where the inputs and outputs are formatted as token sequences. The analysis would test whether the Mamba model's internal representations on these probes exhibit the signatures of attention (linearly decodable accumulated key-value outer products, query representations in specific state dimensions, output gating that multiplies the two). A positive result — evidence of attention-like computation in a model trained on general language data, not specifically on the probe task — would be the strongest possible validation of the paper's central thesis, connecting the constructive theory to the "under the hood" claim in the title. A negative result — no evidence of attention-like structure despite the Mamba model solving the probe tasks — would be equally valuable, suggesting that Mamba's success on language modeling comes from mechanisms other than covert attention, which would redirect the research agenda away from the attention-RNN connection and toward characterizing what alternative algorithms Mamba actually implements.
Practical Applications and Downstream Use Cases
Architecture selection guided by the multiplicative primitive checklist. For practitioners choosing between RNN variants (LSTM, GRU, LRU, Mamba, RWKV) for a sequence modeling application, the paper provides a concrete diagnostic: if your task requires dynamic routing of information between arbitrary past and present positions — the kind of computation that attention excels at — select an architecture that has both an input-side multiplicative gating mechanism (to compute key-value outer products) and an output-side or side gating mechanism (to multiply accumulated key-values by queries). Among the architectures analyzed, this checklist favors: variants of LRU with added input gating (LRU In+Out from the paper), Mamba (which has input-dependent B matrices for input gating and side gating for query multiplication), and stacked standard LRUs (2+ layers where the first layer's output GLU serves as input gating for the second). It advises against: single-layer standard LRUs (missing input gating), GRUs (cannot compute multiplicative outer products at all, even with stacking), and LSTMs (have theoretical capacity but the paper shows standard training doesn't find the attention solution, so they are unreliable for tasks where attention is the target mechanism). For MLP-based alternatives to gating (e.g., feedforward blocks after recurrent layers), the MLP ablation in Figure 4.A shows dramatically worse performance on attention-mimicking, so these should be avoided if attention-like computation is needed. This checklist is directly actionable: it translates the paper's theoretical analysis into an architectural decision rule that practitioners can apply without conducting their own reverse-engineering experiments.
Sizing RNN hidden states for tasks where attention performance is the target. The paper's quantified parameter cost analysis provides a sizing heuristic for practitioners who want to deploy RNNs on tasks currently dominated by Transformers. The compressed construction (exploiting symmetry and low-rank structure) requires approximately r_KQ · (r_V + 1) hidden neurons, where r_V is the effective rank of the value matrix and r_KQ is the rank of the key-query product matrix in the target attention computation. With the ~2× overparameterization threshold observed in Figures 3.A-C, a practitioner should provision at least 2 · r_KQ · (r_V + 1) hidden neurons in their RNN to have confidence that gradient descent can discover the attention solution. For a concrete example: if a task is known to be solvable by a linear attention layer with effective rank r_V = r_KQ = 16 (a realistic estimate for many sequence modeling tasks based on attention matrix analyses in the literature), the practitioner should allocate at least 2 · 16 · 17 = 544 hidden neurons. This is substantially larger than the d² + d = 16² + 16 = 272 neurons that the naive full-rank construction would suggest, but substantially smaller than the 2d(d+1)/2 + 2d ≈ 2 · 136 + 32 = 304 that the symmetric full-rank overparameterized construction would require. Without the paper's compression analysis, a practitioner might naively allocate d² + d = 272 neurons and then be confused when the network fails to learn (since this is below the 2× threshold for the construction to be discovered). The heuristic converts the paper's theoretical compression results into a practical sizing guideline.
Verification of attention implementation in deployed RNNs for safety-critical interpretability. For applications where understanding the model's internal computation is necessary for safety or reliability (medical diagnosis from time series, financial forecasting, autonomous vehicle trajectory prediction), the paper's reverse-engineering methodology provides a template for post-hoc verification of whether a trained RNN is implementing an interpretable algorithm (gradient descent, associative recall via attention) or an opaque computation. The specific steps — use a small probe dataset of known difficulty, extract linear reconstruction scores for hypothesized accumulated quantities, compare polynomial coefficients to known algorithmic solutions, and verify that effective state size matches the theoretical construction — can be integrated into a model audit pipeline. For instance, if an RNN-based medical time series model is making predictions that a clinician needs to trust, the audit could determine whether the model has learned to accumulate something resembling a running regression (interpretable) or is using the recurrent state in a way that cannot be decomposed into known primitives (opaque). The paper's finding that some architectures (the simplified linear-gated RNN) yield interpretable learned weights while others (LSTM) yield opaque solutions can guide architecture selection toward auditable systems. This application is nascent — the verification tools currently work only at small scales and with known teacher algorithms — but the paper establishes that such verification is possible in principle and provides the first demonstration of it.
When to Prefer This Method
The paper does not propose a specific deployable method to prefer over alternatives. It provides a constructive theoretical equivalence and a reverse-engineering analysis framework, not an architecture or training procedure that competes with existing approaches on practical benchmarks. The paper's value is diagnostic and conceptual — understanding when and how RNNs implement attention, not providing a better RNN — so a "prefer X when Y" decision matrix would be inappropriate here. The relevant decision the paper informs is not which model to use, but which architectural elements to include when designing a new RNN if attention-like computation is desired, and the guidance is summarized in the architecture selection checklist in the first downstream use case above.
