MetaLA: Unified Optimal Linear Approximation to Softmax Attention Map

Authors: yuhong  chou
man  yao
kexin  wang
yuqi  pan
ruijie  zhu
yiran  zhong
yu  qiao
jibin  wu
bo  xu
guoqi  li
Venue: NeurIPS_cc_2024_Conference
Published: 2024-04-24 16:49:56

Abstract: Various linear complexity models, such as Linear Transformer (LinFormer), State Space Model (SSM), and Linear RNN (LinRNN), have been proposed to replace the conventional softmax attention in Transformer structures. However, the optimal design of these linear models is still an open question. In this work, we attempt to answer this question by finding the best linear approximation to softmax attention from a theoretical perspective. We start by unifying existing linear complexity models as the linear attention form and then identify three conditions for the optimal linear attention design: (1) Dynamic memory ability; (2) Static approximation ability; (3) Least parameter approximation. We find that none of the current linear models meet all three conditions, resulting in suboptimal performance. Instead, we propose Meta Linear Attention (MetaLA) as a solution that satisfies these conditions. Our experiments on Multi-Query Associative Recall (MQAR) task, language modeling, image classification, and Long-Range Arena (LRA) benchmark demonstrate that MetaLA is more effective than the existing linear models.

Main paper: 
 introduction: 
 Transformer with softmax attention [1] benefits from efficient parallel training and exhibits impressive performance on deep learning applications [2, 3, 4, 5, 6, 7] , but it suffers from the quadratic growth of computation cost to the input length [8] . Linear recurrent models, such as LinFormer [9] , SSM [10] , and LinRNN [11] , are expected to achieve linear substitution of Transformer. The original intention of LinFormer is to replace softmax attention, which exploits the kernel approach to decompose softmax operation; typical work includes TransNormer [12, 13] , RetNet [14] , GLA [15] . On the other hand, SSMs, such as S4 [10] and Mamba [16] , are models inspired by the classical state-space approach, which enjoys sub-quadratic training and inference like either a recurrence or convolution. In contrast, LinRNN is a revival of traditional RNNs, including RWKV-4 [17] , Griffin [18] , LRU [19] , etc., which solves the training difficulties of traditional RNNs due to nonlinear dependencies between hidden states. It is natural to think that they are different types of models, since these LinFormer/SSM/LinRNN models have different origins and forms. This work breaks this perception and abstracts existing LinFormer/SSM/LinRNN models into a unified linear attention form, which has the following significance: i) Facilitates understanding the key designs of existing linear models. Through the unified form, we demonstrate that the main difference between LinFormer/SSM/LinRNN is the hidden state size, how to maintain the hidden state, and how to perform parameter mapping. ii) Links LinFormer/SSM/LinRNN to softmax attention in terms of functionality. The recurrent inference complexity of softmax attention is O(n), which can also be regarded as the maintenance of a hidden state with infinite size. Linear models with O(1) inference complexity are hoping to achieve the same functionality as softmax attention using a fixed hidden state. Since we have unified LinFormer/SSM/LinRNN into linear attention in the form of Query, Key, and Value, we can understand and evaluate existing linear models from the view of "Does the linear attention map have the function of softmax attention map?".To answer this question, we define the necessary conditions for achieving "optimal linear approximation to softmax attention map". First, linear attention needs to satisfy dynamic memory and static approximation to realize the approximation. The former defines memory ability: linear attention with limited hidden states should be able to store the most important information and forget unimportant ones. The latter defines the modeling ability: a linear attention map should be able to approximate any softmax attention map. According to our theoretical analysis, Query and dynamic decay are necessary conditions for approximation. Thus, linear models such as TransNormer [13] , RetNet [14] , RWKV-4 [17] , LRU [19] , HGRN [20] , H3 [21] , S5 [22] , cannot achieve approximation of the softmax attention functions. Second, the Key matrix is not required to achieve approximation, so Mamba [16] and GLA [15] are not optimal parametric approximations.We then propose the MetaLA module, which can satisfy the necessary conditions for optimal linear approximation to softmax attention. MetaLA makes three enhancements: i) Removes the unnecessary Key matrices; ii) Employs self-augmentation to enhance the token's attention to itself, which avoids attention dilution [12] ; iii) Exploits short convolutions to enhance local interactions. We then build a MetaLA Transformer based on MetaLA. Our experiments on associative recall, language modeling, long sequence modeling, and image classification show the effectiveness of MetaLA. Furthermore, we conduct ablation studies to validate the effectiveness of each proposed enhancement in MetaLA. Finally, we discuss two open questions: i) How to further improve linear attention based on the approximation theory introduced in this work? ii) Does the approximation of linear attention to softmax attention imply that it has an upper limit on its capacity? 
