Title: MetaLA: Unified Optimal Linear Approximation to Softmax Attention Map
Venue: NeurIPS 2024
Accepted: yes
Keywords: Self-attention

Linear attention

Efficient

Softmax attention

Foundation Model

Language Model
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.

MetaLA: Unified Optimal Linear Approximation to Softmax Attention Map

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: i) Dynamic memory ability; ii) Static approximation ability; iii) 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. Code: https://github.com/BICLab/MetaLA

1 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?

2 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:

where W Q , W K ∈ R d × d k , W V ∈ R d × d v 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:

At 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]:

where q t , k t ∈ R 1 × d k and v t ∈ R 1 × d v 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 training

which 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:

in 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 kernel K

where ∗ 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] is

where 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 d k × d v , SSM with h t ∈ R N × d , linear RNNs with h t ∈ R 1 × d , where d k > N > 1 in usual. Inspired by this fact, we unify LinFormer/SSM/LinRNN mechanisms in the form of linear attention, formally containing Query, Key, and Value matrices (see Fig. 1).

⊙

element-wise multiplication

⊕

element-wise addition

⊘

element-wise division

⊗

outer-product

⊛

matrix multiplication

𝐐

𝐀

𝐕

𝐊

𝐆

⊙

⊙

∗

𝐗

𝐎

XNorm

𝐘

f 𝑞

f𝛼

f 𝑘

f 𝑣

f𝑔

𝑾𝑂

⊘

∗

⊙

𝐌

Parallel Form

𝐒𝑡

𝐒𝑡-1

𝛂𝑡

𝐤𝑡

𝐪𝑡

𝐯𝑡

𝐠𝑡

⊙

⊙

∗

⊗

𝐱𝑡

⊕

𝐨𝑡

XNorm

𝐲𝑡

f𝛼

f 𝑘

f

𝑣

f 𝑞

f𝑔

𝑾𝑂

Recurrent  Form

General Recurrent Form of LinFormer/SSM/LinRNN is:

̸

where 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 × d v and θ v/g are trainable parameters. In Eq. (13), q t , α t , k t , v t are divided into H partitions (heads), where d ′ k/v = d k/v H , 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 to d from d v through W O ∈ R d v × 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:

K 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):

As 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

relatively large matrix hidden state whose size ( d v d k /H ) is mainly correlated to model dimension d . ii) SSM originate from state space equations. So they focus on how to better maintain the hidden state and optimize f α ; They have a matrix hidden state of moderate size ( d v N ), which is correlated to the fixed expansion N . iii) LinRNN originate from removing nonlinearity in the recurrence of vanilla RNN. So they focus on designing better forget/input/output gates, i.e., to optimize f α , f k , f g ; They have 1D vector hidden state whose size ( d v ) 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 .

4 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.

Definition 4.1. Necessary Conditions of Optimal Linear Approximation to Softmax Attention Map. A function f ( x t , x s | θ ) : R 1 × d ×R 1 × d →R , used to compute attention score between any x t and x s (tokens), with parameters θ , is an optimal linear approximation to softmax attention map if it satisfies: i) Linear complexity . Attention map can be computed in linear time, i.e., O ( n ) space and time complexity during training and O (1) space and time complexity during inference. ii) Dynamic memory ability . When handling inputs sequentially, f ( x t , x s | θ ) with limited hidden states should be able to store the most important information adaptively while forgetting unimportant ones. iii) Static approximation ability : For an arbitrarily given softmax attention map P with scores p ts , there must exists bounded θ such that f ( x t , x s | θ ) = p ts , ∀ t, s = 1 , · · · , n . iv) Least parameter approximation : On the premise that the first three conditions are met, use as few parameters as possible to achieve approximation to softmax attention map.

In definition 4.1, Condition 0 (C0) underlines computational and memory efficiency. Conditions 1 (C1) and 2 (C2) consider memory and modeling ability of linear attention. Due to limited state size d , linear attention can only memorize the history of most important d tokens without information loss and precisely model arbitrary attention map of those d tokens. Condition 3 (C3) is our expectation to seek the least parameters on the premise that previous three conditions are met.

Theoretical Analysis for Optimal Linear Approximation. For the C1 condition, suppose the information about v t 1 , . . . , v t d k is successfully stored in S t ( t 1 , . . . , t d k ≤ t ), we will check whether the model can substitute unimportant v t 1 when the new important input v t +1 arrives.

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

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:

For 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]

lack 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.

5 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.

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:

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

y

M

0

1 - αt, Vt, qt

w

aug

XNorm

St-1

不

St

α

2b

gt

g

Conv1d

|Xt

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]:

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].

6 Experiments

We conduct a comprehensive evaluation of MetaLA to validate its

Seq_len:256

Seq_len:512

1.00

0.75

Based

GLA

0.50

RWKV

0.25

