On the Expressiveness of Softmax Attention: A Recurrent Neural Network Perspective
Authors: Gabriel Mongaras
Eric C. Larson
Venue: Accepted by TMLR ()

Abstract: Since its introduction, softmax attention has become the backbone of modern transformer architectures due to its expressiveness and scalability across a wide range of tasks. However, the main drawback of softmax attention is the quadratic memory requirement and computational complexity with respect to the sequence length. By replacing the softmax nonlinearity, linear attention and similar methods have been introduced to avoid the quadratic bottleneck of softmax attention. Despite these linear forms of attention being derived from the original softmax formulation, they typically lag in terms of downstream accuracy. While strong intuition of the softmax nonlinearity on the query and key inner product suggests that it has desirable properties compared to other nonlinearities, the question of why this discrepancy exists still remains unanswered. This work demonstrates that linear attention is a first-order approximation of the softmax numerator by deriving its full recurrent form. We further show empirically that the denominator's function can be effectively replaced by a simple vector norm. Using this form, each part of softmax attention can be described in the language of recurrent neural networks (RNNs). Describing softmax attention as an RNN allows for the ablation of the components of softmax attention to understand the importance of each part and how they interact. In this way, our work helps explain why softmax attention is more expressive than its counterparts. Code found at: https://github.com/gmongaras/On-the-Expressiveness-of-Softmax-Attention-A-Recurrent-Neural-Network-Perspective