background: 
 For notations in this work, we use bold upper-case letters for matrices (e.g., Q, K), bold lower-case letters for row vectors (e.g., q t , k t ), and italic upper-case for learnable parameter matrices (e.g.,W Q ). We generally use the same alphabet to show the rows of a matrix, e.g., q t is the t-th row of Q.Softmax Attention first calculates an attention map SoftAttMap (Q, K) through Q (Query), K (Key), and use the attention map to weight different tokens V (Value) later:EQUATIONEQUATIONwhere W Q , W K ∈ R d×d k , W V ∈ R d×dv are learnable matrices, n, d, d k , d v are sequence length, model dimension, Key/Query and Value dimension, respectively. X ∈ R n×d refers to the input. M ∈ R n×n is a mask matrix in autoregressive tasks to prevent a token from attending to future tokens. The t-th row of SoftAttMap (Q, K) is a probability distribution that represents the attention scores between token v t to others. Softmax attention in Eq. ( 1) enables efficient parallel training, but suffers from O(n 2 ) time and memory complexity [9] . It uses the recurrent form during inference:EQUATIONEQUATIONAt each time t, token mix is computed between query q t and all the keys, values before k s , v s (s ≤ t). This "KV cache" results in O(n) time and memory complexity per token during inference.Linear Transformer (LinFormer) is a substitute for softmax attention, which can be expressed as a linear dot-product of kernel feature maps [9] :o t = t s=1 F (q t , k s )v s t s=1 F (q t , k s ), F (q t , k s ) = ϕ(q t )ϕ ⊤ (k s ), (5) where q t , k t ∈ R 1×d k and v t ∈ R 1×dv are query, key and value at position t, which are obtained in the same manner as softmax attention. F (•) is the kernel function usually constrained to be non-negative. ϕ(•) is map function applied row-wise to Q and K. By removing the nonlinear softmax operation, LinFormer enables inference with O(1) time and memory complexity per token. LinFormer can also be formulated in the following parallel form during trainingEQUATIONwhich has O(n) time and memory complexity using chunkwise algorithm. Recent advances in LinFormer mainly focus on training acceleration [12, 13, 23] or improving performance [14] .State-Space Model (SSM) come from continuous-time system which maps a 1D function x(t) ∈ R to another function y(t) ∈ R via a hidden state h(t) ∈ R N . In SSM, the continuous parameters can be discretized using a step size ∆ and get discrete parameters A, B, C. The resulting discrete-time system is used to model sequences x, y ∈ R 1×n with elements x t , y t ∈ R via the recurrent form:EQUATIONin autoregressive inference with O(1) time and memory complexity per token. The linear timeinvariant SSM above can be unrolled and computed using the long convolution with kernelEQUATIONwhere * represents casual convolution operation [24] , which enables parallelizable training utilizing Fast Fourier Transforms, resulting in O(n log n) time and O(n) memory complexity during training. When handling vector sequences X, Y ∈ R d×n , SSMs are applied individually on the d channels.Typical SSMs (S4D [25] , DSS [26] , H3 [21] , S5 [22] ) employ data-independent structured transition matrix A or special initialization strategies HiPPO [27] to efficiently enhance long-range dependencies. Mamba [16] advances SSMs by introducing data-dependent parameters and designs a hardware-aware parallel algorithm, further improving prior O(n log n) into O(n) time complexity during training.Linear RNNs (LinRNN) Traditional RNNs suffer from slow sequential training, limited capability in modeling long-term dependencies, and difficulty in scaling. To address these issues, LinRNNs eliminate the nonlinearity within the recurrence and employ element-wise product instead of matrix multiplication [11, 28] . Typical LinRNN such as Gated Linear Recurrent Unit (GLRU) [20, 29] isEQUATIONEQUATIONwhere x t , h t denote input and output, f t , i t are forget and input gates as in traditional RNNs, ⊙ is element-wise multiplication. Linear RNNs have O(1) time and memory complexity per token during inference. Since Eq. ( 10) removes nonlinearity, it enables parallelized training using parallel scan [11] , with only O(n) time and memory complexity. Recent works have made effort to explore effective recurrence (LRU [19] , RWKV [17] ) or gating mechanisms (HGRN [20] , Griffin [18] ).3 General Form of LinFormer/SSM/LinRNN Mechanisms Observing Eq. (3), Eq. ( 5), Eq. ( 7), and Eq. ( 10), we find that their recurrent forms during inference can all be understood from the view of maintaining hidden states. Softmax attention maintains an unlimited hidden state (KV cache). By contrast, LinFormer/SSM/LinRNN have limited hidden states: linear attention with ϕ ⊤ (k t )v t ∈ R General Recurrent Form of LinFormer/SSM/LinRNN is:EQUATION)EQUATIONEQUATIONEQUATIONEQUATIONwhere x t ∈ R 1×d is the t-th input. q t , k t , v t , α t , g t , S t are query, key, value, decay, output gate, hidden state respectively. f q/k/α are functions that map x t from R 1×d to R 1×d k , θ q/k/α are the corresponding parameters to be trained. Similarly, f v/g map x t from R 1×d to R 1×dv and θ v/g are trainable parameters. In Eq. ( 13), q t , α t , k t , v t are divided into H partitions (heads), whered ′ k/v = d k/vH , and h = 1, • • • , H is the index of heads. Each head maintains a hidden state S h t . The diagonal matrix diag(α h t ) denotes the decay of past state. k h t represents the acceptance for the input token v h t . The hidden states are 2D matrix once d ′ k ̸ = 1. In Eq. ( 14), to turn back to 1D shape, the q h t operation is necessary as a dot-product with S h t , then the concat and normalization operations are followed. XNorm denotes any kinds of normalization. In Eq. ( 15), a gate machanism is optional for o t while the dimension should be projected back tod from d v through W O ∈ R dv×d .From a functional view, Eq. ( 13) represents the update process of the hidden state, which contains historical information on keys and values. Eq. ( 14) represents query and weighted sum operations to derive the attention output. Eq. ( 15) represents gate and projection operations to get the final output.General Parallel Form of LinFormer/SSM/LinRNN can be written as follow:EQUATIONEQUATIONK A denotes element-wise division, (Q) t,: is the t-th row of Q, and LinAttMap (Q, K) is the attention map. Each element in the attention map matrix is as follows (heads are omitted for simplicity):EQUATIONAs shown in Tab. 1, the main differences between various linear models are parameter functions f q/k/v/α/g and dimension settings d k , d v , H. We give details in appendix A1 on how to derive LinFormers/SSMs/linRNNs from our unified form, which is termed as "Linear Attention (LinAtt)".LinFormer/SSM/LinRNN models have different origins, so they have different optimization perspectives and various hidden state sizes: i) LinFormer originate from approximation of vanilla softmax attention. They focus on designing better kernel function ϕ, i.e., to optimize f q , f k ; They have Table 1 : From our general form to existing linear models ( * indicates the bias term is omitted).LinFormer LinRNN SSM GLA [15] TrNorm [12] GLRU [20] RWKV-4 [17] Mamba [16] S5 [22] f q (x t , θ q ) ) is relatively small. Despite these differences, they all try to design better parameter functions f q/k/v/α/g and maintain a limited hidden state S t .x t W Q ϕ(x t W Q ) 1 1 x t W C 1 f k (x t , θ k ) x t W K ϕ(x t W K ) σ(x t W i ) * exp (x t W K ) ∆ t (x t W b ) 1 f v (x t , θ v ) x t W V x t W V ϕ(x t W c ) * x t W V x t x t B f α (x t , θ α ) σ(x t W 1 α W 2 α ) * λ exp (jθ) σ(x t W f ) * exp (-W ) exp (∆ t A) exp (∆A) f g (x t , θ g ) σ(x t W r ) * x t W U ϕ(x t W g ) * σ(x t W r ) 1 1 Dimension d k = d/2 d v = d d k = d v = d d k = d v = d = H d v = d = H d ′ k = N d v = N d = H d ′ k =Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: In Sec. 7, we discussed potential concerns about the approximation of softmax attention and the reasons that may cause these concerns.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
theory assumptions and proofs: 
 Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