Mamba

0.00

MetaLA

64 128 256 51264 128 256 512

Model dimensionModel dimension

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.

7 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.

References

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems , pages 5998-6008, 2017.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in Neural Information Processing Systems , 33:1877-1901, 2020.

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 , 2023.

Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 , 2023.

Rohan Anil, Andrew M Dai, Orhan Firat, Melvin Johnson, Dmitry Lepikhin, Alexandre Passos, Siamak Shakeri, Emanuel Taropa, Paige Bailey, Zhifeng Chen, et al. Palm 2 technical report. arXiv preprint arXiv:2305.10403 , 2023.

Chengyi Wang, Sanyuan Chen, Yu Wu, Ziqiang Zhang, Long Zhou, Shujie Liu, Zhuo Chen, Yanqing Liu, Huaming Wang, Jinyu Li, et al. Neural codec language models are zero-shot text to speech synthesizers. arXiv preprint arXiv:2301.02111 , 2023.

Haotian Liu, Chunyuan Li, Qingyang Wu, and Yong Jae Lee. Visual instruction tuning. Advances in Neural Information Processing Systems , 36, 2024.

Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Computing Surveys , 55(6):1-28, 2022.

Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning , pages 5156-5165. PMLR, 2020.

Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In The Tenth International Conference on Learning Representations , 2022.

Eric Martin and Chris Cundy. Parallelizing linear recurrent neural nets over sequence length. In International Conference on Learning Representations , 2018.

Zhen Qin, Xiaodong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong, Nick Barnes, and Yiran Zhong. The devil in linear transformer. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing , pages 7025-7041, 2022.

Zhen Qin, Dong Li, Weigao Sun, Weixuan Sun, Xuyang Shen, Xiaodong Han, Yunshen Wei, Baohong Lv, Xiao Luo, Yu Qiao, and Yiran Zhong. Transnormerllm: A faster and better large language model with improved transnormer, 2024.

Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models. arXiv preprint arXiv:2307.08621 , 2023.

Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficient training. In Forty-first International Conference on Machine Learning , 2024.

Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. In First Conference on Language Modeling , 2024.

Bo Peng, Eric Alcaide, Quentin Gregory Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, et al. Rwkv: Reinventing rnns for the transformer era. In The 2023 Conference on Empirical Methods in Natural Language Processing , 2023.

Soham De, Samuel L Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Muraru, Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, et al. Griffin: Mixing gated linear recurrences with local attention for efficient language models. arXiv preprint arXiv:2402.19427 , 2024.

A. Orvieto, S. L. Smith, A. Gu, A. Fernando, C. Gulcehre, R. Pascanu, and S. De. Resurrecting recurrent neural networks for long sequences. In International Conference on Machine Learning , pages 26670-26698, 2023.

Zhen Qin, Songlin Yang, and Yiran Zhong. Hierarchically gated recurrent neural network for sequence modeling. In Advances in Neural Information Processing Systems , volume 36, pages 33202-33221, 2023.

Daniel Y. Fu, Tri Dao, Khaled Kamal Saab, Armin W. Thomas, Atri Rudra, and Christopher Re. Hungry hungry hippos: Towards language modeling with state space models. In The Eleventh International Conference on Learning Representations , 2022.

Jimmy TH Smith, Andrew Warrington, and Scott Linderman. Simplified state space layers for sequence modeling. In The Eleventh International Conference on Learning Representations , 2022.

Zhen Qin, Weigao Sun, Dong Li, Xuyang Shen, Weixuan Sun, and Yiran Zhong. Various lengths, constant speed: Efficient language modeling with lightning attention. In Forty-first International Conference on Machine Learning , 2024.

Albert Gu. Modeling Sequences with Structured State Spaces . Stanford University, 2023.

Albert Gu, Karan Goel, Ankit Gupta, and Christopher Ré. On the parameterization and initialization of diagonal state space models. Advances in Neural Information Processing Systems , 35:35971-35983, 2022.

Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are as effective as structured state spaces. Advances in Neural Information Processing Systems , 35:22982-22994, 2022.

Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. In Advances in Neural Information Processing Systems , pages 1474-1487, 2020.

Tao Lei, Yu Zhang, Sida I. Wang, Hui Dai, and Yoav Artzi. Simple recurrent units for highly parallelizable recurrence. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing , pages 4470-4481, 2018.

Zhen Qin, Songlin Yang, Weixuan Sun, Xuyang Shen, Dong Li, Weigao Sun, and Yiran Zhong. HGRN2: Gated linear RNNs with state expansion. In First Conference on Language Modeling , 2024.

Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, et al. Rethinking attention with performers. In International Conference on Learning Representations , 2020.

Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah Smith, and Lingpeng Kong. Random feature attention. In International Conference on Learning Representations , 2020.

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.