Published in Transactions on Machine Learning Research (10/2025)
On the Expressiveness of Softmax Attention: A Recurrent Neural Network Perspective
Gabriel Mongaras
gabrel@mmongmma.com
Deportment of Computer Science
Southern Methodist Universily
Eric C. Larson
eclarson @smu. edu
Institute for Computational Biosciences Southern Methodist University
Reviewed on OpenReview: /ttps::/openrreet net/ forum? Pd=P3330
Abstract
Since its introduction, softmax attention has become the backbone of modern transformer architectures due to its expressiveness and scalability across wide range of tasks. However, the main drawback of softmax attention is the quadratic memory requirement and computa- tional complexity with respoct to the sequence length. By replacing the softmax nonlincarity, linear altention and similar methods have been introduced to avoid the quadratic bottleneck of softmax attention. Despite these linear forms of attention being derived from the original softmax formulation, they typically lag in terms of downstream accuracy. While strong intuition of the softmax nonlinearity on the query and key inner product suggests that it has desirable properties compared to other nonlinearities, the question of why this discrepancy exists still remains unanswered. This work demonstrates that linear attention is a first-order approximation of the softmax numerator by deriving its full recurrent form. We further show empirically that the denominator's function can be effectively replaced by simple vector norm. Using this form, each part of softmax attention can be described in the language of recurrent neural networks (RNNs). Describing softmax attention as an RNN allows for the ablation of the components of softmax attention understand the importance of each part and how they interact. In this way, our work helps explain why softmax attention is more expressive than its counterparts
1 Introduction and Background
The formulation of softmax attention was proposed by Bahdanau et al. (2055) as weighting mechanism for aligning recurrent neural networks (RNNs) in encoder-decoder architectures for language translation Sutskever et al. (2014). However, the modern day usage of the attention in a transformer architecture was first employed by Vaswani al. (2017) as a way to do soquence mixing within the context of language translation. This formulation used the softmax activation function model token routing without the use of traditional recurrent network. Since its introduction, the attention mechanism has been widely adopted in various domains such as computer vision Dosovitskiy et al. (2021). generative models Esser et al. (2024). timeseries analysis Nie et al. (2023), audio processing Baevski et al. (2020), graphs Veličković et al. 2018 and many more applications.
While softmax attention is a powerful mechanism that used a variety areas, a major drawback its quadratic complexity with respect to the sequence length, N. Linear attention swaps the softmax nonlinearity
(Code found at: hhtp:///am....mmmooaoooaoeeet--e
Attention-A Naanreett-Nurral-eewwwooo
Perspective

---

Published in Transactions on Machine Learning Research (10/2025)
with a decomposable kernel function, reducing the complexity from quadratic to linear with respect to the sequence length. The original introduction of linear attention Katharopoulos et al. (2020) replaced the nonlinearity with an elu(x) + 1 kernel. Since the original formulation, other methods have been developed that replace the nonlincarity on the QK-softmax inner product with various other functions such as ReLU Xic
et al. (2025), , cosine similarity Mongaras et al. (2025), cosine reweighting QQn et al. (2022) or by decomposing the QK-softmax inner product Wang et al. (2020); Zhuoran et al. (2021). Although these methods are linear in complexity, none are as performant as softmax attention in terms of downstream accuracy.
In this work, we propose recurrent reformulation of softmax attention and use this to describe the elements of softmax attention that make it more performant compared to its linear and sparse counterparts. While softmax attention is the driving method used in many architectures, the question of why the discrepancy between softmax attention and linear attention exists remains unanswered. This work provides a principled analysis of softmax attention by deriving a recurrent formulation using its Taylor series expansion. Through this formulation, we experiment to reveal how each recurrent component contributes to downstream accuracy using targeted ablations. We demonstrate that softmax attention is not merely a heuristic construction but. structured process with interpretable, sequential dynamics. In doing so, we bridge the gap between the observed empirical performance for lincarized attention (and its variants) and the thcorctical underpinnings of softmax attention.
2 Related Work
Expressiveness of Linear Attention Many approaches have attempted to make linear attention more expressive to match softmax attention, while still retaining linear complexity. Choromanski et al. (2021) used linear approximations of the softmax kernel to achieve more performant linear attention mechanisms.
As linear attention can be interpreted an RNN Katharopoulos et al. (2020). Peng et al. (2025) proposed receptance weighted key values (RWKV) attempting to enhance the expressivity of RNNs with receplance vector for time mixing. While this method helped to address computational complexity of RNN training, it still relied on linear attention in its formulation. Mamba Dao and Gu (2024) employed state space models Gu ltt al. 2022) to develop an efficient and expressive form of linear attention. Sun et al. (2025) formulated linear atlention as a step in gradient decent of a hidden view. Another approach treats modeling the hidden state as step in gradient decent [Sun et al.) (2025) which can be viewed as type of linear attention. [Berruzz et al. (2024) proposed Titans, which built upon this gradient formulation creating different variants of the hidden view gradient descent. ATLAS Behrouz et al.
(2025) develops new kind of recurrent models, leveraging test-time compute as in Behrouz et al. (2024). They include a proof show that recurrent models using higher order hidden states are more expressive. Sieber et al. (2024) makes a similar derivation as this work, but in the context of control systems, but did not apply this derivation to a recurrent architecture. [auen et al. (2024) makes a simple derivation, but only explores second-order models. Although all of these methods are more expressive than linear attention, they are difficult to implement efficiently in hardware, akin to RNNs, and fall short to the accuracy of softmax attention.
Softmax Attention Improvements Another direction of research explores how to improve softmax attention. A number of works attempt to improve softmax attention by increasing sparsity to reduce dependence on context length such as Longformer [Beltagy et al. (2020), BigBird Zaheer et al. (2020), and randomized feature attention Peng et al. (2021). RoFormer (Su et al. (0024) adds relative positional encodings to improve long sequence modeling and is used in most modern transformer models. Zhaii et al. 2002)
prevents attention score entropy collapse by reparameterizing the weight matrices, modern transformers use norms to fix this issuc. Newer improvements such as the Forgetting Transformer Lin et al. 20055 and
DeepSeek Liu el al. (2024) make changes to the input of the attention mechanism such as combining the key and value matrices reduce the KV cache size, or adding "forget" gate. Many of these improvements make softmax attention slightly less computational, less memory intensive, or more expressive. However, the core mechanism driving the modeling capability is still softmax attention.
Why Softmax? There are several works that also attempt to address the question: What makes softmax attention performant? This fundamental question has eluded researchers for some time. For example,
2

---

Published in Transactions on Machine Learning Research (10/2025)
noticed softmax attention cannot "zero out" attention heads. More specifically, when -im for all tokens, Hi, it is desirable for the head output to produce zero, thus giving no weight to any tokens. In actuality, it produces uniform distribution over the tokens instead. To fix this issue, Miller (2023) adds a constant to the denominator. This work emphasizes how softmax has subtle problems even though it works
Miller 2023
well practice. (mmith (025)) asks why attention works, building upon Miller (2023) to change the attention mechanism. However, the focus this work on the improvements to attention, rather than exploring exactly what makes attention work. Deng et al. (2023) examines the performance gap between linear and softmax attention, but only in classification and empirical context. Why linear attention is lacking compared to softmax attention, in general, remains not. fully understood. Han et al. (2024) shows that linear attention is lacking an injective property and cannot model local features while softmax attention is injective and can model local features. Collins et al. (2024) uses Lipschitzness to explain in-context learning in softmax. These provide intuition but, ultimately, these properties do not fully explain the expressiveness gap.
Katharopoulos al. (2020) introduced a formulation for lincar attention using recurrent network components. In their work; they also mention that their "formulation does nol impose any constraind on the fenlure funclion and it can be used for representing any transformer model, in theory even those using softmaz attention." This inspired our current work to formulate softmax attention as RNN and use this formulation to describe what makes softmax attention expressive, in general. Moreover, we attempt to explain exactly why linear attention is not as expressive as softmax attention by showing linear attention approximates a subset. of elements that comprise softmax attention.
3 Methodology
We first motivate our approach using the traditional formulation of causal softmax attention for calculating the output of the layer, O,, at time t:
O: = Softmax(Q; K1a) Va. Exn eo.k? V.
(1)
En=r 20.KT
We would like to rearrange this equation such that one can interpret the attention mechanism in recurrent form. However, the exponentiation in softmax couples Qt and K,, preventing a direct regrouping. For instance, in linear attention the attention function is replaced by ((Q.) (KK.)) This decoupling allows causal linear attention to be reformulated into a recurrent structure, as follows:
O, = ((:)) WV(K))..V = ((Q.) H,
(2)
where
H1 = Σ (K,)) .V = =(KK) +1++
In order to derive a recurrent representation for softmax, we analyze the numerator and denominator of softmax attention separately. We first examine causal softmax attention without the denominator and show there exists a recurrent form. We use the Taylor series expansion of softmax to achieve this formulation (Section 3.1 Having a recurrent form, Scction 3.2 shows linear attention, as equation is a first order approximation of softmax allention. Appendix B shows another motivating example the quadratic case for intuition of how the recurrent form expands for higher order Taylor series approximations. In Section 3.3 the denominator is reinterpreted, using the language of RNNs with results provided in Section
3.1 Recurrent Softmax Attention
Causal softmax attention, equation I composed of exponentials of inner products between the query vector, Qt, and key vectors, Kg- Using decomposed inner product form, causal softmax attention be rewritten
3

---

Published in Transactions on Machine Learning Research (10/2025)
RNN by taking the Taylor series expansion of the exponential function. Here, set the multiplicative inverse of the denominator to Gt for simplicity. Although the derivation is for the causal case, it can be extended to the bidirectional case as in Appendix
Note that, in the explanation below, we make use of several properties:
1. Decomposed Inner Product Property: (A -B" is equivalent to (A@"). (Ben) by equation equation full derivation of this property is available in Appendix
Inner Product Equivalence: (A, B) = Ei (AQ B); = - A,B;
3. Kronecker Products Shorthand: A A = X i= A = A 8 A 3
2 A IRd
where A, B E Rd, A B E Rd is the Hadamard product, A Bd E R is the inner product, AT B E Raxd is an outer product, and the lack of an explicit operation denotes either: (1) scalar multiplication with another scalar, vector, or matrix or (2) matrix-vector multiplication, similar to "slicing multiplication" in most numerical packages. With these properties, we can now reformulate softmax attention as follows:
O, = Gt Σ
Qrk
8 G.
(Q: KV
By definition of the Taylor Series of
n!
n=0
d"
G
(0th), ((KSN)²), V.
By equation
n=0
80
G.
(Qtt) (KK? V.
By rearranging sums
4
G.
Σ (K)
By factoring out Q By inner product equivalence
s=1
G.
(K V.
n=0
s=1
G. Σ
HH, H. = [(KKmn)) V.
Define hidden state
n=0
Where Q R.N, K M.R.d, V RRRM,GERR, Rd,K, V e
N and M are the sequence dimensions, indexed by t and respectively.
d (and e in Figure I are the embedding dimensions where d is indexed by i. n is the th order term in the Taylor expansion.
Thus, the softmax attention numerator does have a recurrent formulation. This formulation can also seen visually in Figure [ Rather than the output coming from a single recurrent equation, softmax a sum of infinite recurrent outputs, cach weighted by n!- Each of the RNNs that comprise the output of softmax attention have a hidden state of shape dd due to the th order Kronecker product on the queries and keys. The th order Kronecker product can be thought as creating th order multiplicative interactions between dimensions of the keys and queries. The hidden state for each of the infinite RNNs can be thought of accumulating information of n order interaction terms the dimension of Q and K.
4

---

Published in Transactions on Machine Learning Research (10/2025)
Ier Podd
16
CAAA thod
Coter Prest
Chnobrade
Higer
Figure 1: Softmax attention as an RNN. We define G, in place for the softmax denominator. Linear attention is equivalent to the 1, first order, term. Expanded plots can be found in Appendix G
3.2 Linear Attention is First Order Approximation
With this recurrent formulation, notice that taking the term of the Taylor Series sum results in the well known form of linear attention, as seen in Equation 2
oftinear)
1 (PP)) Σ ((KK))) 1!
V, Q Σ V. V. H, Σ KT V. s=1
(3)
Equation and the left portion of Figure shoow that linear attention is a linear approximation of softmax attention when == + ls. As linear attention is just a single term in softmax attention, this derivation shows how lincar attention is subsct of softmax attention and provides linear attention is typically less performant. Even when other functions are used for é and tb they can only manipulate single recurrent chain-they they cannot model the higher order terms from soft maxt These terms in softmax attention allow it to model combinatorial interactions between inner product dimensions thus, it is at least as expressive as linear attention. Also notice that each subsequent RNN comprising softmax uses a larger hidden state, modeling higher-order interactions on the combinatorial dimension of Q and K. Linear attention works with single, smaller hidden state operating only on the dimension of Q and K, without higher-order interactions. To highlight the difference between softmax attention and linear attention, we also derive quadratic approximation in Appendix B which models pairwise interactions the dimension and has hidden state quadratically larger than linear attention.
intuitive explanation as to why
This result raises an important difference between linear atlention, with functions on Q and K, versus nonlinear attention, which uses a non-decomposable function on the inner product of Q and K. Despite the two forms appearing similar, they are fundamentally different in their ability model higher order interactions. function on each of the vectors, o(Q,) and vi(K)) restricts the vector space dimensions of Q, and Kg. For example, ReLU restricts vectors to the positive portion of the vector space. On the other hand,
5

---

Published in Transactions on Machine Learning Research (10/2025)
the exponential function the inner product space does not restrict the vector space. Instead, this function creates th degree multiplicative factors between dimensions of Qt and Kg for all n E (0. Practically, these multiplicative dimensions become less influential for larger n, as they are weighted by 1 cannot be disregarded entirely.
mt but they
3.3 Reinterpreting the Denominator
In traditional softmax attention the denominator is calculated as:
G. E C9.
(4)
where G, can have values from O to 1 1 Typically, these values are calculated from the same Q2 and K2 values as the numerator. However, we hypothesize that the exact normalization not the most crucial aspect of Gt. Rather, we observe that G, could be interpreted like a gate or norm that stabilizes the numerator, especially for long contexts when becomes large. We that this re-interpretation of the denominator may not be strictly accurate, but does provide convenient mechanism for capturing some aspects of the softmax denominator. This stabilization aspect is hypothesized to be the crucial function of G₁, rather than the exact form of G,. Therefore, we reformulate Ge as a gate via:
(Gate)
1 Gt t
Q K
V.
V== Σ ( (Q")) H n=0
H = =- Σ (K V.
(5)
This assumption of the role of G, allows us to approximate like an output gate a recurrent structure time t. This representation also ties the softmax attention mechanisms traditional properties of expressive recurrent nelworks, such as the LSTM Hochreiter and Schmidhuber (1997) and GRU Cho et al. (2014). However, this interpretation grow without bound. This issuc can be mitigated by dividing by the sequence length, clamping the inner product (for example, we clamp the pre-exponential value to 5), and clipping the gradient. While these tricks help ensure the gate is numerically stable, they are not particularly elegant solutions and complicate the overall implementation. Alternatively, G, can be interpreted as norm at time t, normalizing the numerator according its length and values. This can be realized via:
gate may be too broad of an assumption it also allows the sequence to
1-
_
0e Σ
(Norm)
Σ CCKKT
KT V3
(Q HH
H. = Σ (K V.
(6)
where ||-|| denotes vector norm. The optimal type of norm is not immediately clear, and could be any number of formulations such L2, RMS, others. We test each interpretation of G, under a recurrent perspective, as either a gate or a norm, to understand its equivalence to the traditional softmax denominator.
4 Results
To evaluate our hypothetical softmax alternative, we train multiple Llama 2 Touvron et al. (2023 Touvron ct al. (2023) models for next token language modeling. Keeping the rest of the architecture constant, we replace the attention mechanism with the proposed variations. We show log loss to emphasize the differences between each model. Section 4.1 shows our proposed replacement is empirically equivalent to softmax. Section 4.2 examines the scalability of the proposed replacement. Section 4.3 evaluates linear attention against normal softmax and our proposed methods. As our method uses a Taylor expansion, Section 4.4 looks at the performance different linear attention variations via progressive additions of higher order powers. Section ablates various elements of the recurrent softmax attention.
2Gs can he greater than the denominator than 1, which rare, only occurring for small

---

Published in Transactions on Machine Learning Research (10/2025)
4.1 Softmax Equivalence
Experimental Sctup In our experiments, test both replacing the denominator with gate in equation 5 and with a norm as in equalion p aloongside normal softmax attention. To evaluate the applicability to various domains, we retrain on three datasets: The Pile Gao et al. (222]) SlimPajama (Shen et al. (023) and Wen Penedo et al. 2024/. The Pile is a dataset created by Elcuther AI composed of 825 GiB of English text on various domains such as code, technical papers, math, and articles. SlimPajama is a 627 billion token dataset that a cleaner subset of RedPajama Weber al. (2024). a dataset of various web crawl data. Feb is cleaned and de-duplicated 5-trillion token dataset compiled from 96 different common crawl snapshots. Each tested model is about 300 million parameters and trained on sequence length of 1024. Adding gate or norm requires up to an additional d parameters per layer, which is insignificant relative to the model sizc. More hyperparameters for our models can be found in Appendix D
13
133
12
Figure 2: Test and train loss on various datasets for softmax attention and the proposed methods with gate or norm replacements. Expanded plots can be found in Appendix G
Result Figure shows the resulting loss curves for each dataset and model variant. The loss for the normed model variant follows the model trained with native softmax attention precisely while the model with a gate performs slightly worse. We employ the Le vector norm for this analysis and find that it is numerically stable. The gate was semi-unstable during training. We noticed recoverable loss spikes during training where the model loss would spike and resume after some number of steps. Dividing by the sequence length, clipping the inner product before exponentiation, and performing gradient clipping, helped produce fewer spikes, however some instability was still observed. This result implies that the function of G, is most similar to norm operation. Moreover, the norm does not need to mirror the traditional softmax exponentiation a simple vector norm approximates this well. We note, however, that this does not prove the denominator can be a simple recurrent gate or norm we only claim that performance is similar.
4.2 Scaling
To investigate if scaling laws hold, scale the model in two ways. We scale the model from 300M parameters to 2B parameters, keeping all other hyperparameters constant to evaluate the scalability of the model itself. We also scale the sequence length from 1024 to 4096, keeping all other hyperparameters constant to examine how our method scales with longer sequences. These scaling are tested upon the FineWeb Penedo et al. (2024) dataset. The scaled comparisons are found Figure 3. showing our proposed attention formulation scales identically softmax (with both the sequence length and the model size). Again we observe that the morm better mirrors softmax performance and is more numerically stable
7

---

Published in Transactions on Machine Learning Research (10/2025)
Large Model (28 Params| Fineweb Log Test Loss
Large Model (28 Params) Fineweb Log Troin Loss
Gole kore
Seltrie
Gata
Horn
1.3
1.6
104
1.02
2.06
Step
30C
106
306
666
4099 Sequence Length Fineweb Log Test
4095 Sequence Length Fineweb Train Loss
155
1.25
133
Stpp
0.3
500
10k
25k
20k
3X<
Figure 3: Test and train loss on the Web dataset for large models (about 2B parameters) on 1024 sequence length and small models (about 300 million parameters) on 4096 sequence length for sofumax attention and the proposed methods with gate or norm replacements. Some experiments were cut short due to time constraints
4.3 Linear Attention
As the proposed softmax decomposition has a direct relationship to linear attention, we evaluate the method against several variations of linear attention. Figure S shows that softmax attention and our proposed mothods outperform each linear allention variant by a significant gap. Additional details are provided in Appendix F This behavior supports the assertions in Section 3.2 showing linear attention subset of softmax attention.
Fineweb Log Test Loss
Finewab Log Train Loss
Solynce
1.4
1.3
Lenay Lecar Re.
1.2
5412
1001
Solunat
Sbep
29k
40k
Figure 4: Test and train loss on the FineWeb dataset for various linear attention methods, softmax attention, and the proposed methods with gate or norm replacements.
4.4 Taylor Series Terms
As our method uses a Taylor series, we investigate how performance behaves as higher order terms are added to the less expressive atlention variations. That is, we use the recurrent formulation for a particular lype of linear attention and gradually add in higher order terms from the softmax Taylor expansion of softmax. The results are shown in Figure 5 for three types of linear attention: cosine similarity Mongaras et al. (2025), ReLU Xie al. (2025), and elu(x) + kernel Katharopoulos et al.) (2020). As more terms are added to the softmax approximation, the performance smoothly transitions from "linear attention performance"
8

---

Published in Transactions on Machine Learning Research (10/2025)
to "softmax attention performance." After adding in terms up to 10, we observe that the recurrent approximation mirrors soft with negligible differences. For the linear attention variants, however, we find that adding more terms (described in Appendix Fh improves performance but never fully reaches the performance of softmax attention (even for n = 10). We hypothesize this performance gap exists because linear atlention variants have functions on the independent vectors, ((Q) and ((K), restricting the reachable vector space when combined; whereas softmax does not restrict this vector space. We leave exploring this observation to future work. We note that cosine attention does not gain any benefit from additional terms and hypothesize that due to inner product values being between 0 and 1, the resulting higher order interaction terms are less than 1, limiting the magnitude of higher order terms and therefore the impact of these higher order terms on the resulting output.
Softmax Log Train Loss
Cosine Log Train Loss
FullSeftruw
Dider Sohmes 4th Order Soboee ee Dider Satnaa Ordar Sofures Order Sotrean
Liver Coliee
Ordlrr
3nd Order Coslne
Uran Codea
1.8
Dolo Coaltz
2nd Order Softror
Onder Seftroe
3et Order Coslee
L0th Ceder Cod
16
411 Dider Softrnn
#7h softa
1.2
10th Order Saborax
3.2
Stop
Step
RelU Log Loss
Log Train Loss
Lineut Sall
Orcer RuUU
Order Re.3
4th Grter Beili
Order
Sofmnex
Limer ELU
2nd ELU
Orde EUU
op Orderde ULU
EnthCrtteelL
Sofrnac
User
RUUU
1.8
3v1 Order Bell)
3.5
Order
1.
1.6
14
Order
Utde tell
Softres
12
Sohruu
54
100
15k
206
25k
30k
5k
13k
154
20k
25k
30k
Figure Log train loss for softmax attention and various linear attention method when summing more powers of the inner product. The nth order denotes the sum of powers from 0 to n.
4.5 Ablation Analysis
To further investigate various recurrent elements, we conduct an ablation study using the 300M parameter model, varying both the gated and normed variants of our method. The results of these ablations are plotted in Figure G The leftmost plots ablate several elements of the recurrent softmax denominator, G,, by removing, detaching, and dividing by the sequence length, S (with varying combinations of all three). We find that detaching the denominator from the computation graph significantly hurts downstream performance. Removing the denominator and dividing by the sequence length gives similar performance to softmax, but also creates instability. Adding a gale helps to somewhat stabilize the training, though not entirely. Finally, we find that removing the denominator and adding a norm is what appears to mirror softmax most closely,
Gate The middle column of plots in Figure ablates the gate method employed by combining sequence length normalization with an input gate or output gate (more fully explained in Appendix E. and investigating ReLU linear attention with a gale. The results shows that having a gale helps to stabilize training loss, but must be accompanied by sequence length normalization for good test performance. Furthermore, replacing the exponential with a decomposable ReLU kernel significantly hurts performance, regardless of gating or sequence length normalization. This result emphasizes the importance of the exponential function on the inner product of Q and K, as well as indicating that some sort of sequence length normalization is required.
9

---

Published in Transactions on Machine Learning Research (10/2025)
Faneeb Tokl LoM
Log Trais
Fizeab Log TAME LOG
og
11
Denominator Ablations
Gate Ablations
Norm Ablations
Figure 6: Test and train loss on various datasets for softmax attention and the proposed methods with gate or norm replacements. Expanded plots can be found in Appendix G
Norm The rightmost column in Figure G varies the norm method employed, investigating the L2 norm, RMS norm, layer norm, and Lg norm with sequence normalization. We also investigate the use of an L2 norm with the decomposable ReLU kernel. The results shows that any type of normalization, with or without. learnable parameters, works just as well as softmax. However, replacing the exponential with a decomposable ReLU kernel significantly degrades performance. This graph further emphasizes that the important aspects of softmax attent ion are an exponential plus a vector norm. The choice of norm does not appear to influence training stability or test performance. However, there are differences between a vector norm (c.g., Lz2) and a sequence length normalization (e.g., division by 5 or gating). The vector norm appears to achieve the best performance, implying is a necessary aspect of softmax, but the exact form seems to be unimportant.
5 Conclusions and Limitations
This work connects linear and softmax attention under a unifying recurrent formulation. A Taylor series expansion of the softmax attention numerator was employed to develop recurrent form and competing hypotheses for approximating the softmax denominator were investigated. Linear atlention was shown to be a first order approximation of softmax. Using the recurrent softmax attention formulation, equivalence was shown empirically and the crucial elements of softmax attention were analyzed in an ablation study. Additional experiments showed that tenth order Taylor series approximation is sufficient similar to softmax. Finally, different types of normalization and gates were explored showing that any vector norm is sufficient to approximate the functionality of the softmax attention denominator. The theory developed in this work is crucial to understanding the performance bounds of softmax compared to other forms of attention. Morcover, this theory may be employed to uncover more performant or efficient attention mechanisms.
Limitations The current formulation covers only linear attention and softmax attention, future work can expand it to more complicated recurrent architectures such as RWKV Peng et al. (2025) and state-space models like Mamba Dao and Gu (2024). As recurrent architectures similar to RWKV and Mamba are extensions of lincar attention, the theory developed in this work should extend to the additions made in these works. Only the causal next loken prediction task was investigated. While the derivations should generalize lo other domains, both bidirectional and causal, further investigation is necessary to ensure this generalization. Future work can also incorporate changes from these recurrent models to softmax attention to potentially improve efficiency or performance of softmax attention implementations. Finally, the hyperparameters for the architectures presented were not exhaustively investigated, which could influence the final loss values.
10

---

Published in Transactions on Machine Learning Research (10/2025)
References
Alexei Baevski, Yuhao Zhou, Abdelrahman Mohamed, and Michael Auli. wav2vec 2.0: A framework for self-supervised learning of speech representations. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 12449 12460. Curran Associates, Inc., 2020. URL https://prroeeddngs.. hww.rrooos.dddn............p.... courrps.cc/paper files/paper/ ht0p:///weaaaaeeeeccaaaaaa222660
pdf
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In Yoshua Bengio and Yann LeCun, editors, 3rd International Conference Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. URL http://arxiv. hwnggabss/ee 1409///41..... .0733.
Ali Behrouz, Peilin Zhong, and Vahab Mirrokni. Titans: Learning to memorize at test time. ar preprint arXiv:2501.00663 2024.
Ali Behrouz, Zeman Li, Praneeth Kacham, Majid Daliri, Yuan Deng, Peilin Zhong, Meisam Razaviyayn, and Vahab Mirrokni. Atlas: Learning to optimally memorize the context at test time, 2025. URL
https://arxiv. http://axxivv..rabbbb. argg/bb/2505 23733
Iz Beltagy, Matthew E. Peters, and Arman Cohan. arXiv:2004.05150. 2020.
Longformer:
The long-document transformer.
Kyunghyun Cho, B van Merrienboer, Caglar Gulcehre, F Bougares, H Schwenk, and Yoshua Bengio. Learning phrase representations using rmn encoder-decoder for statistical machine translation. Tn Conference on Empirical Methods in Natural Language Processing (EMNLP 2014), 2014.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Tuey J Colwell, and Adrian Weller. Rethinking attention with performers. In International Conference Learning Representations, 2021.
Liam Collins, Advait Parulekar, Aryan Mokhtari, Sujay Sanghavi, and Sanjay Shakkottai. In-context learning with transformers: Softmax attention adapts to function lipschitzness. Advances in Neural Information Processing Systems, 37:92638 92696, 2024.
Tri Dao and Albert Gu. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality, In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp, editors, Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 10041 10071. PMLR, 21-27 Jul 2024. URL httop::/prrceeddings. mlr. www..prreeet.e..... catt..press/v3335 /23esa//235//ao24a dao24a. .html
Yichuan Deng, Zhao Song, and Tianyi Zhou. Superiority of softmax: Unveiling the performance edge over lincar attention. arXiv preprint ar arXio:2810.1168: 2023.
Alexcy Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image worth 16x16 words: Transformers for image recognition at scale. International Conference on Learning Representations, May 2021. URL
wtwp:ewwwwwwieeeeedeeeeeeeeee mettpadfidd-VicbaNNTY
Patrick Esscr, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, Dustin Podell, Tim Dockhorn, Zion English, and Robin Rombach. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first International Conference on Machine Learning, 2024. URL
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The pile: An s00gb dataset of diverse text for language modeling. CoRR, abs/2101.00027 2021. URL https: btttps//arrxvv.. /tt////wwww.rrg//////oo....
11

---

Published in Transactions on Machine Learning Research (10/2025)
Albert Gu, Karan Goel, and Christopher Re. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Regresentations, 2022. URL Id=uYLFozIvIAC
Dongchen Han, Yifan Pu, Zhuofan Xia, Yizeng Han, Xuran Pan, Xiu Li, Jiwen Lu, Shiji Song, and Gao Huang. Bridging the divide: Reconsidering softmax and linear attention. In The eiihty- Annual Conference on Neural Information Processing Systems, 2024. URL httppp///ppaaaaeees..eeee//ooood
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Newral Compulation, ((8):1735-178)) Generated from Scopus record by
November 1997. ISSN 0899-7667. nei/10.1622/eco..iee KAUST IRTS on 2022-09-14.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are RNNs: Fast. autoregressive transformers with linear attention. In Hal Dauné III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 5156 5165. PMLR, 13 18 Jul 2020. URL http::ttpp/p/e https::/proceeddngs. https//w11/kkaaaarpooooooo html
Zhixuan Lin, Evgenii Nikishin, Xu He, and Aaron Courville. Forgetting transformer: Softmax attention with a forget gate. In The Thirteenth International Conference on Learning Representations, 2025. URL
https://openrewse
Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chenggi Deng, Chenyu Zhang, Chong Ruan, et al. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19497 2024.
Evan Miller. Attention is off by one, 2023. URL http:///wwww.aaall....e.oo/ttttr html
Gabriel Mongaras, Trevor Dohm, and Eric C Larson. Cottention: Linear transformers with cosine attention. Proceedings of the 2025 Computing Conference, 2025.
Tobias Christian Nauen, Sebastian Palacio, and Andreas Dengel. Taylorshift: Shifting the complexity of self-attention from squared to linear (and back) using taylor-softmas In ICPR (6), pages 1-16, 2024. URL
Yuqi Nie, Nam H Nguyen, Phanwadee Sinthong, and Jayant. Kalagnanam. A time series is worth 64 words: Long-term forecasting with transformers. In The Eleventh International Conference on Learning http:/ppweerieeeeeeee/rrrrdddddd
Representations, 2023. URL
Guilherme Penedo, Hynek Kydlíček, Loubna Ben allal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, and Thomas Wolf. The fineweb datasets: Decanting the web for the finest text data at scale. In The Tiihty-eight Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2024. URL https://openrrevier http://ww.eeeooommmmmoodddoorraa
Bo Peng, Ruichong Zhang, Daniel Goldstein, Eric Alcaide, Haowen Hou, Janna Lu, William Merrill, Guangyu Song, Kaifeng Tan, Saiteja Utpala, et al. Rwkv-7" goose" with expressive dynamic state evolution. arXiv preprint ar Xrx:225003. 14456, 2025.
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A. Smith, and Lingpeng Kong. Random feature attention. In 9th International Conference on Learning Represenlations, ICLR 2021, 2021.
Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Lingpeng Kong, and Yiran Zhong. cosformer: Rethinking softmax in attention. In International Conference on Learning Representations, 2022. URL httpps:/ http://ewwweeeeeeeoooootoo//ee/a
Zihan Qiu, Zekun Wang, Bo Zhong, Zeyu Huang, Kaiyue Wen, Songlin Yang, Rui Men, Le Yu, Fei Huang, Suozhi Huang, Dayiheng Liu, Jingren Zhou, and Junyang Lin. Gatod attention for large language models: Non-linearity, sparsity, and attention-sink-free, 2025. hhttp::/aardv..orgabb////a
12

---

Published in Transactions on Machine Learning Research (10/2025)
Noam Shazeer. GLU variants improve transformer. CoRR, abs/2002.05202, 2020. URL https://arxiv
org/ httpa/bbs222.05220
Zhiqiang Shen, Tianhua Tao, Liqun Ma, Willie Neiswanger, Zhengzhong Liu, Hongyi Wang, Bowen Tan, Joel Hestness, Natalia Vassilieva, Daria Soboleva, and Eric Xing. SlimPajama-DC: Understand- ing Data Combinations for LLM Training. arXiv e-prints, art arXiv:2309.1811 September 2023.
doi 1oi100..00// 230.1.2.2.00.811..
Jerome Sieber, Carmen Amo Alonso, Alexandre Didier, Melanic Zeilinger, and Antonio Orvicto. Understanding the differences in foundation models: Attention, state space models, and recurrent neural networks. In The Thirby-cighth Annual Conference on Neural Information Processing Systems, 2024. URL https:
Ethan Smith. Softmax attention is a fluke, 2025. URL http:://ww..ccaassttttt......mm
softmaa-attention attention-is-a-fluss
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced
transformer with rotary position embedding. Newrocomputing, 568:127063, 2024. ISSN 0925-2312.
wttp://c.oo........... 2023.127063 atticile/p//S00222211220181
URL http:::/ww..ccienneelree...
Yu Sun, Xinhao Li, Karan Dalal, Jiarui Xu, Arjun Vikram, Genghan Zhang, Yann Dubois, Xinlei Chen, Xiaolong Wang, Sanmi Koyejo, Tatsunori Hashimoto, and Carlos Guestrin. Learning to (learn at test time):
RNNs with expressive hidden states, 2025. URL https: //oppre////eereiee.oo wttp://r....eerrrrrmroooa/a/oooo
Nya Sutskever, Oriol Vinyals, and Quoc V Le. Sequence to sequence learning with neural networks. Advances in neural information processing systems, 27, 2014.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurclien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. LLaMA: Open and Efficient Foundation Language Models. arXio e-prints, art. 139V February 2023. doi 10.48550/ 2ww.:i.3332111.19.7
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Tukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Maric-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurclien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023. URL https://arxiv org/ hbs///b///// wwwaa///..0.02 09288
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, L ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan,, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https://proceedinng www.eediinnnnnuuissscom BB1/E1tt//ee2255//eeellaaaaaaaa2 Paperarpd
files/paper/
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liô, and Yoshua Bengio.
Graph attention networks. In International Conference Learning Representations, 2018. URL https
7/openreview.co httpa/rrrroo.tttttooormddddddd/
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. Xiv preprint or o00:::006. 04 768, 2020.
13

---

Published in Transactions on Machine Learning Research (10/2025)
Maurice Weber, Daniel Fu, Quentin Anthony, Yonatan Oren, Shane Adams, Anton Alexandrov, Xiaozhong Lyu, Huu Nguyen, Xiaozhe Yao, Virginia Adams, Ben Athiwaratkun, Rahul Chalamala, Kezhen Chen, Max Ryabinin, Tri Dao, Percy Liang, Christopher Ré, Irina Rish, and Ce Zhang. Redpajama: an open dataset for training large language models, 2024. URL petps://arxivy org/ org/ars/ /a///bb//2444 12372
Enze Xie, Junsong Chen, Junyu Chen, Han Cai, Haotian Tang, Yujun Lin, Zhekai Zhang, Muyang Li, Ligeng Zhu, Yao Lu, and Song Han. SANA: Efficient high-resolution text-lo-inage synthesis with linear diffusion transformers. In The Thirteenth International Conference on Learning Representations, 2025.
URL https://openreviee
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in Newral Information Processing Systems, 33, 2020.
Shuangfei Zhai*, Tatiana Likhomanenko* Etai Littwin*, Dan Busbridge*, Jason Ramapuram" Yizhe Zhang, Jiatao Gu, and Josh M. Susskind. Stabilizing transformer training by preventing attention entropy collapse. In ICML, 2023. URL https://arxiv. org/abs/23033 200.b..2990006
Shen Zhuoran, Zhang Mingyuan, Zhao Haiyu, Yi Shuai, and Li Hongsheng. Efficient attention: Attention with linear complexities. In 2021 IEEE Winter Conference on Applications of Computer Vision (WACV), pages 3530-3538, 2021. 1o.110. 1119//99/// SEACV4863)
14

---

Published in Transactions on Machine Learning Research (10/2025)
A Inner Product Decomposition
Let A, E Rd
Define A Rd as the Hadamard product A and B.
Define A B E IR as the inner/dot product of A and B
d
4
Note that the inner product defined as A.B B Σ (AOB), = Σ A,B:
Define A
A = A E as Kronecker products
The lack an explicit operation denotes either: (1) scalar multiplication with another scalar, vector, or matrix. or (2) matrix-vecte vector multiplication, similar "slicing multiplication" in most numerical packages.
Using the notation mention in Section 3.1 show the decomposed form for an inner product space to the th power is equivalent to the inner product of nh order Kronecker products of each vector.
(A Byn
(AO B);
By equation equation
E(AO B);
[(AOB)
Σ O B):
Power to product of terms
(), (AO B)₁
By rearranging the sums
combinatorial multiplication from
(A O B) 8mn
By equation equation
[AA00) o (B*"),
Rearrange the nth order Kronker Product
(7)
= (10m). (B8n)
Change to inner product by equation equation 2
(8)
15

---

Published in Transactions on Machine Learning Research (10/2025)
B Quadratic Derivation
To provide simple example of the recurrent form derived in Section 3.1 and show the recurrent form.
note that A A A E R d
0,
Σ
Q+
KT)
2
QQ K2),
V.
By equation equation
&
Σ
Q1 Q1), (K, K.),
By equation equation
1=
d
Q 8 Q.), (K, 2 V.
rearrange summations
==
Σ QQ 21, Σ (K, @ K.), V.
By factoring out Q
(Q 8 Q.)
(KK)) @ (K?)) V.
By equation equation 2
(Q, 0 Q.)
H. Σ (( @ (K!) V, V. Rd
Define hidden state
C Bidirectional Derivation
N
0: G. Σ
9.-K V.
G,
Q KV)"V.
By definition of the Taylor Series of
N
G Σ
71 (2*0), ((K
By equation equation
0=
G.
(Q""), ((KSH)T), V.
By rearranging sums
n=0
G,
Σ ((K@N)?), Vn
By factoring out Q
G. n
(K V,
By equation equation 2
G.
H. = Σ (K V, V
Define hidden state
16

---

Published in Transactions on Machine Learning Research (10/2025)
The only difference between the causal and bidirectional formulations is the range the sequence sum indexed by s. In the causal case, the index ends at f while the bidirectional case sums to the end of the sequence. This means the hidden state operates on the entire sequence for each token rather than just the past
D Model Parameters
Unless otherwise mentioned, the below are the parameters we used in our models. As our base model is llama
Touvron et al. (2023). RoPE Su et al. (2024] used on the attention matrix and the MLPs follow SwiGLU
Shazeer
20200]
batch size 36
2. learning rate le-4
warmup steps 10,000
4. warmup type linear warmuu from 0, linear decay
mum steps 100,000
6. precision float32 and bfloat16 mixed precision
Weight decay 0.01
8. Max sequence length 1024 for general experiments, 4096 for length scaling experiment
Test percentage 0.001
10. Optimizer AdamW
11. Adam betas 0.9 and 0.999
12. Hidden size- 1024 (3072 for the large model)
13. MLP intermodiate size 2048 (6144 for the large model)
14. Num attention heads 16
15. Num hidden layers 20
16. Tokenizer llama2-7b-hf
17. Gradient clipping clipping for gated models, no clipping for all other experiments
Each model was trained for maximum of days. For most experiments, use distributed data parallel processing to train two 80 GB, A100 GPUs with the exception of the large model, trained on 4 GPUs, and 4096 sequence length, trained on 6 GPUs.
E
Additional Gate Information
The input and output gates referenced in this paper are very similar to that of an LSTM Hochreiter and Schmidhuber (1997) and similar to that used by Qiu et al. (2025). The input gate controls how much information is being added to the LSTM hidden state while the output gale modulates the output of the LSTM cell. As seen in Figure 7 these gating mechanisms can be translated to the original linear attention formulation where the input gate is scalar between [0, 1) at time t, modulating the KTV outer product while the output gate is a scalar between Jo, 1] at time modulating output after the Q,H, inner product. These ideas extend to max attention, where it be computed infinite of these RNNs RNNs-although the
17

---

Published in Transactions on Machine Learning Research (10/2025)
implementation is intractable. However, the gates can be translated to multiplicative values on the QK? attention matrix. As the rows/queries define the output at time 4, the output gate can be defined along the rows. Similarly, as the columns define the input at time the input gate can be defined along the columns.
This idea equivalent to applying the output gate after performing the full attention operation and applying the input gale the values. Mathematically, this can be expressed as follows:
O2
Σ
Gi Cieeeuk Qukk Gouut
=G
Σ e
Cook" [G
V.]
M O Gout 6 3 QOKT
V G [eek" M
Gour V
Where Q e R.. K M.d, V RM, G [0, 11N, Gout E [0, 1j
And Qi Rd, K, Rd V R , (0, G) [0.1]
GPaet
Inner Prod
Hj
(K)
Outer Prod
V.
Figure 7: Linear attention an RNN with an input and output gate.
18

---

Published in Transactions on Machine Learning Research (10/2025)
F Linear Attention Expansions
This section provides details on how linear methods are augmented from their base implementation, presented in Section 4.3 For linear attention variants the activation functions, & and v, are applied to Q and as in the native linear attention.
(Softmar)
O₂ =
Σ n!
(QP")
(( V
(Q: KV),V,
n!
n=0
#=1
8=1 R=0
(Linear)
O₂ n=
(o(2)) (00)
((b(K) V.
(o(Q), (K))
n!
==1
19

---

Published in Transactions on Machine Learning Research (10/2025)
G Expanded Figures
Fineweb Log Test Loss
Fineweb Train Loss
Saftme
Cate
Norm
Softre
1.3
1.5
12
1.15
Selyrac
1.0
Step
0.9
20k
401
60k
80k
20R
45k
65k
D0X
SlimPajama Log
SlimPajama Log Train
Sotteas
Gate
Kome
Softres
1.35
1.5
Worn
122
1.15
Step
0.9
20k
Log Test
Log Train
Gold
Marn Solliion
Cola
Hores
Solliun
1.
12
12.
11
08
3tep-
Stop
Figure 8: Expanded test and train loss plots on various datasets for softmax attention and the proposed methods with gate or norm replacements. (expanded Figure]
20

---

Published in Transactions on Machine Learning Research (10/2025)
Fineweb Test Loss
Fineweb Train
Dosoin Sulunee Derien Gate
Sofreas Detach Denon:
Softrax- Drach Desoi Sote
Det
Sulmace
DVd Site Sertmact Senors -orm
1.
166
Softno Dract Dovon
14
13
Sottwac Denore DiS
12
1.2
Suftmee Deners DvS
Softrag No becorr Norm
Sattnar
Subbrae to Sencn 10 as tate
Sultmes Ceroer Bud
Stap
Settmer Ho Denzm:
Step
0.8
20k
40k
10x
Log
Fineweb Log
Softmax
Cate
Did
Gate
Be feen faus
Softeeu Eato
DV-
reu Gate Did
Gew DNS
Uneor 6
14
Liner U (atc--at---- Did
1.2
Unee falll Gaa D&S
Insil Bad
Cate Bid
Selmma
Step.
Fineweb Log Test Loss
Fineweb Log Train Loss
Sultrec
12 kom Exp
Horn Ducks RVS Horn
Laver Norm Do
Namn RAUU
Sohrras
12:00m Bx
homn Der DNS
RyShom Cep Layer Nomm Eeo
13 Warm Re.U)
166
144
LE Mara bo
L2 Norn to
boooks
1.3
195 Ex
Lowel Marn Exp
12
Syttrax
12 ReUU
1
1.1
20k
40c
50k
0
20k
40%
00k
80x
100k
Figure 9: Expanded test and train loss on various datasets for softmax attention and the proposed methods with gate norm replacements. (expanded Figure G
21

---

Published in Transactions on Machine Learning Research (10/2025)
Softmax attention an RNN, We define G. in place for the softmax denominaton Linear attention is
Kronedter Frod (0uBe
Irase Phod
51
Higha order m0)
Doome W
Kronakeer Prod 0kg
6
122
H4
Kroreccke 60
Ineel Rood
M
Kooeecker Frod 0igiit
Figure 10: Expanded version of Figure [ equivalent to the = 1, first order, term.
Linear Attention
Cader-ra Phoo
Lival mer
Mos
22


Reviews:
[{"review_id": "qvqF8exnam", "reviewer": "Reviewer_LbKq", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper derives a recurrent reformulation of softmax attention via a Taylor-series expansion of the exponential in the QK inner product, yielding an interpretation as an infinite sum of RNNs with growing-order hidden states. From this lens, it argues (and demonstrates empirically) that (i) standard “linear attention” corresponds to the first-order term and is therefore less expressive, and (ii) the softmax denominator primarily serves as a stabilizing factor that can be replaced by a simple vector norm (with a gate being possible but less stable). Experiments on next-token LM with LLaMA-like models (~300M and ~2B) trained on The Pile, SlimPajama, and FineWeb suggest the proposed “normed” formulation matches softmax training curves, while linear-attention variants lag; adding higher-order terms (up to n=10) smoothly bridges the gap toward softmax.", "claims_and_evidence_yn": "Yes", "claims_explanation": "Overall, the paper presents substantial analysis and derivations; although some assumptions are occasionally loose, the conclusions appear largely sound.", "audience_yn": "Yes", "audience_explanation": "Unifying perspective: Clear derivation that places softmax and linear attention in one framework, making the “first-order approximation” view explicit and intuitive.\n\nEmpirical support: Multiple datasets and two model scales; scaling experiments (longer context, larger model) show the normed variant tracks softmax closely; thorough ablations (Taylor order, gating vs norms, linear-attention variants)\n\nDidactic exposition: The RNN view and Kronecker-product formulation help reason about “higher-order interactions” and why linear attention underperforms.", "requested_changes": "Task breadth is limited: Results focus on next-token LM loss curves; there are no standard perplexity benchmarks, downstream tasks, in-context learning evaluations, or long-range reasoning tasks to confirm functional parity with softmax beyond train/test loss.\n\nDenominator “norm” claim needs stress tests: The observation that a simple vector norm suffices is shown on a few setups; stability/robustness across wider context lengths (e.g., 32K+), architectures, and domains (vision, speech) remains unclear. The gate variant’s instability is acknowledged but not deeply analyzed.\n\nFor journal papers, it’s best to use vector graphics for figures.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "02JPeHKvbV", "reviewer": "Reviewer_7CMB", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper presents a novel theoretical analysis of the softmax attention mechanism from the perspective of recurrent neural networks. The primary contribution is the derivation of a recurrent formulation for the numerator of the softmax attention calculation. This is achieved by applying a Taylor series expansion to the exponential function, which reframes the attention mechanism as an infinite weighted sum of RNN-like components. It provides a principled explanation for the performance gap between softmax and linear attention, showing that linear attention is equivalent to only the first-order term of this infinite series, thus lacking the capacity to model higher-order feature interactions. This paper also deconstructs the expressiveness of softmax attention into an ensemble of recurrent structures, each capturing progressively higher-order interactions between queries and keys. The authors also recognize that the softmax denominator does not lend itself to a similar recurrent derivation. They hypothesize that its primary function is normalization and stabilization, rather than its exact mathematical form being crucial. They test this by replacing the denominator with simpler mechanisms, most notably a standard vector norm.\n\nPros: The theoretical connection established between softmax and an infinite ensemble of RNNs is novel, providing a clear and unifying perspective for understanding one of deep learning's most important mechanisms. The work offers clear explanations for why softmax attention is more expressive than its linear counterparts. Extensive experiments on multiple large-scale datasets and model sizes support the claims. The finding that a simple L2 norm can replace the softmax denominator with virtually no performance degradation is a significant and noteworthy result.\n\nCons: The paper's primary weakness is that it does not derive a complete recurrent equivalent for the entire softmax mechanism. The treatment of the denominator is a functional approximation, not a mathematical derivation. While the authors are transparent about this in the methodology section, the title and abstract could be interpreted as claiming a full, exact equivalence, which might be misleading.", "claims_and_evidence_yn": "Yes", "claims_explanation": "The paper's claims are well-supported, with evidence of two distinct types that the authors handle appropriately.\n1. The central claim that the softmax numerator can be expressed as an infinite sum of recurrent components is supported by a rigorous and sound mathematical derivation in Section 3.1 and Appendix A. The steps, involving a Taylor series expansion and the application of an identity involving Kronecker products, are mathematically correct.\n2. The second major claim—that the denominator's function is primarily normalization and can be replaced by a vector norm—is supported by extensive and convincing empirical evidence. Figure 2 shows that the \"Norm\" variant's loss curves are nearly indistinguishable from the standard Softmax baseline across three different large-scale datasets. Furthermore, Figure 3 demonstrates that this equivalence holds when scaling up both the model size and the sequence length. Finally, the ablation analysis in Section 4.5 and Figure 6 further strengthens this claim by showing that various types of vector norms perform well, while removing normalization entirely or using improper gating leads to instability or performance degradation.", "audience_yn": "Yes", "audience_explanation": "This research addresses a high-impact question at the heart of modern deep learning: \"What makes softmax attention so effective?\" The work provides a new theoretical foundation for understanding the expressiveness of the core component of Transformer models, making it highly relevant to Transformer and Attention Researchers. Moreover, the paper builds an explicit bridge between the Transformer and RNN worlds. Researchers working on state-space models and recurrent architectures like RWKV, Mamba, and other SSMs will likely be interested in this unifying perspective.", "requested_changes": "I suggest the authors slightly revise the language in the Abstract and Introduction to more precisely reflect the distinction between the exact derivation for the numerator and the functional approximation for the denominator. For example, instead of \"This work demonstrates that linear attention is an approximation of softmax attention by deriving the recurrent form of softmax attention,\" a more precise phrasing might be: \"This work demonstrates that linear attention is a first-order approximation of the softmax numerator by deriving its full recurrent form. We further show empirically that the denominator's function can be effectively replaced by a simple vector norm.\"\n\nIn Section 3.1, the dot product symbol is used to represent several different operations. While context makes it understandable, explicitly defining the operations or using different symbols in the final recurrent formula would improve mathematical clarity.\n\nSome labels and lines in the figures are small and can be difficult to distinguish.", "broader_impact_concerns": "NA.", "additional_comments": "NA.", "extra_scores": {}, "extra_text": {}}, {"review_id": "cB3GEhpslz", "reviewer": "Reviewer_PB6Z", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper provides both theoretical and empirical insights into why softmax attention is more expressive than linear or kernelized variants. By deriving a recurrent formulation via Taylor expansion, the authors show that linear attention corresponds to the first-order approximation, while softmax can be interpreted as a sum of infinite recurrent chains that capture higher-order query–key interactions. They further analyze the role of the denominator, proposing gating and normalization perspectives, and validate these hypotheses through experiments across multiple datasets and model scales. The results suggest that softmax’s superiority arises from its capacity to model higher-order interactions and the stabilizing effect of normalization. In addition, the authors train several LLaMA-2 variants (300M to 2B) to support their empirical findings.", "claims_and_evidence_yn": "Yes", "claims_explanation": "Strength：\n1. The paper addresses a meaningful and timely question.\n2. It provides both theoretical analysis and empirical validation of the differences between softmax and linear attention, with solid and comprehensive content.\n\nWeakness\n1. Although the experiments align learning rates across LLaMA-2 variants for fair comparison, prior work suggests that the optimal learning rate for linear attention may differ from that of softmax. Under this setting, can the reported results truly reflect the best achievable performance of linear attention?", "audience_yn": "Yes", "audience_explanation": "The paper provides detailed theoretical proofs and comprehensive experiments to substantiate its claims", "requested_changes": "1. Could you discuss the impact of hybrid parameters on the experimental validation?", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)