optimal linear approximation to the softmax attention map: 
 We here discuss the optimal approximation of LinAttMap to SoftAttMap based on its general form. The function of softmax attention is two-fold: i) Memorizing information, all the current and historical information can be stored in KV cache; ii) Modeling relationships, softmax attention can calculate arbitrary attention scores of stored information. Unfortunately, such a powerfully expressive attention map generated by softmax attention requires infinite hidden states. By contrast, linear attention expects to exploit limited hidden states to achieve the same functionality as softmax attention.Some existing linear models, such as Performer [30] , RFA [31] , etc., optimize the model with the goal of approximating the value of SoftAttMap. In contrast, this work investigates the functional approximation of SoftAttMap, which is the basis for value approximation. Specifically, we here attempt to answer two key questions: i) Can linear attention realize the function of softmax attention? ii) If it can be achieved, what kind of linear attention approximation is better? To achieve this goal, we first give the definition of necessary conditions of optimal linear approximation. Then we categorize the existing linear models based on the conditions of the optimal linear approximation. For the C2 condition, Eq. ( 18) illustrates the LinAttMap only relate to q t = f q (x t , θ q ), k t = f k (x t , θ k ), α t = f α (x t , θ α ). Denote decay Λ t = diag(α t ). Assuming the inputs are good enough Table 2 : A review of existing linear models. According to definition 4.1, existing linear models all have some deficiencies: i) Models without dynamic decay Λ t have no ability to memorize dynamically (not satisfying C1); ii) LinRNNs lack the selection ability brought by Q (not satisfying C2), and the approximation ability is poor owing to the small hidden state; iii) Models with K have redundant parameters (not satisfying C3), which probably leads to higher learning difficulty.Softmax Attention -Transformer [1] , GPT [2] , Llama [3] Linear Attention Three Parameter Groups Q, K, Λ RetNet [14] , TransNormer [12] , S4D [25] , H3 [21] , DSS [26] Q, K, Λt GLA [15] , Mamba [16] Two Parameter Groups Q, K linear Transformer [9] , RFA [31] , Performer [30] , cosFormer [32] EQUATIONΛt HGRN [20] , Griffin [18] , T-RNN [33] and the functions (f q , f k , f α ) are expressive enough, we can shift from solving (θ q , θ k , θ α ) to solving (Q, K, Λ t ). We focus on approximating the attention scores between stored tokens, and the problem is simplified via: i) setting query dimension d k = 1; ii) considering only a given time t and its attention distribution p t = [p ts , s = 1, . . . , t] ∈ R 1×t . Then, C2 is proved by the following equations holding with bounded parameters, as a foundation conclusion:EQUATIONEQUATIONFor bounded inputs X, bounded parameters (θ q , θ k , θ α ) are equivalent to bounded (Q, K, Λ t ). Afterwards we will generalize to vector version with d k > 1 and consider distribution of all time (p t , t = 1, . . . , d k ). This is done by viewing q t as a channel selector.For the C3 condition, least parameters mean the fewest parameter groups (Q, K, Λ t ) when d, d k , d v are fixed. Due to space constraints, the detailed analysis in this Section is provided in appendix A2.Conclusions of Optimal Linear Approximation Analysis. i) Linear approximation. The necessary conditions (C1 and C2) for LinAttMap to achieve approximation to SoftAttMap is that its implementation must include Q and dynamic decay Λ t . Both (Q, K, Λ t ) and (Q, Λ t ) can achieve approximation. ii) Least parameter approximation. (Q, Λ t ) has fewer parameters (i.e., K is not necessary), if the model dimensions are fixed. iii) Function of dynamic decay. Λ t is the key to achieve dynamic memory. iv) Function of Query. Q can be seen as a channel selector which selects several channels of Hadamard product of Λ t and K to approximate attention map.In Tab. 2, we review some existing linear models and judge whether they meet the necessary conditions for optimal approximation. Linear attentions can be classified into three types based on the parameter groups: i) Using (Q, K, Λ t ) all together, ii) Exploiting (Q, K), (Q, Λ t ) or (K, Λ t ), iii) Employing only one of Q, K, Λ t .Considering decay can be either dynamic or fixed, here we use subscript t to distinguish, i.e., Λ/Λ t denote fixed/dynamic decay. According to definition 4.1, they have different degrees of deficiencies: i) Models without dynamic decay such as RetNet [14] , TransNormer [12] , RFA [31] , cannot memorize dynamically; ii) LinRNNs such as RWKV-4 [17] , HGRN [20] , Griffin [18] Table 3 : From general recurrent linear form to our MetaLA.Models f q (x t , θ q ) f k (x t , θ k ) f v (x t , θ v ) f α (x t , θ α ) f g (x t , θ g ) Dimension MetaLA x t W Q 1 -α t x t W V σ(x t W α ) ϕ(x t W G + b G ) d k = d/2, d v = dlack the selection ability brought by Q and the approximation ability is poor due to the small hidden state; iii) Models with K such as Mamba [16] , GLA [15] have redundant parameters, which probably leads to higher learning difficulty. Thus, none of the existing linear models meet all C1/C2/C3 conditions. These analyses also inspire us that ignoring the functional approximation of softmax attention does not enable the approximation of softmax attention values. 
metala transformer: 
 Transformer is stacked by a series of Encoder/Decoder blocks. Generally, each block is composed of two modules in sequence: token mixer and channel mixer [34, 35] . Softmax attention plays the role of the token mixer. In this work, we follow the Transformer architecture as a whole but use our proposed MetaLA module as the token mixer. Due to space constraints, the architecture of MetaLA Transformer is given in detail in appendix A3. Here we focus on describing the three enhancements of MetaLA relative to the general linear attention in Sec. 3 (see Fig. 2 ): i) The Key matrix is not used, which is based on our theoretical analysis. ii) Self-augmentation and iii) Short convolution are two other optional techniques to further enhance the modeling ability of our model. i) The Key matrix is not used. We exploit 1 -α t to replace k t , based on theoretical analysis in Sec. 4 and appendix A2, i.e., dynamic decay Λ t is the key mechanism to achieve dynamic memory and static approximation, and K is not required. As shown in Tab. 3, compared with Eq. ( 13), the main improvement is:EQUATION)which can be trained in a parallel form in Eq. ( 16). The only difference is that K is replaced by B and B t,: = 1-α t . With usage of Λ t and Q, MetaLA can achieve linear approximation to SoftAttMap. Without usage of K (W K ), we can allocate more parameters and utilize full-rank matrix W α to produce dynamic decay rather than low-rank matrix used by GLA, such that we do not sacrifice expression capacity of f α and approximation performance of MetaLA.ii) Self-augmentation can enhance a token's attention to itself, avoiding attention dilution [12] :EQUATION)Without changing the hidden state S h t in Eq. ( 21), the proposed self-augmentation (the second term on the right side of the equation) is only added on the output process, with a learnable parameter w aug ∈ R 1×d k . The proposed design has two advantages (more analysis in appendix A3.2): First, it maintains parallel computing; Second, it augments the information of current token itself and does not affect future output through inner state. iii) Short convolution. An additional short convolution can be inserted before entering the MetaLA layer to enhance local interaction further, motivated by Mamba [16] and Griffin [18] . 