David Balduzzi and Muhammad Ghifary. Strongly-typed recurrent neural networks. In International Conference on Machine Learning , pages 1292-1300. PMLR, 2016.

Ilya O Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, and Alexey Dosovitskiy. Mlp-mixer: An all-mlp architecture for vision. In Advances in Neural Information Processing Systems , volume 34, pages 24261-24272, 2021.

Weihao Yu, Chenyang Si, Pan Zhou, Mi Luo, Yichen Zhou, Jiashi Feng, Shuicheng Yan, and Xinchao Wang. Metaformer baselines for vision. IEEE Transactions on Pattern Analysis and Machine Intelligence , 46(2):896-912, 2024.

Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Ré. Zoology: Measuring and improving recall in efficient language models. arXiv preprint arXiv:2312.04927 , 2023.

Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027 , 2020.

Alex Wang, Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. Superglue: A stickier benchmark for general-purpose language understanding systems. Advances in Neural Information Processing Systems , 32, 2019.

Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers. In International Conference on Learning Representations , 2020.

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , pages 248-255, 2009.

Michael Poli, Armin W Thomas, Eric Nguyen, Pragaash Ponnusamy, Björn Deiseroth, Kristian Kersting, Taiji Suzuki, Brian Hie, Stefano Ermon, Christopher Re, et al. Mechanistic design and scaling of hybrid architectures. In Forty-first International Conference on Machine Learning .

Xuyang Shen, Dong Li, Ruitao Leng, Zhen Qin, Weigao Sun, and Yiran Zhong. Scaling laws for linear complexity language models. arXiv preprint arXiv:2406.16690 , 2024.

Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, and Christopher Ré. Simple linear attention language models balance the recall-throughput tradeoff. arXiv preprint arXiv:2402.18668 , 2024.

Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, et al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning , pages 2397-2430. PMLR, 2023.

Alex Andonian, Quentin Anthony, Stella Biderman, Sid Black, Preetham Gali, Leo Gao, Eric Hallahan, Josh Levy-Kramer, Connor Leahy, Lucas Nestler, Kip Parker, Michael Pieler, Jason Phang, Shivanshu Purohit, Hailey Schoelkopf, Dashiell Stander, Tri Songz, Curt Tigges, Benjamin Thérien, Phil Wang, and Samuel Weinbach. GPT-NeoX: Large Scale Autoregressive Language Modeling in PyTorch. https://www.github.com/eleutherai/gpt-neox , 2023.

Zhen Qin, Xiaodong Han, Weixuan Sun, Bowen He, Dong Li, Dongxu Li, Yuchao Dai, Lingpeng Kong, and Yiran Zhong. Toeplitz neural network for sequence modeling. In The Eleventh International Conference on Learning Representations , 2022.

Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Luke Zettlemoyer. Mega: Moving average equipped gated attention. In The Eleventh International Conference on Learning Representations , 2022.

Yuhong Li, Tianle Cai, Yi Zhang, Deming Chen, and Debadeepta Dey. What makes convolutional models great on long sequence modeling? In The Eleventh International Conference on Learning Representations , 2022.

Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Herve Jegou. Training data-efficient image transformers & distillation through attention. In International Conference on Machine Learning , volume 139, pages 10347-10357, 2021.

Michael Zhang, Kush Bhatia, Hermann Kumbong, and Christopher Re. The hedgehog & the porcupine: Expressive linear attentions with softmax mimicry. In The Twelfth International Conference on Learning Representations , 2023.

Ido Amos, Jonathan Berant, and Ankit Gupta. Never train from scratch: Fair comparison of long-sequence models requires data-driven priors. In The Twelfth International Conference on Learning Representations , 2024.

Jungo Kasai, Hao Peng, Yizhe Zhang, Dani Yogatama, Gabriel Ilharco, Nikolaos Pappas, Yi Mao, Weizhu Chen, and Noah A Smith. Finetuning pretrained transformers into rnns. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing , pages 10630-10643, 2021.

Krzysztof Marcin Choromanski, Han Lin, Haoxian Chen, Arijit Sehanobish, Yuanzhe Ma, Deepali Jain, Jake Varley, Andy Zeng, Michael S Ryoo, Valerii Likhosherstov, et al. Hybrid random features. In International Conference on Learning Representations , 2021.

Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströmformer: A nyström-based algorithm for approximating self-attention. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 35, pages 14138-14148, 2021.

Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In International Conference on Machine Learning , pages 9355-9366. PMLR, 2021.

Wolfgang Maass. Networks of spiking neurons: the third generation of neural network models. Neural Networks , 10(9):1659-1671, 1997.

Man Yao, Guangshe Zhao, Hengyu Zhang, Yifan Hu, Lei Deng, Yonghong Tian, Bo Xu, and Guoqi Li. Attention spiking neural networks. IEEE Transactions on Pattern Analysis and Machine Intelligence , 45(8):9393-9410, 2023.

Man Yao, Ole Richter, Guangshe Zhao, Ning Qiao, Yannan Xing, Dingheng Wang, Tianxiang Hu, Wei Fang, Tugba Demirci, Michele De Marchi, et al. Spike-based dynamic computing with asynchronous sensing-computing neuromorphic chip. Nature Communications , 15(1):4464, 2024.

Guoqi Li, Lei Deng, Huajin Tang, Gang Pan, Yonghong Tian, Kaushik Roy, and Wolfgang Maass. Brain-inspired computing: A systematic survey and future trends. Proceedings of the IEEE , 112(6):544-584, 2024.

Man Yao, Jiakui Hu, Zhaokun Zhou, Li Yuan, Yonghong Tian, Bo Xu, and Guoqi Li. Spikedriven transformer. Advances in Neural Information Processing Systems , 36, 2024.

Man Yao, JiaKui Hu, Tianxiang Hu, Yifan Xu, Zhaokun Zhou, Yonghong Tian, Bo XU, and Guoqi Li. Spike-driven transformer v2: Meta spiking neural network architecture inspiring the design of next-generation neuromorphic chips. In The Twelfth International Conference on Learning Representations , 2024.

Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc Le. Transformer quality in linear time. In International Conference on Machine Learning , pages 9099-9117. PMLR, 2022.

Huanru Henry Mao. Fine-tuning pre-trained transformers into decaying fast weights. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing , pages 10236-10242, 2022.

Tobias Katsch. Gateloop: Fully data-controlled linear recurrence for sequence modeling. arXiv preprint arXiv:2311.01927 , 2023.

Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-time models with linear state space layers. Advances in Neural Information Processing Systems , 34:572-585, 2021.

Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces. In International Conference on Learning Representations , 2023.

Junxiong Wang, Jing Yan, Albert Gu, and Alexander M Rush. Pretraining without attention. In Findings of the Association for Computational Linguistics: EMNLP , pages 58-69, 2023.

Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. In International Conference on Machine Learning , pages 28043-28078. PMLR, 2023.

Eric Nguyen, Karan Goel, Albert Gu, Gordon Downs, Preey Shah, Tri Dao, Stephen Baccus, and Christopher Ré. S4nd: Modeling images and videos as multidimensional signals with state spaces. Advances in Neural Information Processing Systems , 35:2846-2861, 2022.

Ramin Hasani, Mathias Lechner, Tsun-Hsuan Wang, Makram Chahine, Alexander Amini, and Daniela Rus. Liquid structural state-space models. In The Eleventh International Conference on Learning Representations , 2022.

Junyoung Chung, Caglar Gulcehre, Kyunghyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. In NIPS Workshop on Deep Learning , 2014.

A. Graves and A. Graves. Long short-term memory. In Supervised Sequence Labelling with Recurrent Neural Networks , pages 37-45. Springer, 2012.

Felix A Gers, Jürgen Schmidhuber, and Fred Cummins. Learning to forget: Continual prediction with lstm. Neural Computation , 12(10):2451-2471, 2000.

Shiva Kaul. Linear dynamical systems as a core computational primitive. Advances in Neural Information Processing Systems , 33:16808-16820, 2020.

Antonio Orvieto, Soham De, Caglar Gulcehre, Razvan Pascanu, and Samuel L Smith. On the universality of linear recurrences followed by nonlinear projections. arXiv preprint arXiv:2307.11888 , 2023.

James Bradbury, Stephen Merity, Caiming Xiong, and Richard Socher. Quasi-recurrent neural networks. In International Conference on Learning Representations , 2016.

Shuangfei Zhai, Walter Talbott, Nitish Srivastava, Chen Huang, Hanlin Goh, Ruixiang Zhang, and Josh Susskind. An attention free transformer. arXiv preprint arXiv:2105.14103 , 2021.

Lin Zheng, Jianbo Yuan, Chong Wang, and Lingpeng Kong. Efficient attention via control variates. In International Conference on Learning Representations , 2023.

Songlin Yang and Yu Zhang. Fla: A triton-based library for hardware-efficient implementations of linear attention mechanism. https://github.com/sustcsonglin/ flash-linear-attention , 2024.

Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring a broad discourse context. arXiv preprint arXiv:1606.06031 , 2016.

Jian Liu, Leyang Cui, Hanmeng Liu, Dandan Huang, Yile Wang, and Yue Zhang. Logiqa: a challenge dataset for machine reading comprehension with logical reasoning. In Proceedings of the Twenty-Ninth International Conference on International Joint Conferences on Artificial Intelligence , pages 3622-3628, 2021.

Hector Levesque, Ernest Davis, and Leora Morgenstern. The winograd schema challenge. In Thirteenth International Conference on the Principles of Knowledge Representation and Reasoning , 2012.