experiments: 
 We conduct a comprehensive evaluation of MetaLA to validate its capabilities as a foundation model. i) MQAR [36] . Performance on the Multi-Query Associative Recall (MQAR) task is closely linked to language modeling and can also imply the effectiveness of our theory in modeling hidden states and retrieving information. ii) Autoregressive language modeling on the Pile [37] dataset and evaluation on Common-Sense Reasoning and SuperGLUE [38] zero-shot benchmarks are conducted. iii) LRA [39] . We execute experiments on the Long Range Arena (LRA) benchmark [39] to investigate MetaLA's ability in long sequence modeling. iv) ImageNet [40] . Generalization ability in visual tasks. Due to space constraints, we put some additional experiments in appendix A5, including: v) Scalability. We extend MetaLA to a 3B parameter scale and a 300B data scale for preliminary validation. vi) Retrieval and long context abilities. We evaluated MetaLA's retrieval performance on the MAD tasks [41] , and its effectiveness in handling long contexts on the Needle in a Haystack task [42] . vii) Training efficiency. We provide comparative results on training throughput and GPU memory usage across various models. Detailed experimental setup and further discussion are given in appendix A4.Associative Recall. The synthetic MQAR task [36] is exploited to evaluate MetaLA's memory ability.In the task, given multiple queries, the model must recall the corresponding key-value pairs before. We follow default settings in [36] to generate datasets. Fig. 3 shows that MetaLA outperforms other linear models, which have three parameter groups (Mamba [16] , GLA [15] , Based [43] ) or fixed decay (RWKV-4 [17] ), well supporting our theoretical analysis and module design. The attention baseline achieves optimal results (> 99.0) under both conditions. The additional experiments in appendix A5 show that MetaLA outperforms Mamba on more challenging settings.Language Modeling. We train two scales of MetaLA: 360M/1.4B on the Pile dataset. For baselines of 360M MetaLA, we train them from scratch aligned with our configurations. For the 1.3B MetaLA, we compare it with publicly available models [14, 15, 16, 20, 44] . We implement all the pre-train experiments with GPT-Neox [45] . The zero-shot evaluation results on SuperGLUE and Commensense Reasoning benchmarks are reported in Tab. 4 and Tab. 5. Specifically, compared to the LinRNN model HGRN [20] , MetaLA expands hidden state dimensions and uses Query matrix; Compared to LinFormer model RetNet [14] with fixed decay, MetaLA uses dynamic decay; Compared to SSMs like Mamba [16] and LinFormer with dynamic decay GLA [15] , MetaLA omits the Key matrix in computation. Results indicate that MetaLA has better performance than these linear models and the Transformer-based Pythia [44] . See appendix A5 for more task results. Long Sequence Modeling. LRA is used to evaluate the model's ability in long sequence modeling. We compare MetaLA with Transformer, linear foundation models, and models specifically designed for long sequence modeling. Tab. 7 shows that MetaLA achieves comparable results with SOTA linear models, demonstrating that our model effectively preserves the ability to model long sequences.Image Classification. We compare MetaLA with Deit [49] (Transformer), HGRN [20] (LinRNN), GLA [15] (LinFormer) and Mamba [16] (SSM) on ImageNet. As shown in Tab. 6, MetaLA performs better than other typical linear models at both scales of 6M and 23M.Ablation Studies. We conduct ablation studies on the 360M model trained with 15B tokens and compare the results in zero-shot experiments. First, restoring the Key matrix in linear attention does not improve performance while increasing parameters, supporting our theoretical result that K is not necessary for approximation, and its functional role can be replaced by dynamic decay. Second, the ablations of self-augmentation and short convolution demonstrate the effectiveness of our model design, i.e., enhance tokens' own attention and local interactions. 