Christopher Clark, Kenton Lee, Ming-Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers) , pages 2924-2936, 2019.

Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 34, pages 7432-7439, 2020.

Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics . Association for Computational Linguistics, 2019.

Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM , 64(9):99-106, 2021.

Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05457 , 2018.

Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing , pages 2381-2391, 2018.

Leo Gao, Jonathan Tow, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Kyle McDonell, Niklas Muennighoff, et al. A framework for few-shot language model evaluation. Version v0. 0.1. Sept , page 8, 2021.

Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Processing Systems , 35:16344-16359, 2022.

A1 Related Work

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.

A1.1 LinFormer

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 ] .

A1.2 SSM

SSM 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 = Nd = 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 ∆ ) .

A1.3 LinRNN

LinRNN 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 d k = d v = H = d except LRU [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, y 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 · σ ( x t W a + b a ) where θ α = [ A , W a , b a ] . Then k t = √ 1 -α 2 t 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 y t = [ 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 = Nd = H, d k H = 1 , where N is a small expansion. α t = A is the fixed decay and v t = f v ( x t , θ v ) = x t B where θ v = B . Furthermore, q t = k t = g t = 1 and W O = C .

A1.4 Approximation to Softmax Attention

Approximation 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.

A2 Analysis of Optimal Linear Approximation Conditions

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.

A2.1 Analysis of Dynamic Memory Ability (C1)

Proposition A2.1. Only models with dynamic decay can satisfy C1 (Dynamic memory ability). Let S t ∈ R d k × d v be hidden state of general linear attention (see Sec. 3). At a time t , the information about v t 1 , . . . , v t d k is successfully stored in S t ( t 1 , . . . , t d k ≤ t ), which means ( S t ) i, : = v t i . When the new important input v t +1 arrives, only models with dynamic decay can substitute historical unimportant v t 1 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):

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 let k d k +1 = 1 .

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.

A2.2 Analysis of Static Approximation Ability (C2)

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 t 1 , . . . , 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.

where 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 ) .

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, we need to prove following equations hold (heads are omitted for simplicity):

Simple case. We first simplify the problem via i) setting dimension size of Q , i.e. d k = 1 and ii) considering only a given time t ( 1 ≤ t ≤ d k ) and its autoregressive attention distribution p t = [ p ts , s = 1 , . . . , t ] ∈ R 1 × t . We need to prove the following equations hold with bounded parameters, as a foundation conclusion:

We discuss models with fixed/dynamic decay or without decay. For cases with dynamic decay, we further consider using dynamic decay to replace key. Query's role is discussed later.

i) Models with fixed decay α result in unbounded parameters . We can solve Eq. (A6) and derive:

Because α ∈ (0 , 1) , the last term 1 α t -s is unbounded, which leads to unbounded k s . Hence Eq. (A7) can not be satisfied. The result is intuitive because the fixed exponential decay makes it hard for a token to attend to distant information, leading to an excessively focused attention map rather than an arbitrary attention map. This conclusion about fixed decay is general and is unrelated to the usage of Q or K .

ii) Models without decay ( α = 1 ) can satisfy Eq. (A6) . We can solve Eq. (A6) and derive:

At any time s , fix an appropriate | q t | ≤ C q , and we can obtain bounded solutions for all k s .

iii) Models with both dynamic decay α t and key k t can satisfy Eq. (A6) , but have parameterredundancy. We can solve Eq. (A6) starting from time t to time 1 :

At any time t -k , parameters α t -k +2 , · · · , α t have already been determined. It is equivalent to use two free parameters k t -k , α t -k +1 to concurrently approximate just one variable p t t -k q t , which is redundant. Dynamic decay and key actually have similar function, i.e., to balance historical memory and new input. Fix a | q t | ≤ C q , we can obtain many possible bounded solutions of all k s and α s .

iv) Models with dynamic decay α t and without key k t can satisfy Eq. (A6) . Inspired by the parameterredundancy in iii), we further consider a linear model with only query and dynamic decay, i.e., using dynamic decay 1 -α t to replace k t 's role to approximate softmax attention map. We further assume q t = 1 , which is bounded, and define S t = ∑ t s =1 p ts , which satisfies S t ∈ [0 , 1] and S t ≥ S t -1 .

Solve Eq. (A6) starting from time t to time 1 :

At any time t -k , we can derive a closed-form and bounded solution α t -k = S t -k -1 S t -k ∈ [0 , 1] .

General case. Aiming to solve Eq. (A4) and Eq. (A5), we generalize to vector version with d k > 1 and consider distribution of all time p t , t = 1 , . . . , d k . In an extreme case, let one dimension of q relate to one time, which means q t = [0 , · · · , q t , · · · , 0] . Through the analysis above for one dimension (Eq. (A6) and Eq. (A7)), we can ensure:

Actually, we can further relax the condition and view q t as a selector, which selects several channels of Hadamard product of α j and k s , and uses these channels to approximate its attention distribution. One-hot assumption of q t is an extremely hard selector. This means models without Q can not approximate softmax attention map, because without selection they can only fit one token's attention distribution theoretically.

In summary, through analysis of simple case i) to iv), models with fixed decay Λ cannot ensure bounded parameters. Through analysis of general case, Q is necessary. So we can conclude that: only models with parameters ( Q , K , Λ t ) , ( Q , K ) or ( Q , Λ t ) can satisfy C2. (For models with only Q , without decay and k t = 1 for all t , a token's attentions to other tokens are all the same, which means such models can only approximate one special attention map.)

A2.3 Analysis of Least Parameter Approximation (C3)

When model dimension d, d k , d v are fixed, parameter counts of θ q , θ k , θ α are fixed too. Thus fewest parameters means fewest parameter groups. Linear attentions with general form 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 .

Proposition A2.3. Only models with Query and dynamic decay ( Q , Λ t ) can satisfy C3 (Least parameter approximation). Models with ( Q , Λ t ) can meet C0, C1, C2 simultaneously with fewest parameters.

Proof. i) C0: all the linear models with general form (Sec. 3) own linear complexity and satisfy C0.

ii) C1: according to proposition A2.1, only models with dynamic decay Λ t have dynamic memory ability.

iii) C2: according to proposition A2.2, only models with parameters ( Q , K , Λ t ) , ( Q , K ) or ( Q , Λ t ) have static approximation ability.

So we claim that only models using ( Q , K , Λ t ) or ( Q , Λ t ) can meet C0, C1 and C2 simultaneously. And ( Q , Λ t ) with two parameter groups is least parameter approximation. That is, dynamic decay can replace Key's function to balance historical information and new input when approximation is performed.

A2.4 Conclusions of Optimal Linear Approximation Analysis

In this section, we summarize our conclusions of theoretical analysis. It is worth noting that our theory studies functions of LinAttMap, but not whether these functions will be successfully learned.

Linear approximation. The necessary conditions (C1 and C2) for LinAttMap to achieve approximation to SoftAttMap is that its implementation must include at least two parameter groups: Query Q and dynamic decay Λ t . Both ( Q , K , Λ t ) and ( Q , Λ t ) can achieve approximation.

ii) Least parameter approximation. LinAttMap utilizing ( Q , Λ t ) can achieve linear approximation to SoftAttMap with theoretically fewest parameters.

iii) Function of Query. Q can be seen as a channel selector, which selects several channels of Hadamard product of α t and k t , and uses these channels to approximate attention map. Thus it is indispensable for approximation and its dimension size (also the model memory size) is the guarantee of approximation ability.

iv) Function of dynamic decay. Dynamic decay Λ t is the key to achieve dynamic memory and precise approximation, and can substitute the role of K when approximation is performed.

A3 MetaLA Architecture

We here introduce Meta Linear Attention (MetaLA), a linear approximation with least parameters to softmax attention map. We design three enhancements of MetaLA relative to the general linear attention in Sec. 3: 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 approximation ability of our model. appendix A3.1 explicates enhancement i) and design of basic MetaLA layer. appendix A3.2 introduces enhancements ii) and iii), and design of overall MetaLA architecture.

A3.1 MetaLA Layer

In MetaLA layer, we exploit 1 -α t to replace k t based on theoretical analysis in appendix A2, i.e., LinAttMap utilizing ( Q , Λ t ) is the linear approximation with least parameter redundancy to SoftAttMap and K is redundant.

General Recurrent Form. Using marks in Eq. (11)-Eq. (15), the improvement lies in Eq. (A25):

where x t ∈ R 1 × d denotes the current input. W Q , W α ∈ R d × d k , W V , W G ∈ R d × d v and W O ∈ R d v × d are learnable parameters. d ′ k/v = d k/v H and h = 1 , · · · , H is the index of heads. Here d k represents Query/Decay dimension. LayerNorm is chosen as the normalization operation. Furthermore, we choose σ = Sigmoid( · ) 1 /τ following [15], where τ = 16 is used to control the value of dynamic decay and ϕ = SiLU .

As for parameter allocation, we simply set d v = d and d k = d 2 , similar to GLA [15]. Thus, the whole number of parameters used by a MetaLA layer is 4 d 2 , the same as the softmax attention layer and smaller than other popular attention-like subquadratic models such as RetNet [14] ( 8 d 2 ) and GLA [15] ( 4 d 2 +24 d ), etc. Without usage of 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 do not sacrifice expression capacity of f α and approximation performance of MetaLA.

General Parallel Form. The recurrent form of MetaLA can be written in a parallel mode as follows:

B A in Eq. (A28) denotes element-wise division and ( Q ) t, : means the t -th row of matrix Q . In practical implementation, we utilize chunkwise form proposed in [15] for hardware-efficient training.

Attention map. Here the attention map is (( ( Q ⊙ A )( B A ) ⊤ ) ⊙ M ) . The element in the attention map matrix is as follows (heads are omitted for simplicity):

A3.2 MetaLA Transformer

Improved MetaLA. After deriving the basic MetaLA layer shown above, we consider further optimization when designing complete MetaLA block. One observation is that in most cases, setting α t = 0 and completely discarding historical information is catastrophic. Actually, in real-world scenario, the learned dynamic decay α t always closes to 1 because of many tasks' strong need to capture long-term dependencies. Hence we propose a technique called self-augmentation to enlarge a token's attention to itself while do not disrupt the model's state constitution, in order to better avert attention dilution. Moreover, an additional short convolution can be inserted before entering MetaLA layer to further enhance local interaction, motivated by Mamba [16] and Griffin [18]. Combining these two techniques, the improved MetaLA is shown as follows, the main improvements lie in Eq. (A31) and Eq. (A35):

As for self-augmentation, without changing the composition of hidden state S h t , it is only added on the output process through a learnable parameter w aug ∈ R 1 × d k , which is then divided into heads like other parameters do. σ aug ( · ) is a nonlinearity to control the magnitude of augmentation term and avoid covering with the original attention. We choose Sigmoid( · ) in this paper.

The proposed design has two advantages: First, it maintains parallel computing as shown in Eq. (A38); Second, it only enhances current token's own attention and does not change the attention of future tokens to the current token ( p ts , s < t ), because we only change o h t while maintaining S h t unchanged like that in original MetaLA layer Eq. (A25). Thus the separation of output and memory is realized.

where W t, : = w aug and sum( Q , dim = 1) means calculate sum for each row of matrix Q . Other marks are defined same as Eq. (A29).

MetaLA block. Regard each layer as a function mapping the input X ∈ R n × d to the output Y ∈ R n × d where:

following Transformer structure, now the Token Mix mechanism Eq. (A31)-Eq. (A37) can be integrated as follows:

and the Channel Mix part (GLU [3]) can be written as follows:

which can be integrated as:

where X l and X l +1 ∈ R n × d refer to the input and output of block l in MetaLA. We choose XNorm = LayerNorm . Stacking several MetaLA blocks above, we can derive complete MetaLA Transformer as a Linear Foundation Model, as shown in Fig. A1.

A4 Experimental Details

Associative Recall (AR). Following [36], we train two layer models with a Transformer backbone that interleaves token mixer and channel mixer. Learning rates are swept by np . logspace( -4 , -2 , 4) for sequence length 256 and additional np . logspace( -5 , -3 , 4) for length 512, and maximum test accuracy is reported. For GLA [15] and MetaLA, we set H = 2 and d k = d v = d . The kernel size of short convolution of MetaLA is 2 .

Language Modeling. For 360M/1.4B model, we train it from scrach with a total of 15B/300B tokens on 16/32 A100 GPUs at a learning rate of 3e-4/2e-4 with batch size 0.5M/2M. Both models maintain a length of 2048 and are trained using fp16. The training setup of baselines [15, 16, 44] of 360M MetaLA are aligned with MetaLA configurations. For the 1.3B MetaLA, we compare it with publicly available models [14, 15, 16, 20, 44]. To maintain a fair comparison between linear models, we trained Mamba from scratch using the same settings with MetaLA on 100B tokens. For GLA and Retnet, we adopted the open-source checkpoints in FLA [79]. For HGRN and Pythia, we used the official open-source checkpoints. We implement all the pretrain experiments with GPTNeox [45]. We evaluate our models on SuperGLUE benchmark [38] and Common-Sense Reasoning benchmark including LAMBADA [80], LogiQA [81], Winograd Schema Challenge (WSC273) [82], BoolQ [83], PiQA [84], HellaSwag [85], WinoGrande [86], ARC-easy (ARC-e), ARC-challenge (ARC-c) [87], OpenBookQA [88]. We report perplexity (ppl) and accuracy (acc) on LAMBADA, accuracy normalized by length on HellaSwag, ARC-challenge and OpenbookQA, and acc on the other subtasks. For SuperGLUE benchmark, we report F1 score on CB, Exact-Match (EM) score on MultiRC, and accuracy on the other subtasks, following the original work. The LM evaluation harness [89] is used to implement all evaluations.

Long Range Arena. We evaluate the long sequence modeling capability of our model on LRA task. We use the adamW optimizer and cosine warmup scheduler. We set the head size be 4 for group

×N

Add

Channel

GLU

Mixer

XNorm

Add

MetaLA

Token

Layer

Mixer

XNorm

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.

A5 Additional Experiments