conclusion and discussion: 
 Conclusion. We unify LinFormer/SSM/LinRNN models into the form of linear attention with Query, Key, Vaule matrices, and then analyze whether they can achieve the optimal approximation to the softmax attention function. Theoretical analysis shows that the existing LinFormer/SSM/LinRNN cannot achieve optimal approximation. Consequently, we propose the MetaLA architecture, which can achieve functional approximation of softmax attention with least parameters. The performance on various types of tasks verifies the effectiveness of MetaLA.Discussion and Limitation. Here, we discuss two key questions about approximation perspectives. i) How does an optimal approximation to softmax attention inspire linear attention design? In this work, we mainly remove the Key matrix, use dynamic decay, and enhance local interactions and the token's own attention. This is clearly not the end of linear attention optimization. This work focuses on functional approximation, previous studies about the value approximation [30, 31, 50 ] can be further investigated on the basis of our functional approximation theory as well as MetaLA architecture. Additional optimization may include improving the recall ability of limited hidden states or designing better parameter functions. ii) Does approximation to the softmax attention imply that linear attention has an upper capacity limit? Taken literally, approximation seems to imply that linear attention cannot exceed softmax attention. However, we found better results for linear attention than softmax attention in some experimental results, such as zero-shot and LRA. Similar findings were also reported in previous work [13, 15, 16] . We argue that this issue deserves further exploration. For the time being, evaluation metrics that do not adequately reflect the model's capabilities, insufficient training [51] , and linear attention that does have advantages in certain abilities [15] are all possibilities. 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: This paper is foundational research and not tied to particular applications.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [No] Justification: Impact Statements. This paper presents work whose goal is to advance the field of Language Modeling. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: The data, code, and models used in this paper have all been properly cited.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [Yes]Justification: In Sec. 5, we provided details of the model we proposed.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] Justification: the paper does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA]Justification: the paper does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: We overview several architectures of three types of linear models, i.e., LinFormer, SSM, and linRNN, and show how they can be specialized from the general form in Sec. 3.LinFormer abandons the softmax form of vanilla attention, instead leveraging the dot product between Keys and Values to achieve linear complexity [9] . Advances in LinFormer mainly focus on designing better kernel function [30, 32, 52, 53, 54, 55, 50] such as the spiking neurons in a spiking neural network [56, 57, 58, 59 ] can be understood as efficient kernel functions [60, 61] , exploring architectural optimization [13, 14, 62] , introducing gating mechanism [15, 31, 63, 64] , etc. We show the general form of several LinFormer architectures in Tab. A1. For the general form of LinFormers,d k = d v = d is the most common choice.• Vanilla LinFormer [9] : Vanilla linear Transformers lack explicit decay. Instead, They choose to maintain two hidden states containing a denominator to normalize attention scores. As for their general form, q t /k t = f q/k (x t , θ q/k ) = ϕ(x t W Q/K ) where θ q/k = W Q/K and ϕ is the nonlinearity. v t = x t W V . There is no decay or output gate, i.e., g t = 1 and α t = 1.• RetNet [14] : Different from vanilla LinFormer, RetNet uses positional encoding exp (jnθ) and fixed decay λ to control the hidden state. For its general form,q t /k t /v t = f q/k/v (x t , θ q/k/v ) = x t W Q/K/V where θ q/k/v = W Q/K/V . α t = λ exp (jθ)is the fixed decay vector, where j is imaginary unit. Furthermore,y t = [o t ⊙ SiLU(x t W G )]W O where θ g = W G .• TransNormer [12] : TransNormer also uses positional encoding exp (jnθ) and fixed decay λ to control the hidden state. For its general form, q t /k t = f q/k (x t , θ q/k ) = ϕ(x t W Q/K ) where θ q/k = W Q/K and ϕ is the nonlinearity. v t = x t W V , and α t = λ exp (jθ) is the fixed decay vector, where j is imaginary unit. The normalization layer chosen by the paper is SRMSNorm. Furthermore,y t = [o t ⊙ (x t W U )]W O where θ g = W U .• GLA [15] : GLA considers a data-dependent gating mechanism for LinFormer. For its general form,q t /k t /v t = f q/k/v (x t , θ q/k/v ) = x t W Q/K/V where θ q/k/v = W Q/K/V . α t = sigmoid(x t W 1 α W 2 α + b α ) 1/τ is a dynamic decay, where W 1 α and W 2 α are low-rank matrices. Furthermore,y t = [o t ⊙ SiLU(x t W r + b r )]W O where θ g = [W r , b r ].Table A1 : From general form to LinFormers.Vanilla LinFormer [9] RetNet [14] TransNormer [12] GLA [15] f q (x t , θ q ) ϕ(x t W Q ) x t W Q ϕ(x t W Q ) x t W Q f k (x t , θ k ) ϕ(x t W K ) x t W K ϕ(x t W K ) x t W K f v (x t , θ v ) x t W V x t W V x t W V x t W V f α (x t , θ α ) 1 λ exp (jθ) λ exp (jθ) σ(x t W 1 α W 2 α + b α ) f g (x t , θ g ) 1 σ(x t W G ) x t W U σ(x t W r + b r ) Dimension d k = d v = d d k = d/2 d v = d A1.2 SSMSSM represents an alternative linear architecture to Transformer, based on the state space equations [24, 65] . Typical SSMs [25, 26, 66, 67, 68, 69] employ structured transition matrix and special initialization strategies such as HiPPO [27] to efficiently enhance long-range dependencies. Mamba [16] advances SSMs by introducing selection mechanism, i.e. input-dependent parameters [70] and designs a hardware-aware parallel algorithm. We show the general form of several SSM architectures in Table A2 . All of them are SISO (Single-Input, Single-Output) except S5 [22] . We omit original S4 [10] model because it focuses on Diagonal Plus Low-Rank (DPLR) transition matrix A rather than fully diagonal one, which is difficult to implement due to computational inefficiency. For most SSMs, d v = H, which means each channel uses independent parameters (q h t , α h t , k h t ). • DSS [26] : DSS is the first research to show the effectiveness of diagonal structured SSM with fixed parameters, which means all parameters are constant through time. For the general form of DSS,d v = d = H, d k H = N, where N is an expansion. Q, K, Λ are independent of x t and are determined by learnable parameters [A, B, C, ∆], using ZOH method to discretize. Furthermore, x t = v t , g t = 1 and W O = I, where I denotes the identity matrix.• S4D [25] : S4D is also a diagonal structured SSM with fixed parameters. It theoretically explain and expand the effectiveness of DSS and focus on parameterization and initialization of diagonal SSM. It can be discretized using ZOH or Bilinear method, resulting different f k/α compared with DSS. Similarly, x t = v t , g t = 1 and W O = I.• H3 [21] : H3 combines S4D with LinFormer and adds some architectural optimization such as local convolution and output gates. The core layer of H3 is S4D, and thus it shares the same general form with S4D.• S5 [22] : S5 uses MIMO (Multi-Input, Multi-Output) SSM with fixed parameters. For its general form,d v = N d = H, d k H = 1,where N is an expansion much smaller than that of SISO models. α t = exp (∆A) is the fixed decay and v t = f v (x t , θ v ) = x t B where B = (∆A) -1 (exp (∆A)-I)•∆B, i.e., it chooses ZOH method. Similar to previous SSM models, all the parameters are independent of x t and are determined by learnable parameters [A, B, C, ∆]. Furthermore, q t = k t = g t = 1 and W O = C.• Mamba [16] : Mamba advances SSMs with selection, which means all parameters are time-varying and input-dependent. For the general form of Mamba,d v = d = H, d k H = N . q t = C t = x t W C . Similarly, k t = B t = x t W B ∆ t .The decay term can be written asα t = exp(∆ t A). x t = v t , g t = 1 and W O = I. Furthermore, ∆ t = softplus(x t W lora + b ∆ ).Table A2 : From general form to SSMs.DSS [26] , S4D [25] , H3 [21] (ZOH) S4D [25] , H3 [21] (Bilinear) Mamba [16] S5 [22] f q (x t , θ q ) C C x t W C 1 f k (x t , θ k ) (∆A) -1 (exp (∆A) -I)∆B (I -∆ 2 A) -1 ∆B ∆ t (x t W B ) 1 f v (x t , θ v ) x t x t x t x t B f α (x t , θ α ) exp (∆A) (I -∆ 2 A) -1 (I + ∆ 2 A) exp (∆ t A) exp (∆A) f g (x t , θ g ) 1 1 1 1 Dimension d v = d = H d k /H = N d v = N d = H d k /H = 1 A1.3 LinRNNLinRNN is a variant of the vanilla Recurrent Neural Network (RNN) model [71, 72, 73] that eliminates the nonlinearity within the recurrence and employs element-wise product instead of matrix multiplication [11, 28, 74, 75] . Recent works have made efforts to explore more effective recurrence [17, 19, 76, 77] and design more advanced gating mechanisms [18, 20, 33] . We show the general form of several LinRNN architectures in Tab. A3. LinRNNs generally implement [19] , which means each channel uses independent parameters (q h t , α h t , k h t ). They maintain a smaller unexpanded hidden state, which is a 1D vector and can be seen as a special case with head dimension size d ′ k = 1. • RWKV-4 [17] : RWKV-4 is a linear RNN model which includes fixed decay and output gates. Moreover, RWKV-4 treats the present token differently and uses token shift operation (we omit for clear illustration). For its general form, k t = exp (x t W K ) where θ k = W K . v t = x t W V while q t = 1. α t = exp (-W ) is fixed decay with the learnable parameter θ α = W . Furthermore,d k = d v = H = d except LRUy t = [o t ⊙ sigmoid(x t W r )]W O where θ g = W r .• GLRU [20] : Gated linear recurrent unit (GLRU) mentioned in HGRN is a linear RNN model including input/forget/output gates. For the general form of GLRU,α t /k t = f α/k (x t , θ α/k ) = σ(x t W α/K + b α/K ) where θ α/k = [W α/K , b α/K ] and σ is the sigmoid function. v t = SiLU(x t W V + b V ) while q t = 1. Furthermore, y t = [o t ⊙ SiLU(x t W g + b g )]W O where θ g = [W g , b g ].• Griffin [18] : Griffin is an RNN model with gated linear recurrence. Moreover, it ties dynamic decay and Key as we do. For the general form of Griffin,α t = f α (x t , θ α ) = A c•σ(xtW a +ba) where θ α = [A, W a , b a ]. Then k t = 1 -α 2t to replace Key's role. v t = σ(x t W x + b x ) ⊙ x t which contains an input gate additionally. q t = 1 and yt = [o t ⊙ GeLU(x t W g + b g )]W O where θ g = [W g , b g ].• LRU [19] : LRU is a linear and diagonal RNN model with fixed decay, inspired by deep SSMs' success. For its general form,d v = N d = H, d k H = 1,where N is a small expansion. α t = A is the fixed decay andv t = f v (x t , θ v ) = x t B where θ v = B. Furthermore, q t = k t = g t = 1 and W O = C.Table A3 : From general form to LinRNN.RWKV-4 [17] GLRU [20] Griffin [18] LRU [19] f q (x t , θ q )1 1 1 1 f k (x t , θ k ) exp (x t W K ) σ(x t W i + b i ) 1 -α 2 t 1 f v (x t , θ v ) x t W V ϕ(x t W c + b c ) σ(x t W x + b x ) ⊙ x t x t B f α (x t , θ α ) exp (-W ) σ(x t W f + b f ) A c•σ(xtW a +ba) A f g (x t , θ g ) σ(x t W r ) ϕ(x t W g + b g ) σ(x t W g + b g ) 1 Dimension d k = d v = d = H d v = N d = H d k /H = 1Approximation to softmax attention includes two parts: value and functional approximation. Some previous works of LinFormer [31, 30, 53, 78] were devoted to designing better kernel map ϕ to approximate the value of SoftAttMap via randomized features, i.e., ϕ(q t )ϕ ⊤ (k s ) ≈ exp(q t k ⊤ s ). Recent work [50] further introduces attention weight distillation loss to minimize the cross-entropy loss between the computed linear attention weights and those that would have been computed via softmax attention, in order to achieve strict value approximation. Most of them follow typical designs in [9] and use only Q and K. However, in this work, we find that only considering value approximation is insufficient. Different from previous studies, we investigate functional approximation of SoftAttMap. Through theoretical analysis, we show that these works cannot achieve the functionality of softmax attention due to lack of dynamic decay Λ t . Therefore, their purpose of achieving value approximation cannot be achieved as well. It is worth noting that this is not a denial of the significance of previous works about value approximation. We can also get inspiration from them. Functional approximation is the basis and prerequisite for value approximation.In this section, we show that existing linear models are not optimal linear approximation to softmax attention map, according to definition 4.1, and prove that linear models with only Query and dynamic decay can satisfy the optimal linear approximation conditions. In definition 4.1, C0 underlines computational and memory efficiency. C1 and C2 consider memory and modeling ability of linear attention. C3 is our expectation to seek the least parameters on the premise that previous three conditions are met. We first discuss the requirements to satisfy each conditions, based on the general form of linear attention in Sec. 3. All the linear models own linear complexity and satisfy C0, so we analyze C1, C2 and C3 in appendix A2.1,appendix A2.2 and appendix A2.3, respectively. Then we synthesize all the discussions and present the conclusions in appendix A2.4.For notation we use Q, K, V, Λ t to denote Query, Key, Value and Decay matrices respectively. Considering decay can be either dynamic or fixed, here we use subscript t to distinguish, i.e., Λ/Λ t denote fixed/dynamic decay. Corresponding q t , k t , v t , α/α t are query, key, value and fixed/dynamic decay of each timestep.Proposition A2.1. Only models with dynamic decay can satisfy C1 (Dynamic memory ability). Let S t ∈ R d k ×dv be hidden state of general linear attention (see Sec. 3). At a time t, the information about v t1 , . . . , v t d k is successfully stored in S t (t 1 , . . . , t d k ≤ t), which means (S t ) i,: = v ti . When the new important input v t+1 arrives, only models with dynamic decay can substitute historical unimportant v t1 by v t+1 successfully, i.e., obtain (S t+1 ) 1,: = v t+1 .Proof. Without loss of generality, suppose t = d k and t i = i(i = 1, . . . , d k ), where d k is Key dimension. That is, the information about v 1 , . . . , v d k is successfully stored, which means the i-th row of S d k , (S d k ) i,: = v i . Suppose v 1 is unimportant information need to be substituted. Now the new and important input v d k +1 arrives, models with general form Eq. ( 13) update S d k as follows (heads are omitted for simplicity):S d k +1 = diag(α d k +1 )S d k + (k d k +1 ) T v d k +1 .(A1)We will check whether the model can obtain(S d k +1 ) 1,: = v d k +1, with fixed/dynamic decay or without decay. In the following discussion we letk d k +1 = 1.i) Models with fixed decay, which means α d k +1 = α ∈ (0, 1). The model can update and obtain(S d k +1 ) 1,: = αv 1 + v d k +1. This means the model with fixed decay can only store most recent several tokens rather than store most important tokens without information loss, i.e., it has ability to forget but no ability to forget and memorize dynamically.ii) Models without decay, which meansα d k +1 = 1. This case results in (S d k +1 ) 1,: = v 1 + v d k +1, which means they have no mechanism to erase old information. So such linear models have no ability to forget and memorize dynamically, and what they can only do is information blending. Thus the attention at time t is distracted by all the information before, making it hard to focus on and approximate most important ones. After normalization, the attention distribution will have a relative high entropy along time dimension, which is often referred as attention dilution problem.iii) Models with dynamic decay, which means α d k +1 ∈ [0, 1]. The model can easily erase and substitute information. Settingα d k +1 = [0, 1, • • • , 1] can obtain (S d k +1 ) 1,: = v d k +1. So such linear models have ability to memorize and forget dynamically, making historical unimportant information have few effect on new one.Eq. ( 18) illustrates the LinAttMap only relate to q t = f q (x t , θ q ), k t = f k (x t , θ k ) and α t = f α (x t , θ α ). Assuming the inputs are good enough and the functions (f q , f k , f α ) are expressive enough, we can shift from solving (θ q , θ k , θ α ) to solving (Q, K, Λ t ). Based on definition 4.1, we focus on approximating the attention scores between stored tokens x t1 , . . . , x t d k .Proposition A2.2. Only models with parameters (Q, K, Λ t ), (Q, K) or (Q, Λ t ) can satisfy C2 (Static approximation ability). For an arbitrarily given softmax attention map P with scores p ts , only above models can ensure Eq. (A2) and Eq. (A3) hold with bounded parameters.EQUATIONEQUATIONwhere C q , C k are constant and || • || denotes vector norm. In our general form in Sec. 3, f (x t , x s |Q, K, Λ t ) = LinAttMap (Q, K) t,s (see Eq. ( 18)). For bounded inputs X, bounded parameters (θ q , θ k , θ α ) are equivalent to bounded (Q, K, Λ t ).normalization. In Retrieval, Image, Pathfinder and Path-X we use a bid-model. The hyperparameters for all tasks can be found in Table A4 Image Classification. We evaluate our models on ImageNet [40] . The input size of ImageNet is 224 × 224. Following Deit [49] , the batch size is set to 2048 during 300 training epochs with a cosinedecay learning rate whose peak value is 2.4 × 10 -3 . We choose AdamW (β 1 = 0.9, β 2 = 0.98) with 0.05 weight decay as the optimizer. Note that we do not use cutmix or mixup during the training. More challenging settings for MQAR. We evaluate some models with sequence length 512 and with more retrieval key-value pairs (80, default is 64). The attention baseline benefits from global modeling capabilities, achieving optimal results. The results in Tab. A7 show that Met-aLA outperforms Mamba (does not converge under the same training conditions), and there is still a significant gap compared to transformers.Results on the synthetic MAD task. we evaluate several models on MAD [41] , a collection of six synthetic tasks predictive of scaling laws, Figure A2 : Training efficiency evaluations. The throughput and memory usage on a single A800 GPU of Transformer and various linear models. Transformer++ is implemented using FlashAttention [90] and SwiGLU.including recall, memorization, and compression. As shown in Tab. A6, MetaLA achieves the best results across various linear complexity models.Results on the Needle in a Haystack (NIAH) task. We also present experimental results on the NIAH task following [42] , which is designed to evaluate the in-context retrieval capabilities of LLMs. Retrieval ability in long texts is a significant challenge for linear models, as all current linear models lack good solutions to this problem. Nonetheless, Tab. A8 shows that MetaLA has achieved satisfactory results in comparisons among linear models. Compared to Transformer models, this performance is still insufficient. This is precisely the issue we hope to address next, following the unification of linear model forms.The scalability with respect to model size and training tokens. For preliminary validation, we further evaluate our model ranging in size from 380M to 3B, trained with up to 300B tokens, on the CommonSense Reasoning benchmark. The strong results in Tab. A9 demonstrate the potential of our model when scaling up the parameter scale and training dataset.Training efficiency evaluations. The comparative results on training throughput and GPU memory usage across various 1.3B-sized models are shown in Fig. A2 . The report indicates that: (1) Our model demonstrates good linearity, maintaining processing speed and memory efficiency with increasing sequence length, unlike the Transformer, which experiences a sharp drop in token processing speed as sequence length increases.(2) Our model matches the computational efficiency of linear models like GLA [2] in both latency and memory, and is significantly faster than Mamba [5] , which also has linear complexity. Justification: In the abstract, we summarized the theory we proposed, our solution MetaLA, and the results obtained from our experiments.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 


Decision: This submission receives scores of 7, 7, 8, 6, indicating an acceptance of this submission. The reviewers generally agree that the paper presents a valuable contribution to the field of efficient attention mechanisms. The unified framework, theoretical analysis, and the proposed MetaLA model are novel and promising. The additional experiments and clarifications provided in the author's rebuttal further strengthen the paper's contributions.

It is important to address the reviewers' concerns regarding the clarity of the presentation and the scope of the experiments in the final version. A more thorough discussion of the limitations and trade-offs between MetaLA and softmax attention, particularly in recall-intensive tasks, would also enhance the paper.

Strengths:
- Unified Framework: The paper provides a valuable unified framework for understanding different linear attention models, including LinFormer, SSM, and LinRNN.
- Theoretical Analysis: The theoretical analysis of the optimal linear approximation to softmax attention is well-received.
- Novel Model: The proposed MetaLA model is seen as a novel and promising approach to improving the efficiency of attention mechanisms.
- Empirical Results: The empirical evaluation on various tasks demonstrates the effectiveness of MetaLA compared to existing linear models.

Area to Improve:
- Presentation: The presentation of the theoretical analysis and the "optimal linear approximation" concept could be improved for better clarity.
- Experimental Scope: Some reviewers suggest expanding the experiments to include more challenging tasks, longer sequences, and a wider range of model scales.
- Justification: Further justification is needed for certain design choices, such as the removal of the Key matrix and the use of self-augmentation.
- Comparison with Softmax Attention: A more in-depth discussion of the limitations and trade-offs between MetaLA and softmax attention, particularly in recall-intensive tasks, is recommended.