Additional subtasks for Common-Sense Reasoning. We extend more subtasks of Commonsense Reasoning for 1B4 MetaLA. Additional experimental results are shown in Tab. A5.

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 MetaLA 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 eval-

uate several models on MAD [41], a collection of six synthetic tasks predictive of scaling laws,

Training throughput

20

(S/>)

15

Tokens per second(

10

Transformer++

5

Mamba

GLA

MetaLA

0

2048/8

4096/4

8192/2

16284/1

Traininglength/Batchsize

GPUmemoryusage

80

60

(GB)

40

Transformer++

20

Mamba

GLA

MetaLA

2048/8

4096/4

8192/2

16284/1

Traininglength/Batchsize

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.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

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.

2. Limitations

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.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: In Sec. 4 and appendix A2, we provided the definitions and detailed proofs of our theory.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: In Sec. 6, we showcased our experimental setup and the results, and in appendix A4, we further elaborated on the details of all experimental setups.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [No]

Justification: We have not open-sourced the code, but we have clearly outlined the details of our experiments in Sec. 6 and appendix A4 to ensure the reproducibility of the results.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We presented the detailed settings for training and test in Sec. 6 and appendix A4.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: We did not conduct experiments that required reporting significance.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: In Sec. 6 and appendix A4, we provided details about the computer resources used for the experiments.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The research conducted in the paper complies with the NeurIPS ethical guidelines in all aspects.

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).

10. 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).

11. 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.

13. 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.

14. 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.

15. 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nProposes 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.\n\nSTRENGTHS:\n- Addresses an  important and relevant topic of analyzing efficient sequence models that attempt to replace Softmax Attention\n- Systematically comparing  and unifying the formulations of the various linear attentions and linear RNNs/SSMs is useful \n- The tables classifying similarities and differences of the models are nice\n- The theoretical results in the appendix appear to be sound\n- The experiments performed suggest promising results\n\nWEAKNESSES:\n- The general presentation of the \"optimal linear approximation\" and theoretical analysis is dense and confusing and obscures the contributions. \n  - 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.\n  - Once you have moved these Proposition statements to the Section 4, I would recommend rewriting/reorganizing Section 4 framed around these proposition statements\n  \n- 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. \n- 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. \n- 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.\n\nQUESTIONS:\nMy main questions and issues are listed above.\n\n- Line 287: typo, \"Value\" is misspelled in the text\n\nLIMITATIONS:\nSufficient", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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. \nThese enhancements aim to improve the approximation accuracy of softmax attention and outperform existing linear attention methods.\n\nSTRENGTHS:\n- The paper presents a comprehensive unification of various linear attention models (LinFormer, SSM, LinRNN), offering a deeper understanding of their underlying mechanisms and differences.\n- The proposed MetaLA model is theoretically grounded based on the proposed framework.\n- The authors provide empirical evaluation of the MetaLA model on NLP tasks, demonstrating its effectiveness and robustness.\n\nWEAKNESSES:\n- 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.\n- Authors provide empirical results mainly for classification tasks. It is not clear how their model performs in the text generation scenario.\n- 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.\n- 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.\n- Since the authors do not provide code for the MetaLA model, it is hard to assess the practicality and ease of implementation.\n\nQUESTIONS:\n- 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.\n- 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?\n- 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?\n\n\n[1] Linear Log-Normal Attention with Unbiased Concentration\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-28"}

comments: 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.

STRENGTHS:
Strong build-up toward their method from well-curated baselines. Their method shows high performance empirically.

WEAKNESSES:
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:
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:
It feels minor improvement on top of previous linear attention alternative studies.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-09

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1, This work unifies three widely-applied linear-complexity substituitions (LinFormer, SSM, LinRNN) of Transformer into the same framework, including a recurrent form and a parallel form, which enables higher-level studies and analysis.\n\n2, This work proposes three criteria for evaluating a linear-complexity model (memory, approximation, parameter), which brings insights and targets for future developments of linear-complexity models.\n\n3, Combining 1 and 2, this work evaluates whether each of the three existings models satisfies each criteria, and shows that every model does not fulfill all requirements. Based on this insight, they propose a MetaLA Transformer model that satisfies all three criteria and performs well in multiple tasks.\n\nWEAKNESSES:\n(After rebuttal: The authors carefully addressed my main concerns.)\n--------------\n\nPlease correct me if there are any mistakes.\n\n1, 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*.\n\n2, 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).\n\nQUESTIONS:\nThe authors are encouraged to disagree or discuss on my questions/comments.\n\n1, 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$?\n\n2, 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?\n\n3, On Page 5 Line 147, why this softmax-generated expressive attention map requires infinite states?\n\n4, On Page 5 Line 177, what is the condition (or consequence) of \"good enough\" inputs?\n\n5, 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?\n\nLIMITATIONS:\nLimitations are adequately discussed in the Section of Conclusion and Discussion.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-13"}