Reviews: review: paper_summary: paper_summary: This paper discusses the development and evaluation of linear complexity models for Transformers, which aim to replace the conventional softmax attention mechanism. The authors first unify existing models including LinFormer, SSM, and LinRNN into the framework of Linear Attention. Then they establish three (actually four) criteria for optimal linear attention: (0) linear complexity, (1) dynamic memory capability, (2) static approximation ability, and (3) minimal parameter usage. This paper shows that the three existing models mentioned above do not meet all these conditions and therefore perform suboptimally. In response, they introduce a new model, Meta Linear Attention (MetaLA), designed to fulfill these criteria. MetaLA is tested across various tasks including the Multi-Query Associative Recall (MQAR) task, language modeling, image classification, and the Long-Range Arena (LRA) benchmark, where it outperforms existing linear models, demonstrating its effectiveness.
strength_weakness: strength_weakness: (After rebuttal: The authors carefully addressed my main concerns.)
--------------

Please correct me if there are any mistakes.

1, The main concern originates from the three criteria of "**optimality**" for linear attention (LA). In general, "optimality" indicates a situation that any possible method/result cannot be substantially better. However, it seems like these three criteria are all *necessary* conditions. As far as I noticed, this work has not stated whether these criteria are *sufficient* for optimality. As a consequence, even if MetaLA satisfies all three criteria, it is not proper to state that MetaLA is *optimal*.

2, The presentations could be improved, especially in Section 2. It is important to define or denote the dimensionality of not only each matrix/vector (which the authors have done) but also each function and operation (which they have not).
questions: questions: The authors are encouraged to disagree or discuss on my questions/comments.

1, On Page 2 Line 67, what is the mask matrix $M$ like? From the context, I guess $M$ is a lower triangular matrix with all valid elements $=1$?

2, On Page 3 Line 78 (top), why does the LinFormer consume $O(1)$ time and memory complexity? Also, would you please elaborate on the "Chunkwise algorithm" mensioned in Line 80?

3, On Page 5 Line 147, why this softmax-generated expressive attention map requires infinite states?

4, On Page 5 Line 177, what is the condition (or consequence) of "good enough" inputs?

5, On Appendix A 2.3 (and also A 2.4), it is claimed that (Q, $\Lambda_t$) can meet C0, C1, C2 with fewest parameters. Is this concept of "fewest" restricted in a certain domain of model selection? In other words, do you have a lower-bound proof showing that any model utilizing fewer parameters than (Q, $\Lambda_t$) will fail?
limitations: limitations: Limitations are adequately discussed in the Section of Conclusion and Discussion.
paper_summary: paper_summary: They proposed MetaLA, which solved problems of previous attention alternatives (LinRNN, SSM, LinFormer). They start to build their MetaLA by deriving general form of linear alternative of softmax attention.

1. They remove the K matrice redundant parameters and achieve better training efficiency
2. Add self-augmentation. This allows the input value to immediately affect the output depending on the query to prevent forgetting current token information.
3. Adding a convolutional layer at input X. This enhances the local interaction further, motivating Mamba and Griffin.
strength_weakness: strength_weakness: 1. The experiment is not done with various scales. It is hard to know the effect of the scaling model and training dataset with MetaLA.
2. The performance of MetaLA is mainly shown with only benchmark scores, and there are no latency reports. Therefore, it is hard to know if this performance is Pareto optimal in the trade-off between latency and performance.
questions: questions: 1. Is approximating softmax attention the optimal solution for sequential modeling? This paper aims to approximate the various characteristics of the computational and model aspects very well. However, I wonder if softmax attention is not the optimal solution for sequential modeling. Is there any justification for it?
2. Can you provide the code?
limitations: limitations: It feels minor improvement on top of previous linear attention alternative studies.
paper_summary: paper_summary: The paper presents a theoretical analysis of existing linear attention methods such as LinFormer, SSM, and LinRNN. Building on this analysis, the authors propose a unified framework that combines the strengths of these methods. Utilizing this framework, authors develop a novel linear attention model called MetaLA. The key innovations of MetaLA include the elimination of the Key matrix, the introduction of dynamic decay for achieving dynamic memory and static approximation, and the integration of self-augmentation with short convolution. 
These enhancements aim to improve the approximation accuracy of softmax attention and outperform existing linear attention methods.
strength_weakness: strength_weakness: - The only result for long sequences provided by the authors is on the LRA benchmark. This is not sufficient to justify the scalability of the method for real LLMs trained on sequences longer than 1k tokens since LRA is an artificial and non-representative benchmark.
- Authors provide empirical results mainly for classification tasks. It is not clear how their model performs in the text generation scenario.
- The authors neither mention nor present any experiments with low-precision formats like FP16 and BF16. Many linear attention methods fail when trained in these popular formats.
- The authors do not provide any measurements or comparisons of computational efficiency and memory requirements of their method compared to softmax attention and other linear methods.
- Since the authors do not provide code for the MetaLA model, it is hard to assess the practicality and ease of implementation.
questions: questions: - How does MetaLA perform on tasks with longer sequences compared to standard softmax attention and other linear attention mechanisms? For example, on NarrativeQA or other tasks from LongBench.
- Can the authors provide more information on the setting of hyperparameters for their method, specifically whether the decay factor should be selected by the user?
- Another linear attention work [1], which the authors did not cite, proposes evaluating the concentration ability of linear attention using entropy and spectral gap metrics. What is the authors' opinion on these concentration metrics in the context of the MetaLA framework analysis?


[1] Linear Log-Normal Attention with Unbiased Concentration
limitations: limitations: N/A
paper_summary: paper_summary: Proposes a unifying view of recent linear attention/SSM/linear RNN methods and compares these models to Softmax Attention. Proposes MetaLA to address the shortcomings of the prior methods in approximating Softmax attention. Performs experiments on MQAR, language modeling, image classification and LRA.
strength_weakness: strength_weakness: - The general presentation of the "optimal linear approximation" and theoretical analysis is dense and confusing and obscures the contributions. 
  - I would recommend moving the Proposition statements to the main text (the proofs can remain in the appendix). E.g. Currently I get to line 218 and am told that the key value is being dropped based on theoretical analysis performed in Appendix A2, but I would rather already have an idea of this result before getting to this point. Similarly, the "self-augmentation" change seems to come out of nowhere.
  - Once you have moved these Proposition statements to the Section 4, I would recommend rewriting/reorganizing Section 4 framed around these proposition statements
  
- The MQAR analysis is weak. I would like to see more difficult versions of this task (larger vocab sizes, more retrievals, longer contexts) and an attention baseline to better understand the limitations and differences between the different methods. 
- I think more downstream tasks should be included that require the recall/retrieval abilities that are known to affect all these fixed state size methods (e.g. https://arxiv.org/abs/2402.01032, https://arxiv.org/abs/2402.04248, https://arxiv.org/abs/2402.18510v3 to name a few, more generally these and other works should be cited) would make the proposed approach and value of the proposed framework more convincing. This would help the reader understand how much progress as actually been made toward approximating Softmax attention with more efficient methods. Even simple needle in the haystack, passkey retrieval, or phonebook retrieval (as in the Griffin paper) would help give a better sense of this. 
- In the conclusion and discussion, some potential weaknesses and questions regarding the differences between Softmax attention and other methods are mentioned and proposed as future work (model capabilities, insufficient training, or eval issues). However, I would suggest some of these questions should be better explored in this work since that would help provide a better sense of how useful the proposed framework and method is.
questions: questions: My main questions and issues are listed above.

- Line 287: typo, "Value" is misspelled in the text
limitations: limitations: Sufficient

Comments: title: Thanks for the replies
comment: Thank you for kind and insightful response. I do not have any more concerns and questions anymore. It was really worth it to discuss with authors and feel honor to reviewing this paper. I also wonder what will be happened next to linear vs. quadratic attention wars. I hope your future research will reveal the parato front of these trade off soon! Again, thank you.
Thank you for your feedback, suggestions, and insightful discussion, which have had a positive impact on us. We continue to address your concerns below:

This is an interesting question. The current linear models can match or even outperform softmax attention in most tasks while being more efficient. However, they fall short in terms of memory capacity compared to softmax attention. As we can see, there is a conflict between the increasing input sequences and the fixed memory space. Therefore, both approaches have their advantages and disadvantages, making it difficult to achieve Pareto optimality. To develop a model that balances efficiency with good memory capability, we think that relaxing the complexity of processing each token (to log-linear) could be a promising solution. This will be left for our future work.
Thank you for your thoughtful feedback and for increasing your confidence in our paper. We genuinely appreciate your time and effort in reviewing our paper and responses. Your support and understanding are invaluable to us. If you have any further questions or need additional clarification, please let us know. Thank you once again for your help.
The reviewer thanks the authors for your thorough and informative explanations. I am not an expert in this area, and therefore I am not able to advocate for this paper. However, I will increase my confidence from 2 to 3 as a reflection of what I have learned from your response.
Thank you sincerely for your thoughtful feedback and for increasing your score. We are grateful that our responses and additional experiments have addressed your concerns. Your support and encouragement mean a lot to us, and we truly appreciate the time you’ve taken to review our work.
Thank you very much for your thoughtful feedback and for raising your score. We sincerely appreciate your valuable insights, and we will make sure to thoroughly discuss the gap with global attention in recall-intensive tasks as you suggested. Additionally, we will include a softmax attention baseline for the MAD tasks in the final version to enhance the comparison. Your comments have been instrumental in improving our work, and we are grateful for your time and consideration.
Thank for for your thoughtfully written rebuttal. I gave the borderline accept because I could not fully understand this paper, therefore my reviewing confidence was pretty low. 

My concerns was all resolved (performance scalability and latency optimality) by general response so I want to raise my score to accept (however I am not still sure about the content of paper. E.g. equations and background theories). However, after reading the general response, now I can understand how this paper is novel and valuable in more depth. 

The response of my question about optimality of softmax attention is interesting. In my opinion, the softmax attention solves the problem of RNN and human intelligence which is limitation of memory and computation for each token. I think we should improve the softmax attention while maintaining it's non constant time and space complexity to achieve AGI that performs better than human. In this sense, I have a question about linear attention mechanisms. What if we increase memory space of linear attention mechanism as non constant (linear or log linear)? As far as I know, all linear attentions are strictly limit their space and time complexity as constant for each tokens. But I wonder that is it really pareto optimal if we scale it compare to quadratic attention. Do you think current implementation of linear attention can handle this? (Including this paper)
Thank you for the rebuttal. The authors have addressed my concerns regarding long sequences and numerical formats. Additionally, they provided further experimental results, comparing their method to the LLaMa2 and MAMBA models. Overall I think it's a good paper and following the additional information provided by the authors, I have decided to increase my score.
- Thank you for your response. I appreciate the additional recall intensive experiments. Please ensure the gap with global attention that these experiments expose for recall intensive tasks are discussed thoroughly. This discussion of limitations and the additional experiments will strengthen the paper.

- Please also include a softmax attention baseline for the MAD tasks in a final version to provide better comparison

I have raised my score.