MetaLA: Unified Optimal Linear Approximation to Softmax Attention Map
Authors: Yuhong Chou
Man Yao
Kexin Wang
Yuqi Pan
Rui-Jie Zhu
Jibin Wu
Yiran Zhong
Yu Qiao
Bo XU
Guoqi Li
Venue: NeurIPS 2024 oral (2024)
Keywords: Self-attention
Linear attention
Efficient
Softmax attention
Foundation Model
Language Model
Primary area: deep_learning_architectures

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
Yuhong Chou 2, Man Yao²-, Kexin Wang', Yuqi Pan?, Ruijie Zhu?, Yiran Zhong Yu Qiao", Jibin Wu Bo Xu², Guoqi Lift
The Hong Kong Polytechnic University
Institute of Automation, Chinese Academy of Sciences UC Santa Cruz Shanghai AI Lab
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 answer this question by finding the best linear approximation to softmax attention from 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: pettp://githhee com/ BIIIaabIaabbML...
1 Introduction
Transformer with softmax attention ① benefits from efficient parallel training and exhibits impressive performance deep learning applications 2 3 [5 7). but it suffers from the quadratic growth of computation cost to the input length 8. Linear recurrent models, such as LinFormer (9. SSM [00, and LinRNN [11), are expected to achieve linear substitution of Transformer. The original intention of LinFormer to replace softmax attention, which exploits the kernel approach to decompose
softmax operation; typical work includes TransNormer [12] [3], RetNet [14,, GLA (15). On
the other hand, SSMs, such as S4 (1O] and Mamba [16), are models inspired by the classical state-space approach, which enjoys sub-quadratic training and inference like either 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 nonlinear dependencies between hidden states. It is natural to think that they are different types of models, models have different origins and forms.
since these LinFormer/SS//Lnneee
This work breaks this perception and abstracts existing LinFormer/SSM/Lnneee unified linear attention form, which has the following significance: i) Facilitates understanding
models into a
Equual contribution. yuhong.chooe www.@@/oonnet.ooy..uuu man.yao@iaaac.co Cooreesponding author, guoqi li@ ia..c..nn
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

the key designs of existing linear models. Through the unified form, we demonstrate that the main difference between Lifoommr/SSM/LinRRY the hidden state size, how to maintain the hidden state, and how to perform parameter mapping. ii) Links LinFormer/SSM/Linnee in terms of functionality. The recurrent inference complexity of softmax attention is ((n), which can also be regarded as the maintenance of hidden state with infinite size. Linear models with (1) inference complexity are hoping to achieve the same functionality as softmax attention using a fixed hidden state. Since we have unified LinForree/SSM/LinRNY 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 softmax attention
To answer this question, we define the necessary conditions for achieving "optimal linear approxima- tion to softmax attention map". First, linear attention needs to satisfy dynamic memory and static appreximation 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: linear attention map should be able to approximate any softmax attention map. According to our theoretical analysis, Query and dynamic decay necessary conditions for approximation. Thus, linear models such as TransNormer [3], RetNet [14), RWKV-4 177. LRU [19), HGRN 200). H3 [21] S5 [22]. cannot achieve approximation of the softmax attention functions. Second, the Key matrix is not required to achieve approximation, so Mamba [66 are not optimal parametric approximations.
and GLA [5]
We then propose the MetaLA module, which can satisfy the necessary conditions for optimal linear approximation to softmax attention. Metal.A makes three enhancements: i) Removes the unnecessary Key matrices; ii) Employs seff-augmentation to enhance the token's attention itself, which avoids attention dilution [2: 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.. qt. kt), 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., qe the t-th row of Q.
Softmax Attention first calculates an attention map Sttt (Q, K) through Q (Query), K (Key), and use the attention map to weight different tokens V (Value) later:
QKT
SoftAttMap (Q, K) V softmax
① .MM
V
E RR
(1)
dd
Q K = O. E R Rexds
V XWW V
Ruxds
(2)
where W Q: WK E Rexd: Wv E Rdxd. are learnable matrices, n, d, dk. dy are sequence length, model dimension, Key/Query and Value dimension, respectively. X E Raxd refers to the input. M E Rax" is a mask matrix in autoregressive tasks to prevent token from attending to future tokens. The t-th row of MoftAttMap (Q, K) is a probability distribution that represents the attention scores between token V, to others. Softmax attention in Eq. I enables efficient parallel training, but suffers from On??) time and memory complexity [9.
It uses the recurrent form during inference:
O, = E'= explqik,) Jv.
E RIxd.
(3)
E exp(qik.)
qe, k; = x, W W E Rixds;
Vi W
E RIxd.
(4)
At each time I, token mix is computed between query q: and all the keys, values before k,. v.(s S (). This "KV cache" results in O(n) time and memory complexity per token during inference.
2

---

Linear Transformer (LinFormer) is substitute for softmax attention, which can be expressed as linear dot-product of kernel feature maps
E Flg,,k,vv 0t E F(gr,k,)
F(q,,,, o(q:) (ko);
(5)
where qt, k; R1xdd and V, R are query, key and value at position which obtained in the same manner as softmax attention. F(.) is the kernel function usually constrained to be non-negative. o(-) is map function applied row-wise to Q and K. By removing the nonlinear
softmax operation, LinFormer enables inference with (1)) time and memory complexity per token. LinFormer can also be formulated in the following parallel form during training
O (
(K)
M) V
E R
(6)
which has (nn) time and memory complexity using chunkwise algorithm. Recent advances in LinFormer mainly focus on training acceleration [ [3] or improving performance [14.
State-Space Model (SSM) come from continuous- time system which maps a ID function xxt) e R another function (t) E R via hidden state h(t) E R.V In SSM, the continuous parameters can be discretized using step size A and get discrete parameters A, B, C. The resulting discrete-time system is used to model sequences x, E xn with elements 24, Ye E R via the recurrent form:
h: Ah,-1 Bri;
y/ Ch.,
in autoregressive inference with (1) time and memory complexity per token. The linear time- invariant SSM above can be unrolled and computed using the long convolution with kernel K
K = CB,CAB "A
B
E R11x, y K= x.
(8)
where represents casual convolution operation
24 which enables parallelizable training utilizing Fast Fourier Transforms, resulting in On log n) time and ((n) memory complexity during training,
When handling vector sequences X, Y E Rdxx, SSMs are applied individually on the d channels.
Typical SSMs (S4D (25) DSS 26). 21 2 employ fata-indeppndent structured transition matrix A or special initialization strategies HiPPO 27 to efficiently enhance long-range dependencies. Mamba [16] advances SSMs by introducing data deteaddenndent 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, LinR NNs eliminate the nonlinearity within the recurrence and employ element- wise product instead of matrix
multiplication (8).8)) Typical LinRNN such as Gated Linear Recurrent Unit (GLRU) (20 29 is
G, ir, C. o(x W + by), (x, W; + b.), W bb) E R Ixd
(9)
h, - b, + i, c,, lxd
(10)
where x,, h, denote input and output, fr, i, are forget and input gates as traditional RNNs, O is wlsment-wise multiplication. Linear RNNs have O(1) time and memory complexity per token during inference. Since Eq. (T0) removes nonlinearity, enables parallelized training using parallel scan [[]] with only (n) time and memory complexity. Recent works have made effort to explore effective
recurrence (LRU 19), RWKV [7] or gating mechanisms (HGRN 200, Griffin [8]
3 General Form of LinFormerrSSM//LnRRNY 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, SinFormer/SSM/LinRR
have limited hidden states: linear attention with Φ (k,)v; E Rd SSM with h, R xd, linear RNNs with h, E R1xd usual. Inspired by this fact, we unify LinForrerSSMMLLinRNN mechanisms
where dę N in the form of linear attention, formally containing Query, Key, and Value matrices (see Fig, I
3

---

Parallel Form
Recurrent Form
Wo
W
XNom
XNomm
element-wise multiplication
Sy
S
element-wise addition element-wise division
outer-product
G
k;
V:
9:
8:
matrix multiplication
f
fxx
fg
Eg
4
Figure I: General Form of LinFormer/SSWLinnRNY Mechanisms. The general form equips with two modes parallel and recurrent computation which enjoys both training and inference efficiency.
General Recurrent Form of LinForner/SSMLLnnNY is:
q (x,, k; fx, (x) at fa (x(,
ER!xd
(11)
V; (x,, 0,), g: = fg (xx, 0g) E Exld,
(12)
S diag(a 1 (k') vv
ER
(13)
0, X.Nom(concat q St
.. q (
ERixd,
(14)
Y1 (o, o g.) Wo
E
(15)
where X2 E R Rixd is the t-th input. qr, kg Vz, Qt gt, S, are query, key, value, decay, output gate, hidden state respectively. f//M/o are functions that map x, from R
R xixd, Pa/k/o are the corresponding parameters to be trained. Similarly, fu/g map x, from R R.xd to R R1xdd and 00/g trainable parameters. In Eq. (I3), qe, 9r, k,, V, divided into H partitions (heads), where da/"
H - and = 1, is the index of heads. Each head maintains a hidden state S The diagonal matrix diag(o,) denotes the decay of past state. k. represents the acceptance for the input token v². The hidden states are 2D matrix once d' o 1. In Eq. (14). to turn back to ID shape, the qf dot- dot-product with S/ then the concat and normalization operations are
d//=
operation is necessary followed. XNorm denotes any kinds of normalization. In Eq. (15). gate machanism optional for O, while the dimension should be projected back to d from d, through Wo E Rd.
From a functional view, Eq. (l3) represents the update process 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. (5) represents gate and projection operations to get the final output.
General Parallel Form of LinFormer/SSS//iinRR can be written as follow:
0 LinAttMap (Q,AttMapQQQKK)))) (Q,K) ((QQ))) (
O M)V,
(16)
[1ii5 (Q (g/k/v)) A II ag. M.j. = 1,i>j..
(17)
K denotes element-wise division, (Q).. is the I-th row of Q, and LinAttMap (Q, K) is the attention
A map. Each element in the attention map matrix is as follows (heads are omitted for simplicity):
q: (III as) @ k.) s sstt.
LinAttMap (Q, (Q.
(18)
0,
As shown in Tab. I the main differences between various linear models are parameter functions fa/kfo/a// and dimension settings dę, d,, H. We give details in appendix AT on how to derive LinFormersSSMMsllnRRNNN from our unifted form, which is termed as "Linear Attention (LinAtt)"".
LinFooreer/SM//iinnee models have different origins, so they have different optimization per- spectives and various hidden state sizes: i) LinFormer originate from approximation of vanilla softmax attention. They focus designing better kernel function ø, i... to optimize fq. fx; They have

---

Table 1: From our general form to existing linear models (* indicates the bias term is omitted).
LinFormer
LinRNN
SSM
Modcls
GLA
TrNorm 12
GLRU 20
RWKV-4 17
Mamha 16
25
(x,, 0g)
We Q
((x,Wa) ((xWR) R)
W.
1
[x 0e)
X, K
a(x, W )
exp (x, W R) x, x.W
(x,
1
(x,, 0.)
X, W.W
X. W
o(x, W
X;
x, B
La (x,, Om))
a (xx, W W
exxp (j0) W U
a W 5)*
exp (-W) W)
exp (A,4)
exp (AA)
(x,,x,0))
o(x, w)
6(x, 9)
o ,x;
1
= d/2 d = d
d = d H Nd= d' =1
Dimension
de de d
d₂ d = d= H
dg =N
relatively large matrix hidden state whose size (d / 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 fa: They have matrix hidden state of moderate size Nd,,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 fo. fx. fgi They have ID vector hidden state whose size (d,)) is relatively small. Despite these differences, they all try to design better parameter functions fa/k/v/afs and maintain limited hidden state S..
4 Optimal Linear Approximation to the Softmax Attention Map
We here discuss the optimal approximation of LinAutap to SoftAttMap based on its general form. The function 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[3]]. etc., optimize the model with the goal of approximating the value of SoftAttMap. In contrast, this work investigates the functional approximation of SoftAutaa,, 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 [(x,, x, (0) : R1xd xd R1xd R. used to compute attention score between any x and x, (tokens), with parameters 0, 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., (nn) space and time complexity during training and (1) space and time complexity during inference. ii) Dynamic memory ability. When handling inputs sequentially, f(x,, x, (0) 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 Pts- there must exists bounded 0 such that fx,,, x, (0) Pts. Yt, 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 (CO) 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 Cl condition, suppose the information about v., Via, is successfully stored in S, (l,,..., ld. S f), we will check whether the model can substitute unimportant Vt, when the new important input Vx+1 arrives.
For the C2 condition, Eq. (18) illustrates the LinAuMMp only relate to q = [x( 0z); k, (x, 0x), ag (x, 0.). Denote decay A, = diag(a.). Assuming the inputs good enough
5

---

Table 2: A review of existing linear models. According to definition 4.1
existing linear models all have some deficiencies: Models without dynamic decay A, have no ability to memorize dynamically (not satisfying C1); ii) LinRNNs lack the selection ability brought by Q (not satisfying C2), and the approximation ability is poor owing to the small hidden state; iii) Models with K have redundant parameters (not satisfying C3), which probably leads to higher learning difficulty.
CO
CI
C2
C3
Models
Softmax Attention
X
V
Transformer
GPT 2 Llama
Q,K,
X
X
x
RetNet [1]. TransNormer 12] S4D
Three Parameter Groups
25
H3
21
DSS
26
Q K, A
x
GLA [5]. Mamba
Q K
V
x
linear Transformer RFA Performer (30), cosFormer 32
K, A
V
x
RWKV4 17
Two Parameter Groups
Linear Attention
K, A,
GLRU 20
Q A
x
Q A
MetaLA (This Work)
Q or K
x
One Parameter Group
V
x
x
LRU
19
22
HGRN 20
Griffnn [8]
A
T-RNN
33
0g, 0a) to solving [Q K. Ar). We focus on approximating the attention scores between stored tokens, and the problem simplified via: i) setting query dimension dg 1; ii) considering only given time / and its attention distribution pa [Pts, 1,....4] 1] E Rixt, Then, C2 is proved by the following equations holding with bounded parameters, as foundation conclusion:
and the functions (fg: fa) are expressive enough, we can shift from solving 0g,
K.) = =e
II aj)ks Pes, Vs = 1,...,4 f=0+1
(19)
s.t. t. 1ael Cq. Ikol S Ck, a, [0,1], Vs L....... 02, For bounded inputs X, bounded parameters (0,, 0x 0a) are equivalent to bounded (Q, K, At).
(20)
Afterwards we will generalize to vector version with di > and consider distribution of all time (pr, = I,...,,, dg). This is done by viewing qr as a channel selector.
For the C3 condition, least parameters mean the fewest parameter groups (Q, K, A,) when d, dr, d. are fixed. Due to space constraints, the detailed analysis in this Section is provided in appendix [2
Conclusions of Optimal Linear Approximation Analysis. Linear approximation. The nec- essary conditions (C1 and C2) for LinAuMap to achieve approximation to SoftAutMap is that its implementation must include Q and dynamic decay Ar. Both (Q, K, Az) and (Q, A,) can achieve approximation. ii) Least parameter approximation. (Q, A,) has fewer parameters (i.e., K is not necessary), the model dimensions are fixed. iii) Function of dynamic decay. A, is the key to achieve dynamic memory. iv) Function of Query. Q can be seen a channel selector which selects several channels of Hadamard product of A, and 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 ( QQ, K, As) all together, ii) Exploiting (Q, K). (Q. A.) (K, A,), iii) Employing only one of Q, K, A₁ Considering decay can be either dynamic fixed, here we use subscript to
distinguish, i.e., A/A, denote fixed/dynamic decay. According to definition 4. they have different
degrees of deficiencies: i) Models without dynamic decay such as RetNet 14 TransNormert [2]
[FA cannot memorize dynamically; ii) LinRNNs such as RWKV-4 [7], HGRN 20),
Griffin [8

---

Table 3: From general recurrent linear form to our MetaLA.
Models
fg (g)
fa(x,, (x) Ow) fo (x,,
02)
fg(x,,00) 0g)
Dimension
Metal
x, WWo
1 1-0 x, Wy W.) @(x,Wa bc) d, d
lack the selection ability brought by and the approximation ability is poor due the small hidden state; iii) Models with K such as Mamba[]6]. GLA[I3] have redundant parameters, which probably leads higher learning difficulty. Thus, none of the existing linear models meet all C1/C2/C3 conditions. These analyses also inspire us that ignoring the functional approximation of softmax attention does not enable the approximation of softmax attention values.
MetaLA Transformer
Transformer is stacked by 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 MetaL. module as the token mixer. Due to space constraints, the architecture of Metal. Transformer is given in detail in appendix [A3] Here we focus describing the three enhancements (see Fig. 2 i) The Key matrix not used, which is based on theoretical analysis. ii) Self-angmentation and iii) Short convolution are two other optional techniques to further enhance the modeling ability of our model.
of MetaLA relative to the general linear attention Sec. 3
i) The Key matrix is not used. We exploit 1 replace k,, based on theoretical analysis in Sec. and appendix A2 i.e., dynamic decay A, 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 improve- ment is:
to
Ye
auu
ap
Ve, 9:
S
Se
s = diag(a)) 1 (1 (1-aa) v., E Rd
a,
C2
Vt
9r
8:
La
[p]
fg
fg
(21)
which can be trained in a parallel form in Eq. (T6). The only difference is that K is replaced by B and
Convid
X
Br: = 1- aq With usage of A, and Q, MetaLA can achieve linear approximation to SoftAttMap. Without usage of K (WKl 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 fa and approximation performance of Metal..
Figure Recurrent form MetaLA. We mark all three enhancements in red.
ii) Self-augmentaatinn can enhance token's 's attention to itself, avoiding attention dilution 112
of = q St + Gang (q (w aug O (1 - (o(-a,)))v Tv.)) R R
(22)
Without changing the hidden state S 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
Wang Rlxd 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 addi- tional short convolution be in-
Seq_len 256
Seq_len: 512
1.00
serted before entering the MetaLA
Accuracy
0.75
Based
layer to enhance local interaction further, motivated by Mamba [6] and Griffin [18].
GLA
0.50
RWKV
0.25
Mamba
0.00
MetaLA
64 128
256 512 64
256 512
6 Experiments
Model dimension Model dimension
We conduct comprehensive eval-
Figure 3: Accuracy (%) on the synthetic MQAR task.
uation of MetaLA to validate its
7

---

Table 4: Performance Comparison on SuperGLUE.. PS: parameter size (billion). T: tokens (billion).
I means the results reported by 20) For baselines that need be compared, if they do not have public checkpoints, we train and test them under identical conditions with MetaLA. MetaLA,; MetaLA with tied embedding trained using 100B tokens. MetaLA,, MetaLA trained with 300B tokens.
Models
PS
T
WSC
WIC
RTE
CB
MULTIRC
BOOLQ
COPA
AVG
Pythia
0.41
15
36.54
50.00
52.35
39.29
0.31
61.99
62.00
43.21
Mamba
0.37
36.54
50.31
52.71
42.86
2.52
58.78
64.00
43,96
GLA
0.36
36.54
49.84
53.07
41.07
0.42
53.49
66.00
42.92
Metal.
0.36
15
36.54
50.00
52.71
42.86
0.31
58.96
67.00
44.05
Pythia
1.4
300
36.54
50.00
53.07
35.71
0.94
60.73
72.00
44. 14
HGRN!
1
100
40.38
50.78
53.43
42.86
3.04
58.69
70.00
45.60
Mamba
1.4
100
39.42
50.94
55.23
26.79
1.15
53.27
73.00
42.83
RetNer
1.3
36.54
50.00
52.71
46.43
2.52
60.21
68.00
45.20
GLA+
1.3
100
36.54
50.16
53.07
37.50
0.31
61.04
69.00
43.95
MetaLAa
1.3
100
49.04
51.25
55.60
37.50
1.78
55.50
70.00
45.81
MetaLAa
1.4
300
62.50
51.88
49.10
48.21
1.57
56.27
75.00
49.22
Table Performance Comparison on Commonsense Reasoning. 'idiccates testing using open- source checkpoints. HS: HellaSwag. WG: WinoGrande. OBQA: OpenhookQA.
Models
PS
T
LOGIQA
WSC273
BOOLQ
PIQA
HS
WG
ARC-c
OBQA
AVG
Pythia
0.41
15
21.81
57.51
61.99
63.66
33.15
51.78
22.78
28.60
42.66
Mamba
0.37
15
20.43
56.78
58.78
64.80
33.98
49.80
22.87
29.20
42.08
GL.A
0.36
15
23.04
56.78
53.49
63.55
32.00
52.10
22.78
27.40
41.39
Metal
0.36
22.43
58.24
58.96
63.82
32.18
53.12
23.38
28.00
42.52
Pythia*
1.4
300
21.35
72.89
63.12
70.89
51.98
56.99
28.41
33.20
49.85
HGRN
1
100
22.43
58.97
58.75
71.00
48.05
$1.14
28.07
31.80
46.28
Mamba
1.4
22.73
68.50
53.27
71.44
48.63
53.59
29.01
31.80
47.37
RetNet
1.3
100
22.73
63.74
60.21
69.53
48.39
53.28
26.19
30.80
46.86
GLA=
1.3
100
21.81
63.00
61.04
70.08
48.00
51.93
28.33
31.40
46.95
MetaLAa
1.3
100
21.81
65.93
55.50
70.02
47.32
55.01
27.47
33.00
47.01
MetalA,
1.4 300
21.35
73.63
56.27
72.25
53.58
58.17
30.03
34.60
49.99
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
capabilities as a foundation model. i) MQAR 36..
on the Pile 377 dataset and evaluation 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 AS including: v) Scalability. We extend MetaLA
3B parameter scale and 300B data scale for preliminary validation. vi) Retrieval and long context abilities. We evaluated MetaLA's retrieval performance the MAD tasks [I), and its effectiveness in handling long contexts the Needle 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. shows that MetaLA outperforms other linear models, which have three parameter groups (Mamba 16). GLA [3), 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
challenging settings.
Language Modeling, We train two scales Metal.: 360M/1.44 on the Pile dataset. For baselines of 360M MetaLA, train them from scratch aligned with our configurations. For the 1.3B MetaLA,

---

Table 7: Performances Comparison on the Long Range Arena. We cite baselines from HGRN 20
Method
Oi
Text
Retrieval
Image
Pathfinder
Path-X
AVG.
Transformer
38.37
61.95
80.69
40.57
65.26
47.81
S4 I0]
59.60
86.82
90.90
88.65
94.20
96.35
86.09
DSS-softmax
26
60.60
84.80
87.80
85.70
84.60
87.80
81.88
TNN 46
61.04
87.90
90.97
88.24
93.00
96.10
86.21
SS [22
62.15
89.31
91.40
88.00
95.33
98.56
87
Mega 47
63.14
90.43
91.25
90.44
96.01
97.98
88.21
SGConv 48
61.45
89.20
91.11
87.97
95,46
97.83
87.17
LRU 19
60.20
89.40
89.90
89.00
95.10
94,20
86.30
HGRN
20
59.95
88.14
94.23
88.69
92.92
97.50
86.91
Mamba
38.02
82.98
72.14
69.82
69.26
67.32
66.59
MetaLA(ours)
59.34
89.27
91.28
91.88
91.66
96.57
86.67
Table 8: Ablation studies. Ablation study results on the 360M model trained for 15B tokens. We compared the model variants zero-shot experiments of the Comparison Commonsense Reasoning benchmark. HS: HellaSwag. WG: WinoGrande. OBQA: OpenbookQA.
Models
LOGIQA
WSC273
BOOLQ
PIQA
WG
ARC-c
OBQA
AVG
MetaLA
22,43
58.24
58.96
63.82
32.18
53.12
23.38
28.00
42.52
MetaLA w/o selfaug
21.81
58.61
$7.52
64.47
32.56
49.41
23.89
29.00
42.16
MetaL. w/o conv
22.58
51.65
49.36
52.07
25.82
51.22
26.54
28.80
38.51
Metal. w/ key
21.20
57.88
49.11
63.00
32.99
50.99
23.63
27.60
40.80
we compare it with publicly available models 14 15
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. and Tab. 5
16
20
44
Specifically, compared to the LinRNN model HGRN [20] MetaLA expands hidden state dimensions and uses Query matrix: Compared to LinFormer model RetNet 44] with fixed decay, MetaLA uses dynamic decay; Compared to SSMs like Mamba [6] 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-basee Pythia [4]. 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. shows that MetaLA achieves comparable results with SOTA linear models, demonstrating that our model effectively preserves the ability to model long sequences.
Table 6: Results on ImageNet-1k:
Model
Ace
PS (M)
Acc PS (M)
Deit
72.20
5.7
79.90
22.0
HGRN
74.40
6.1
80.09
23.7
GLA
72.47
6.1
79.23
23.5
Mamba
73.39
79.60
23.7
Image Classification. We compare MetaLA with Deit
MetaLA 75.33
6.1
80.14
23.7
49
(Transformer), HGRN 20 (LinRNN),
GLA
15
(LinFormer) and Mamba [6] (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 the 360M model trained with 15B tokens and compare the results zero-shot experiments. First, restoring the Key matrix linear attention does not improve performance while increasing parameters, supporting not necessary for approximation, and its functional role can replaced by dynamic decay. Second, the ablations of self aoff-augmentttion and short convolution demonstrate the effectiveness of our model design, i.., enhance tokens' own attention and local interactions.
theoretical result that K is
7 Conclusion and Discussion
Conclusion. We unify HinFormer/SS/LiinRR 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 LinFormerSSM/LinRNNY

---

cannor achieve optimal approximation. Consequently, we propose the MetaLA architecture, which can achieve functional upproximation of softmax attention with least parameters. The performance on various types of tasks verifies the effectiveness of Metal.A
Discussion and L Limitation. Here, we discuss key questions about approximation perspectives. 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 approxination to the seftmax 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 L.R.. Similar findings were also reported in previous work [13 15 [6]. 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.
8 Acknowledgements
This work was partially supported by CAS Project for Young Scientists in Basic Research (YSBR- 1 116), National Distinguished Young Scholars (62322603), National Natural Science Foundation of China (62236009, U22A20103, 62441606, 62406322), Beijing Science and Technology Plan (7241 10000422401 1), Beijing Natural Science Foundation for Distinguished Young Scholars (JQ21015), China Postdoctoral Science Foundation (GZB20240824, 2024M753497), and CAAI- MindSpore Open Fund, developed on OpenI Community.
References
[1] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez. Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Informa- tion Processing Systems, pages 5998-6008 2017.
[2] 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: 33:1877-1901, 2020.
[3] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.1399) 2023.
[4] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Nikolay Bashlykov, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint vrXiv:2307.09288, 2023.
[5] 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 preprini arXiv:2305. 10403, 2023.
[6] 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.0211
2023.
[7] Haotian Liu, Chunyuan Li, Qingyang Wu, and Yong Jae Lee. Visual instruction tuning, Advances in Neural Information Processing Systems, 36, 2024.
[8] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey.
ACM Computing Surveys, 55(6):1-28, 2022.
10

---

[9] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mans: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, pages 5156-5165. PMIR, 2020.
[10] Albert Gu, Karan Goel, and Christopher Ré, Efficiently modeling long sequences with structured state spaces. In The Tenth International Conference on Learning Representations, 2022.
[1]] Eric Martin and Chris Cundy. Parallelizing linear recurrent neural nets over sequence length. In International Conference on Learning Representations, 2018.
[12] 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.
[13] Zhen Qin, Dong Li, Weigao Sun, Weixuan Sun, Xuyang Shen, Xiaodong Han, Yunshen Wei, Baohong Lv, Xiao Luo, Yu Qiao, and Yiran Zhong. Transnormerllm: faster and better large language model with improved transnormer, 2024.
[14] 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.
[15] 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.
[16] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, In First Conference on Language Modeling, 2024.
[17] Bo Peng, Eric Alcaide, Quentin Gregory Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, et al. Rwkv: Reinventing mnss for the transformer era. In The 2023 Conference on Empirical Methods in Natural Language Processing, 2023.
[18] Soham De, Samuel I. 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.
2024.
[19] A. Orvieto, S. L. Smith, A. Gu, A. Fernando, C. Gulcehre, R. Pascanu, and S. De. Resurrecting recurrent neural networks for long sequences. International Conference on Machine Learning, pages 26670-26698. 2023.
[20] 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-33222 2023.
[21] 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.
[22] Jimmy TH Smith, Andrew Warrington, and Scott Linderman. Simplified state space layers for sequence modeling. In The Eleventh International Conference on Learning Representations, 2022.
[23] 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.
[24] Albert Gu. Modeling Sequences with Structured State Spaces. Stanford University, 2023.
[25] Albert Gu, Karan Goel, Ankit Gupta, and Christopher Ré. On the parameterization and initialization of diagonal state space models. Advances Neural Information Processing 3ystems,35:3597 35:35571-33999 2022.

---

[26] Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are as effective as structured state spaces. Advances in Neural Information Processing Systems, 15:22982-22994, 2022.
[27] Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. In Advances Neural Information Processing Systems, pages 1474-1487, 2020.
[28] 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 44810 2018
[29] 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.
[30] 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.
[31] Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah Smith, and Lingpeng Kong. Random feature attention. International Conference on Learning Representations, 2020.
[32] 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.
[33] David Balduzzi and Muhammad Ghifary. Strongly-yped recurrent neural networks. Interna- tional Conference on Machine Learning, pages 1292-1300. PMLR, 2016.
[34] Iya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, and Alexey Dosovitskiy. Mip-mixer: An all-mip architecture for vision. In Advances in Neural Information Processing Systems, volume 34, pages 24261-24271 2021.
[35] 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 912669912222 2024.
[36] 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.
[37] 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 1rXiv:2101.000 10000000.0000..01 2020.
[38] 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.
[39] 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.
[40] 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 IEFF/CVF Conference Computer Vision and Pattern Recognition (CVPR), pages 248 248-255, 2009.
[41] Michael Poli, Armin W Thomas, Eric Nguyen, Pragaash Ponnusamy, Björ 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 Machine Learning.
12

---

[42] 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.
[43] 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.1886 18668, 2024.
[44] Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O' O'Brien, Eric Hallahan, Mohammad Aflah Khan, et al. Pythia: A suite for analyzing large language models across training and scaling. International Conference on Machine Learning, pages 2397-2430. PMLR, 2023.
[45] 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-NcoX: Large Scale Autoregressive Language Modeling in PyTorch. https://wwzs github..hu.... com/ wwweeeeeeeeeeeeeee gpt- 2023.
[46] 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.
[47] 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.
[48] Yuhong Li, Tianle Cai, Yi Zhang, Deming Chen, and Debadeepta Dey. What makes convolu- tional models great on long sequence modeling? In The Eleventh International Conference on Learning Representations, 2022.
[49] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Herve Jegou. Training data- efffccfffii image transformers & distillation through attention. In International Conference on Machine Learning, volume 139, pages 10347-10357, 2021.
[50] Michael Zhang, Kush Bhatia, Hermann Kumbong, and Christopher Re. The hedgehog & the poreupine: Expressive linear attentions with softmax mimicry. In The Twelfth International Conference on Learning Representations, 2023.
[51] 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.
[52] Jungo Kasai, Hao Peng. Yizhe Zhang, Dani Yogatama, Gabriel Ilharco, Nikolaos Pappas, Yi Mao. Weizhu Chen, and Noah A Smith. Finetuning pretrained transformers into mns. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing. pages 10630-10644 10643, 2021.
[53] Krzysztof Marcin Choromanski, Han Lin, Haoxian Chen, Anijit Sehanobish, Yuanzhe Ma, Deepali Jain, Jake Varley, Andy Zeng, Michael S Ryoo, Valerïi Likhosherstov, et al. Hybrid random features. In International Conference on Learning Representations, 2021.
[54] 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 Artificial Intelligence, volume 35, pages 14138-14148, 2021.
[55] 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.
[56] Wolfgang Maass. Networks of spiking neurons: the third generation of neural network models. Neural Networks, 10(9):1659-1671, 1997.
13

---

[57] 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.
[58] Man Yao, Ole Richter, Guangshe Zhao, Ning Qiao, Yannan Xing, Dingheng Wang. Tianxiang Hu, Wei Fang, Tugba Demirei, Michele De Marchi, et al. Spike-based dynamic computing with asynchronous censing-computing neuromorphic chip. Nature Communications, 15(1) (1)444444, 2024.
[59] 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 the IEEE, 112(6):544-588)) 2024.
[60] Man Yao, Jiakui Hu, Zhaokun Zhou, Li Yuan, Yonghong Tian, Bo Xu, and Guoqi Li. Spike- driven transformer. Advances Neural Information Processing Systems, 36, 2024.
[61] 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.
[62] 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.
[63] 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.
[64] Tobias Katsch. Gateloop: Fully data cata-controlled linear recurrence for sequence modeling. arXiv preprint arXiv:2311.01927, 2023.
[65] 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.
[66] Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces. In International Conference on Learning Representations, 2023.
[67] Junxiong Wang, Jing Yan, Albert Gu, and Alexander Rush. Pretraining without attention. In Findings of the Association for Computational Linguistics: EMNLP, pages 58 66, 2023.
[68] 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 Machine Learning, pages 28043-28078.
PMLR, 2023.
[69] Eric Nguyen, Karan Goel, Albert Gu, Gordon Downs, Preey Shah, Tri Dao, Stephen Baccus, and Christopher Ré. S4nd: Modeling images and videos multidimensional signals with state spaces. Advances in Neural Information Processing Systems, 35:2846-288 2861, 2022.
[70] 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.
[71] 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.
[72] A. Graves and A. Graves. Long short-term memory. In Supervised Sequence Labelling with Recurrent Neural Networks, pages 37-45. Springer, 2012.
[73] Felix A Gers, Jürgen Sehmidhuber, and Fred Cummins. Learning to forget: Continual prediction with Istm. Neural Computation, 12(10):2451-247. 2000.
14

---

[74] Shiva Kaul. Linear dynamical systems core computational primitive. Advances in Neural Information Processing Systems, 33:16808- 16820, 2020.
[75] Antonio Orvicto, 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.
[76] James Bradbury, Stephen Merity, Caiming Xiong, and Richard Socher. Quasi-recurrent neural networks. In International Conference on Learning Representations, 2016.
[77] Shuangfei Zhai, Walter Talbott, Nitish Srivastava, Chen Huang, Hanlin Goh, Ruixiang Zhang, and Josh Susskind. An attention free transformer. arXiv preprint arXiv:2105. 14/03, 2021.
[78] Lin Zheng, Jianbo Yuan, Chong Wang, and Lingpeng Kong. Efficient attention via control variates. International Conference on Learning Representations, 2023.
[79] Songlin Yang and Yu Zhang. Fla: A triton-based library for hardware-efficient im- plementations of linear attention mechanism. https://github. com/ sus poustssongli/ 2024.
[80] 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 broad discourse context. arXiv preprint arXiv: 06031, 2016.
[81] 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.
[82] Hector Levesque, Ernest Davis, and Leora Morgenstern. The winograd schema challenge. In Thirteenth International Conference the Principles of Knowledge Representation and Reasoning, 2012.
[83] 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-292222 2019.
[84] Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. Piqa: Reasoning about physical commonsense natural language. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 7432-7439, 2020.
[85] Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguisties. Association for Computational Linguistics, 2019.
[86] 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.
[87] 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.052 05557,, 2018.
[88] Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? new dataset for open book question answering. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pages 2381-2399, 2018.
[89] 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 vO. O.I. Sept, page 8. 2021.
[90] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances Neural Information Processing Systems, 35:16344-16359, 2022.
15

---

A1
Related Work
We overview several architectures three types linear models, i.c. LinFormer, SSM, and linRNN, and show how they can be specialized from the general form in Sec.
A1.1 LinFormer
LinFormer abandons the softmax form of vanilla attention, instead leveraging the dot product between Keys and Values to achieve linear complexity 99. Advances LinFormer mainly focus on designing
better kernel function 30
32
52
54
55 50 such the spiking neurons in spiking neural
network [56] 57
588 59 can be understood as efficient kernel functions 600
61 exploring architectural
optimization 13 14 introducing gating mechanism [13 31
63 64 etc. We show the general
form of several LinFormer architectures in Tab, Al For the general form of LinFormers, dy = d, d
the most common choice.
Vanilla LinFormer [9: Vanilla linear Transformers lack explicit decay. Instead, They choose to maintain two hidden states containing denominator to normalize attention scores. As for their general form, q: kr = fq/e((x,, fg//k) = o(x, W Q/K) where 0g/k = W Wo/K and Φ is the nonlinearity. V. X.W. v. There is no decay or output gate, i.c., g 1 and a, = 1.
RetNet [4] Different from vanilla LinFormer, RetNet uses positional encoding exp (jnØ) and fixed decay À to control the hidden state. For its general form, /ı/k//v = fn/k/u (x,, (x,, a/k/o)) = x, A.Wr q/ K/v where f = W W//K a = Aexp (jØ) is the fixed decay vector, where j is imaginary unit. Furthermore, Yt = oo @ Sx, W Wo where 0g WG.
Trrmerormer [2] TransNormer also uses positional encoding exp (jn0) and fixed decay X to control the hidden state. For its general form, q /k, = f//x((x (Xt, 0g/k) = 6(x; W Q/K) where 0a/k W Q/K and & is the nonlinearity. V: = x, y, and a À exp (j0) is the fixed decay vector, where j is imaginary unit. The normalization layer chosen by the paper is SRMSNorm. Furthermore, y: [o, (x) where 0g = u.
GLA 15 GLA considers a data-dependent gating mechanism for LinFormer. For its gen- x.Wq/K/V Q/K/V where lg/k/v W Wo/k/V. a sigmoid(x, W W + bb,) 1/r is a dynamic decay, where W and W are low-rank matrices. Furthermore, yt 0o SiLU(x, W b where fg b,
eral form, qr 8t/k./v
f//k//((x,. (xx/k//((x,0/kk/)))
a
Table A1: From general form LinFormers.
Vanilla LinFormer 9
Models
RetNet 14
TransNormer 12
GLA
f (xe, 0g)
6xx W Q ((x, W
X W Q
6(x, Wo) o(x, W K)
x; x Q x, WWK
f (X,, 0k
X, W K
f (X,, (%))
W V
X, W V
x, W,
X, W
fa (x1, Oa)
1
exp (j0)
expp (j0)
(x, W W + +ba)
fg (Xe, 0g)
1
(x; We)
x; Wu
(x, W b,)
d2 d/2
Dimension
d2 = d, =
d, = d
A1.2 SSM
SSM represents an alternative linear architecture to Transformer, based on the state space equa- tions 24 65). Typical SSMs 25 26 66
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 parame- ters 70 and designs 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 222] We omit original S4
67
[10 model because it focuses Diagonal Plus Low-Rank (DPLR) transition matrix
16

---

A rather than fully diagonal one, which is difficult to implement due to computational inefficiency. For most SSMs, d, H. which means each channel uses independent parameters (q a. k
DSS [26]: DSS is the first research to show the effectiveness of diagonal structured SSM with fixed parameters, which means all parameters
constant through time. For the general form of DSS, d = d H, de N, where N is an expansion. Q. K, A are independent of x2 and determined by learnable parameters [A, B, C, A), using ZOH method to discretize. Furthermore, X2 Vt. g: and Wo- where I denotes the identity matrix.
S4D is also a diagonal structured SSM with fixed parameters. 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//a compared with DSS. Similarly, X Vx, gi and Wo I.
S4D 25
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.
SS [2 S5 uses MIMO (Multi-Input, Multi-Output) SSM with fixed parameters. For its general form, d, Nd = H, + = 1, where N is an expansion much smaller than that of SISO models. a = exp (AA) is the fixed decay and (x,, f₂) X; B where B = (AA) - (exp (AA) - ()) AB. i... it chooses ZOH method. Similar to previous SSM models, all the parameters are independent of
X1 and are determined by learnable parameters [A, B, C, A]. Furthermore,
K, g: = 1 and
Wo = C.
Mamba [6: : Mamba advances SSMs with selection, which means all parameters are time-varying and deptddeeend.nt. For the general form of Mamba, d,- = d= H, de H = N. q. = C. = x; W c. Similarly, k,- = B, x.W nÂz, The decay term can be written as 0t, exp(A.A). X Vt. g: and Wo = I. Furthermore, A. softplus(x, W Woora + ba).
Table A2: From general form to SSMs.
DSS S4D 25). 21 (ZOH)
S4D 23 H3
Models
Mambal 16
25
(Bilinear)
F (xx; 0g)
C
x, We
f (x,, 0%)
( (A I)AB
(I
AA
A)
Ax n)
f, 0.))
X
x,
(Xe, 0a)
exp (AA)
(I 2
A A -
A)
exp (A,A)
exp [AA)
2
(xe, 0g)
1
1
d d = H
d, Nd = H ds/ =1
Dimension
/1 = N
A1.3 LinRNN
LinRNN is 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 755. Recent works have made efforts to explore more effective recurrence
[77 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 dy = d, H d except LRU 19), which means each channel uses independent parameters (qf at k They maintain smaller unexpanded hidden state, which is a ID vector and can be seen special case with head dimension size d' = 1.
* RWKV-4 [7]: RWKV-4 is 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. exp W x) where Oz W k. V, = = Wy while q 1. a₂ = exp (-W) is fixed decay with the learnable parameter 0. = W. Furthermore, y: fo, sigmoid( W, W) W where 0, W. r.
17

---

-GLUU 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, ar/k, = fask (x,, fa/k)) oxx; W a/K + ba/k) where latk Wa/k /a/k] and o is the sigmoid function. Ve (x Wy v bv) while qu 1. Furthermore, Yt = lo @ SiLU(x; W, by) where
0g [W,, g: bgl-
. Griffin [8] Griffin is an RNN model with gated linear recurrence. Moreover, it ties dynamic
decay and Key as do. For the general form of Griffin, a fa (x,, A (x, W. +b.) where 8a = [A, W a. ba]. Then K / a to replace Key's role. Vi o(x; W
b C x
which contains an input gate additionally. q: 1 and y: = [o, @ GeLU W, + bg) Wo where
Bg W g. 6gl.
LRU 19 LRU is a linear and diagonal RNN model with fixed decay, inspired by deep SSMs' success. For its general form, d,- = .d H. de = 1, where N is a small expansion. at = A is the fixed decay and Vi (x, 8.) x, B where A, = B. Furthermore, qi k, gr 1 and
Wo
Table A3: From general form to LinRNN.
Models
RWKV-4
GLRU
Griffin 18
LRU[19]
(x,, 0g)
1
1
1
(x, 0%)
exp (x, W x) W exp(-W W,)
(x, W, + bb))
of
2
1
f, [x,, 0,)
o(x, W. + b,)
(x, W + b,) OX
B
f (Xx, Oa) (x,,, 0g)
(x; W / + +by)
A A tbb
A
o(x, W + bbg)
ox W by)
1
dy Nd H
Dimension
d₁ d, d d= H
d/ H =
A1.4 Approximation to Softmax Attention
Approximation to softmax attention includes two parts: value and functional approximation. Some
previous works of LinFormer [1 30, 53 78 were devoted to designing better kernel map ø to approximate the value of Map via randomized features, i.e., Φ( )6) (k,) z exp(q,k]) 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, 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 A,, Therefore, their purpose of achieving value approximation cannot be achieved 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. CO 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 CO, so we analyze C1, C2 and C3 in appendix A2. appendix A2.2 and appendix A2.3 respectively. Then we synthesize all the discussions and present the conclusions in appendix A2.4
18

---

For notation we use Q. V, A, to denote Query. Key, Value and Decay matrices respectively. Considering decay can be either dynamic or fixed, here we use subscript to distinguish. i.c., A/A, denote fixed/dynamic decay. Corresponingg q,, k,, olor, 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 CI (Dynamic memory ability). Let S. &ddx be hidden state of general linear attention (see Sec. 3). At a time t, the information about Via.....Vta, is successfully stored in St (tt...., .td2 S t), which means (Si)i: = Vt,, When the new important input N441 arrives, only models with dynamic decay can substitute historical unimportant Vta by Vi+1 successfully, ie., obtain (S4+1),,: Ve+1-
Proof. Without loss of generality, suppose t = dg and t; i(i I.....d.). de). where dę is Key dimension. That is, the information about Vi...., Vd, is successfully stored, which means the i-th row of Sa, (Sa, )ē, v,. Suppose Vi is unimportant information need to be substituted. Now the new and important input vd,-I arrives, models with general form Eq. (13 update Sd, as follows (heads are omitted for simplicity):
S = diag(ad++))Sa. + (kda+1) V
Vds+1-
(A1)
We will check whether the model obtain (Sa++1)1; +1,+1, with fixed/dynamic decay without decay. In the following discussion we let kd, = 1.
i) Models with fixed decay, which means Od,+11 = E (0, 1). The model can update and obtain (Sac+1)1.= = av1 Vd,+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 ad = 1. This case results in (Sd++1)1, = =1 + Vd 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 distracted by all the information before, making it hard to focus on and approximate most important ones. After normalization, the attention distribution will have relative high entropy along time dimension, which is often referred as attention dilution problem.
iii) Models with dynamic decay, which means ad,+1 E [0, 1J. The model can easily erase and substitute information. Setting 0dd++1 [0, 1,.... 1] can obtain (S1,+1)1, Vd 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 = [x( 0g), kr = (x,, 0g) and oz = fa (x,, U.). Assuming the inputs are good enough and the functions (f.g fx, fa) expressive enough, we can shift from solving (0,, 0g. 0&) to solving (Q, A,). Based on definition 4.1 we focus on approximating the attention scores between stored tokens Xtys Xta.
Proposition A222. Only models with parameters (Q, Ar), (Q, or (Q, Ac) can satisfy C2 (Static approximation ability). For an arbitrurily given sofimax attention map P with scores Pts. only above models ensure Eq. (A2) and Eq. (33) hold with bounded parameters.
xxx,.x/Q,KKK = qi (o II 0 = -Pes: vs S tt= t1, jmst1
tde
(A2)
b. ||ril| S Ca.lIkal| 5 Ck. ||ae|l E 10, 11, 4,, ldx,
(A3)
where Cg. Ck are constant and || -11 denotes vector norm. In our general form in Sec. f(x,, X, = LinAttMap (Q, K),, (see Eq. (88 For bounded inputs X, bounded
parameters (0,, Ox 0a) equivalent to bounded (Q, K. A.).
19

---

Proof. Without loss of generality, suppose t = dg and t; = i(i 1, ,.....dy,, where dy is Key dimension. That is, we need to prove following equations hold (heads are omitted for simplicity):
x, A,) - q: II
a,) o k.) k.)² = Piss Vs
dRr
(A4)
5=
s.t. q - 1 & Ck, ||aull E [0, 1]. vt = dR.
(A5)
Simple case. We first simplify the problem via i) setting dimension size of Q i.c. dN = and ii) considering only a given time ł (1 t de) and its autoregressive attention distribution
Pr Pts, 8 = ......t E Rixt, We need to prove the following equations hold with bounded parameters, foundation conclusion:
x1x,xx/Q.KKK. Q, = 90 II az) jms+1
Pts. Vs I....,t
(A6)
8. L.
1g.1 S C.. Ikis S Ck2 as E [0, 11, I....,L.
(A7)
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 's role is discussed later.
i) Models with fixed decay a result in unbounded parameters. We can solve Eq. (A6) and derive:
Pis
Pts
[kg
as
(A8)
a
q
C
a
Because a E (0, 1), the last term -- is unbounded, which leads to unbounded k2 Hence Eq. (A7) can not be satisfied. The result intuitive because the fixed exponential decay makes it hard for token to attend to distant information, leading to an excessively focused attention map rather than an arbitrary attention map. This conclusion about fixed decay general and is unrelated to the usage of Q or K.
ii) Models without decay (a 1) can satisfy Eq.
A66. We can solve Eq. A6 and derive:
kg Pte
vs 1,...,t.
(A9)
qt
At any time fix appropriate |g.| Cg. and can obtain bounded solutions for all hig
iii) Models with both dynamic decay ay and key ky can satisfy Eq. (A6), but have parameter- redundancy. We can solve Eq. (A6) starting from time t to time 1:
k 9t
(A10)
10 = Pci-
(A11)
(A12)
P
k k (41-242-...a))
(A13)
qt
(A14)
k102 (a3. ... ,.)))) = Pet
(A15)
At any time / - k. parameters QG-k42, Oit have already been determined. is equivalent to two free parameters ki-k, 0t-k++1 to concurrently approximate just variable P which is redundant. Dynamic decay and key actually have similar function, i.c., to balance historical memory and new input. Fix a lqel & Cy, we can obtain many possible bounded solutions of all K2 and a.,
9.
iv) Models with dynamic decay a and without key kę can satisfy Eq. (A6). Inspired by the parameter- redundancy in iii), we further consider a linear model with only query and dynamic decay, i.e., using dynamic decay 1 to replace h; 's role to approximate softmax attention map. We further assume qe 1, which is bounded, and define S, E' -= Pts. which satisfies S. [0, 1] and S, 2 S,-1-
20

---

Solve Eq. A6 starting from time t to time 1:
at Pu = St-1;
(A16)
0t-1 1 Pti-1 St-2
(A17)
a4
St-1
(A18)
Pti-k
St St-k-1
0t-k
(A19)
0/1 ... ..0gg
S
(A20)
Pti
So
(A21)
02 S1
At any time t k. derive closed-form and bounded solution Oę-k Sc-k-1 8
[0. 1].
General case. Aiming to solve Eq. (A4) and Eq. (A5), we generalize vector version with de 1 and consider distribution of all time Pt, = 1,...,d..
dk- In an extreme case, let one dimension of
q relate to time, which means qe -0....4h
0). Through the analysis above for one
dimension (Eq. (66) and Eq. (A7 ) we can ensure:
q (III))) = ge( II 0z) = 1=8+1
= Prss Vs I....,L.
(A22)
1=
Actually, we can further relax the condition and view qu as a selector, which selects several channels of Hadamard product of a and k., and uses these channels to approximate its attention distribution. One-hot assumption of qi 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 's attention distribution theoretically.
In summary, through analysis of simple case i) to iv), models with fixed decay A cannot ensure bounded parameters. Through analysis of general case, Q is necessary. So we can conclude that: only models with parameters (Q, K, At), (Q, K) (Q, A,) can satisfy C2. (For models with only Q. without decay and k, for all I, 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ı, d, are fixed, parameter counts of 0,, O2. 0, 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/A,)) all together, ii) exploiting (Q, K), (Q,)/A,) or (K, A/A)), iii employing only one of Q K. A/A,.
Proposition A2.3. Only models with Query and dynamic decay (Q, can satisfy C3 (Least parameter approximation). Models with (Q.Ai) can meet CO, Cl, C2 simultaneously with fewest parameters.
Proof. i) CO: all the linear models with general form (Sec. own linear complexity and satisfy CO. úi) Cl: according to proposition A2.1 only models with dynamic decay A, have dynamic memory ability.
iii) C2: according to proposition A2.2 only models with parameters (Q, (Q, or (Q, have static approximation ability.
So we claim that only models using (Q, K, A,) or (Q, A.) can meet CO, CI and C2 simultaneously. And (Q, Ax) 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 performed.
21

---

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 Mip, but not whether these functions will be successfully learned.
i) Linear approximation. The necessary conditions (CI and C2) for LinAttMap to achieve approxi- mation to MoftAttMap is that its implementation must include at least two parameter groups: Query Q and dynamic decay A,. Both (Q, K. A1) and (Q, At) can achieve approximation.
ii) Least parumeter appreximation. Lin utilizing (Q, A,) can achieve linear approximation to MoftAttMap with theoretically fewest parameters.
iii) Function of Query. Q can be Hadamard product of a and k, 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.
channel selector, which selects several channels of
iv) Function of dynamic decay. Dynamic decay A₁ 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), 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.1T explicates enhancement i) and design of basic MetaLA layer. appendix A3.2 introduces enhancements ii) and iii), and design overall MetaLA architecture.
A3.1 MetaLA Layer
In MetaLA layer, we exploit - a to replace k; based on theoretical analysis in appendix A2 ie.. LinAttMap utilizing (Q, A₂) is the linear approximation with least parameter redundancy to SoftAttMap and K is redundant.
General Recurrent Form. Using marks in Eq.
(11 -q..E..
15
the improvement lies in Eq.
A25
q: x,W Q. a: oxx E R lxde
(A23)
V. = x,W v. 6(x bc) ERixd,
(A24)
S/ diag(a,) ) + (1 - a E ERoxd
(A25)
0, XNorm concat XNorm(concat / S.S. ... (#! E Ridd.
(A26)
yi = (gr O E Rid
(A27)
where X RIxd denotes the current input. Wo, W. E Rdxds Wv, Wo Raxd, and
and h = 1,... H the index of heads. Here de represents Query/Decay dimension. LayerNorm is chosen as the normalization operation. Furthermore, we choose a = Sigmoid()) following [1], where - = 16 is used to control the value of dynamic decay and و SILU.
Wo E Rd. xd are learnable parameters. dzyo du/s "
As for parameter allocation, we simply set d, = d and d₃ $. similar to GLA [5]. Thus, the whole number of parameters used by a MetaLA layer is 4d the same as the softmax attention layer and smaller than other popular attention-like subquadratic models such as RetNet [14] (8d²) and GL.A [15] (4d² + 24d), 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 GLAA such that do not sacrifice expression capacity of fa and approximation performance of MetaLA.
22

---

General Parallel Form. The recurrent form of MetaLA can be written in parallel mode as follows:
0 ((((0A)) () о MMV,
(A28)
[1,85
(Q/K/V),:
A II O3 j=1
B Os, M., = (
(A29)
B in Eq. (228) denotes element-wise division and (Q),, means the 1-th row of matrix Q. In practical implementation, we utilize chunkwise form proposed in [15] for hardware-efficient training.
A
Attention map. Here the attention map is ((Q @ () ") ④ M). The element in the attention
map matrix is as follows (heads are omitted for simplicity):
(q. ((T)) 0z) o (1 (1-as)) St..
LinAttMapr,s
(A30)
0, s >
A3.2 MetaLA Transformer
Improved MetaLA. After deriving the basic MetaLA layer shown above, we consider further optimization when designing complete Metal. block. One observation is that in most cases, setting 0g = and completely discarding historical information is catastrophic. Actually, in real-world scenario, the learned dynamic decay 0g always closes because of many tasks' strong need to capture long-term dependencies. Hence we propose a technique called self-augmentation to enlarge token's attention itself while do not disrupt the model's state constitution, 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. (A3I) and Eq. (A35):
X1 Convid(x;)
E Rixd
(A31)
q: x, W Q: a = (xx
Rixde
(A32)
V, = x, W v,g. o(x W E Rixd,
(A33)
S diag(a (
(1 11-af a ERdixd.
(A34)
A S + aug (q: (w @ (1 - (1-a))) v.) E Rixd
(A35)
O, XNorm(concat (of; of II) RIxd,
(A36)
y: (gr の o. Wo R
(A37)
As for self-augmentation, without changing the composition of hidden state S is only added on the output process through a learnable parameter Wang R1x, which is then divided into heads like other parameters do. mgg (-) is 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 (Pt., $ < f), because we only change of while maintaining St unchanged like that in original Metal.A layer Eq. (A25). Thus the separation of output and memory is realized.
= ((QQ O A)
B 5
④ M
V + diag (sum ((Q c Wi, = 1)
(A38)
where Wt, = Waug and sum(Q Q dim marks are defined same as Eq. (A29).
1) means calculate sum for each row of matrix Q. Other
MetaLA block. Regard each layer as function mapping the input X Raxd to the output E RD xd where:
Xi
y
X2
Y2
X
Y
(A39)
Xn
y
23

---

following Transformer structure, now the Token Mix mechanism Eq.
MetaLA Transformer
A31 Eq. A37 can be integrated as
follows:
N
Y Mix X),
(A40)
Add
and the Channel Mix part (GLU B can be written follows:
GLU
Channel Mixer
Y (Swish(XW,, o XW2)W, (A41)
XNorm
which can be integrated as:
Add
X x TokenMix (XNorm' (x')) + X', (A42)
Token Mixer
Layer
X -ChannelMix (XNorm + +x( (x**+ )), (A43)
XNorm
where X and X E R refer to the input and output of block in MetaLA. We choose XNorm LaycrNorm. Stacking several MetaLA blocks above, can derive complete MetaLA Transformer as Lin- ear Foundation Model, as shown in Fig. A1
Figure A1: MetaLA Transformer. Stack- ing N MetaLA blocks, each block is com- posed of two modules in sequence: token mixer and channel mixer.
A4 Experimental Details
Table A4: prper-parameters of MetaLA on LRA. d is the dimension of model. dl is the dimension of d and dæ, d2 is the hidden dimension in GLU. num-warmup and max-step used for cosine warmup.
Task
Depth
d1
d2
dropout
Ir
bs
wd
num-warmup max-step
Listops
6
32
16
64
0.0
0.0005
128
0.01
5000
50000
Text
128
256
0.1
0.004
64
0.0
10000
50000
Retrieval
2
256
128
512
0.1
0.0008
128
0.0001
312
50000
Image
512
256
512
0.0
0.003
128
0.0
30000
50000
Pathfinder
6
128
64
128
0.0
0.002
256
0.0
50000
500000
Path-X
32
0.0
0.00075
256
0.0
15000
500000
Associative Recall (AR). Following [6). we train two layer models with Transformer backbone that interleaves token mixer and channel mixer. Learning rates are swept by np. logspace(- (-aaacc((44, 22 for sequence length 256 and additional np. logspace(-5) -3, 4) for length 512, and maximum test
accuracy reported. For GL.A [15) and MetalA, we set H and de d, = The kernel size of short convolution of MetaLA is 2.
Language Modeling. For 360M/1.4B model, we train it from scrach with total of 15B/300B tokens on 16/32 A100 GPUs at learning rate of 3e-4/2e-4 with batch size 0.5M/2M. Both models maintain length of 2048 and are trained using fp16. The training setup of baselines [5] 16 of 360M MetaLA are aligned with MetaLA configurations. For the 1.3B MetaLA, we compare
with publicly available models 14 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 GPT- Neox [45]. evaluate our models on SuperGLUE benchmark (38] and Common Sense Reasoning benchmark including LAMBADA [80), LogiQA [81) Winograd Schema Challenge (WSC273) [82]
[5]
16
20
44.
To maintain a fair comparison between linear
BoolQ 883, PiQA 84,, HellaSwag 85). WinoGrande [86), ARC-easy (ARC-e), ARC-challenge
(ARC-c) 187. 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 report F1 score 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.
other subtasks. For SuperGLUE benchmark,
Long Range Arena. We evaluate the long sequence modeling capability of our model on LRA task. We use the adam W optimizer and cosine warmup scheduler. We set the head size be 4 for group

---

normalization. In Retrieval, Image, Pathfinder and Path-X we use a bid-model. The hyperparameters for all tasks can be found in T4
Image Classification. We evaluate our models on ImageNet 40 The input size of ImageNet is 224 224. Following Deit [49], the batch size is 2048 during 300 training epochs with cosine- decay learning rate whose peak value is 2.4 x 10-3, We choose AdamW W (ß) = 0.9, 32 = 0.98) with 0.05 weight decay as the optimizer. Note that do use cutmix or mixup during the training.
Table A5: Performance Comparison on Additional subtasks for Common-Sense Reasoning. PS: parameter size (billion). tokens (billion). means the results reported by [20]. i indicates testing using open-source checkpoints. For baselines that need to be compared, if they do not have public checkpoints, we train and test them under identical conditions with MetaLA, LAMB: Lambada. HS: HellaSwag. WG: WinoGrande. OBQA: OpenbookQA. MetaLA,,: MetaLA with tied embedding trained using 100B tokens. Metal.Ag: Metal. trained with 300B tokens. "AVG" refers to the average
result on subtasks other than LAMBADA.
Models
PS
LAMB ppl
LAMB
LOGIQA
WSC273
Pythia
1.4
10.94
49.78
21.35
72.89
HGRN!
1
100
21.81
36.39
22.43
58.97
Mamba
1.4
100
14.02
44.44
22.73
68.50
RetNet
1.3
100
22.65
36.79
22.73
63.74
GLA
1.3
100
20.05
39.92
21.81
63.00
MetaLA,
12.84
45.99
21.81
65.93
MetaLA,
1.4
300
10.06
50.42
21.35
73.63
Models
PS
T
BOOLQ
PIQA
HS
WG
ARC-e
ARC-c
OBQA
AVG
Pythia
1.4
300
63.12
70.89
51.98
56.99
60.56
28.41
33.20
51.04
HGRNI
1
100
58.75
71.00
48.05
51.14
55.51
28.07
31.80
47.30
Mamba
1.4
100
53.27
71.44
48.63
53.59
58.59
29.01
31.80
48.62
RetNet
1.3
100
60.21
69.53
48.39
53.28
54.17
26.19
30.80
47.67
GLA
1.3
100
61.04
70.08
48.00
51.93
54.88
28.33
31.40
47.83
MetaLA,
1.3
100
55.50
70.02
47.32
55.01
56.90
27.47
33.00
48.11
MetaLA,
1.4
300
56.27
72.25
53.58
58.17
61.49
30.03
34.60
$1.26
Table A6: Results on MAD tasks. All architectures are tested according to the MAD protocol.
Models
Compression Fuzzy recall In context recall Memorization Noisy recall Selective Copy Avg
Multihead Hyena
47.79
18.01
97.46
89.48
98.74
90.81
73.72
GL.A
37.70
12.45
91.58
57.05
92.58
88.63
63.33
Mamba
43.95
9.60
87.92
89.45
90.91
81.79
67.27
MetaLA
45.55
15.18
99.87
85.83
99.73
97.71
73.98
A5
Additional Experiments
Additional subtasks for Common-Sense Reasoning. We extend more subtasks of Commonsense Reasoning for IB4 MetaLA. Additional experimental results are shown in Tab, A5
More challenging settings for MQAR. We
Table A7: Results on MQAR with sequence
evaluate some models with sequence length 512 and with more retrieval key-value pairs (80, de- fault is 64). The attention baseline benefits from global modeling capabilities, achieving optimal results. The results in Tab. show that Met- aLA outperforms Mamba (does not converge under the same training conditions), and there is still significant gap compared transformers.
length 512 and retrieval key-value pairs 80.
Models
Model dimension
Acc
Transformer
64
>99.0
Mamba
64
0.0
MetaLA
64
28.5
Transformer
128
>99.0
Mamba
128
0.0
MetaLA
128
90.4
Results on the synthetic MAD task. we eval-
uate several models MAD [4], collection of six synthetic tasks predictive of scaling laws,
25

---

Training throughput
GPU memory usage
1 Mamba
Tokens per Second IKUS
60-
Gigabyye (GB)
30
Transformer+
Transformer++
Mamba
20
GLA
GLA
MetaLA
MetaLA
Training length/Betcc
Training lengttyBatch size
Figure A2: Training efficiency evaluations. The throughput and memory usage single A800 GPU of Transformer and various linear models. Transformer++ is implemented using FlashAttention and SwiGLU.
90
including recall, memorization, and compression. As shown in Tab. A6 MetaLA achieves the best results across various linear complexity models.
Results the Needle in Haystack (NIAH) task. We also present experimental results the NIAH task following 42 which is designed to evaluate the in-context retrieval capabilities of LLMs. Retrieval ability long texts is significant challenge for linear models, as all current linear models lack good solutions to this problem. Nonetheless, 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 model ranging in size from 380M to 3B, trained with up 300B tokens, 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 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.
Table A8: Results on the Needle Haystack (NIAH) task. We introduce accuracy metrics across four context scales and three model scales. The middle columns display accuracies below the 4K and 8K thresholds. The rightmost columns detail both the average accuracy and the weighted average accuracy. All models are trained with sequence length 8K.
Models
PS
Acc@2K 2K
Acc@4K
Acc@8k 8K
Ace@ 16K
Acce=4K Acce=8K
Avg
Weighted Avg
LLaMA2
0.4
100.0
97.1
97.8
0.0
99.3
99.5
56.4
52.3
HGRN2
0.4
8.6
6.3
1.3
0.0
17.0
9.3
4.9
4.8
MetaLA
0.4
25.7
2.9
11.1
4.4
11.3
8.7
8.5
9.0
LLaMA2
1.0
100.0
71.4
73.3
0.0
92.5
90.9
47.8
44.1
HGRN2
1.0
17.1
5.7
2.9
3.5
18.3
13.4
9.7
10.0
MetalA
1.0
7.9
8.3
13.7
17.8
14.3
14.9
12.6
LLaMA2
3.0
97.1
100.0
82.9
0.6
93.9
48.8
HGRN2
3.0
58.4
11.4
2.9
7.3
28.9
18.0
17.9
MetaLA
3.0
48.3
7.0
4.1
18.4
34.8
22.2
19.1
19.5
26

---

Table A9: Scalability tests of MetaLA on the CommonSense Reasoning benchmark. PS: parame- ter size (billion). T: tokens (billion). "AVG" refers to the average result of all subtasks.
Models
PS
T
BOOLQ
PIQA
HS
WG
ARC-E
ARC-C
OBQA
Avg
LLaMA2
0.41
300
54.04
67.19
38.75
52.17
49.24
23.72
30.00
45.02
Cosformer2
0.38
300
57.40
66.27
36.65
50.59
51.81
23.72
29.00
45,06
MetaLA
0.38
300
60.09
67.79
38.51
50.99
52.19
25.60
30.00
46.45
LLaMA2
300
56.42
69.97
47.04
52.72
57.07
28.16
32.60
49.14
Cosformer2
300
44.28
70.73
45.55
50.51
55.22
27.30
31.00
46.37
MetaLA
300
59.05
69.37
46.43
54.38
57.41
26.96
33.00
49.52
ILaMA2
300
61.31
73.18
57.88
59.59
63.93
31.40
34.00
54.47
Cosformer2
300
50.92
74.27
57.38
57.30
63.22
31.40
35.20
52.81
MetaLA
300
62.84
74.16
59.25
58.80
64.52
33.28
35.80
55.52
27

---

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: 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 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 long 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 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 (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn 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 few datasets or with 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, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails 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 faimesss.
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 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant 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 complete (and correct) proof?
28

---

Answer: [Yes]
Justification: In Sec. 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 proof' in the paper should be numbered and cross referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide 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
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
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 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 novel architecture, describing the architecture fully might suffice, or if the contribution 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 accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily 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 way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which 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 have some path reproducing or verifying the results.
5. Open access to data and code
29

---

reproducible, they Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: We have not open-sourced the code, we have clearly outlined the details of our experiments in Sec. 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.cm cc/
public/guides/ htaaeees//ooeeSSemiiiiiinooieee
for more details.
While we encourage the release of code and data, we understand that this might 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 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
/
Coaee//ddeeeSbmmmmma ssi onPolicy) for more details.
The authors should provide instructions 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. only subset of experiments 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) 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, hyper- parameters, 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 ap- pendix A4
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to 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, 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, confi- dence intervals, or statistical significance tests, 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 with given experimental conditions).
30

---

The method for calculating the error bars should be explained (closed form formula, call 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.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma enror bar than state that they have 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 range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: Sec. S 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 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 the paper conform, in every respect, with the NeurlPS Code of Ethics httts:://neur: wwtpeeerrips..cppuueeeeeeeeeeeee http:/pp...c.pp.ccttttccccccceea
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 NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration 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 no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
31

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess 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 partieular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate 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 intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have 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 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
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?
Question:
Answer: [Yes]
Justification:
The data, code, and models used 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 dataset.
The authors should state which version of the asset used and, possible, include 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.
32

---

If assets are released, the license, copyright information, and terms of use in the thcode. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssi
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 not available online, the authors are encouraged to reach out to the asset's 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/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent 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. Crowdsoureing 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 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 the supplemental material is fine, but if the main contribu- tion 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 equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: the paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending 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.
33

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere 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.
34


Reviews:
[{"review_id": "70bk1Cae4d", "reviewer": "Reviewer_rbUC", "summary": "This paper discusses the development and evaluation of linear complexity models for Transformers, which aim to replace the conventional softmax attention mechanism. The authors first unify existing models including LinFormer, SSM, and LinRNN into the framework of Linear Attention. Then they establish three (actually four) criteria for optimal linear attention: (0) linear complexity, (1) dynamic memory capability, (2) static approximation ability, and (3) minimal parameter usage. This paper shows that the three existing models mentioned above do not meet all these conditions and therefore perform suboptimally. In response, they introduce a new model, Meta Linear Attention (MetaLA), designed to fulfill these criteria. MetaLA is tested across various tasks including the Multi-Query Associative Recall (MQAR) task, language modeling, image classification, and the Long-Range Arena (LRA) benchmark, where it outperforms existing linear models, demonstrating its effectiveness.", "questions": "The 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?", "limitations": "Limitations are adequately discussed in the Section of Conclusion and Discussion.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1, 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.", "weaknesses": "(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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your insightful feedback and your time in reading our paper.\n>*Weakness 1:*  Concerns about the optimality for linear attention.\n\n**A:** Your concern is valid. There is no definitive evidence to prove that MetaLA is the optimal approximation of self-attention. Please allow me to explain why we chose to use this term. \n\n- First, we repeatedly state in the paper that we are discussing the necessary conditions for optimality. In fact, these necessary conditions highlight an important issue, namely, what kind of models can approximate softmax attention.\n\n- Second, we use the term \"optimal\" to draw the community's attention to the issue of how to approximate softmax attention. We believe that better approximating softmax attention is one of the potential ways to address cutting-edge issues in linear models, such as retrieval and in-context learning.\n\n>*Weakness 2:* The presentation in Section 2.\n\n**A:** Thank you for your suggestion. We will carefully check and optimize these expressions.\n\n>*Q1:* Explanation of the mask matrix $M$.\n\n**A:** Your understanding is correct. For the matrix $M$, the lower triangular part (including the diagonal) is set to 1 to maintain model causality, ensuring that the output at time $t$ depends only on inputs before $t$. As for the upper triangular part (excluding the diagonal),\n\n- Linear attention: The upper triangular elements are set to 0, as no exponential operation is involved.\n\n- Softmax attention: The upper triangular elements are set to $-\\infty$, ensuring that their exponentiated values become 0, thereby not affecting the normalization of values before $t$.\n\n>*Q2:* About the complexity of linear attention and the chunk-wise algorithm.\n\n**A:** As shown in Line 78, **during inference**, the time and memory complexity per token is $\\mathcal{O}(1)$, meaning that computational and memory costs remain constant. From the parallel formula Eq.5 in Line74, we can express the inference in a recurrent form (operating separately on the numerator and denominator):\n\n$\nS_{t} = \\sum_{i=1}^{t}\\phi^{T}(k_{i})v_{i} = S_{t-1} + \\phi^{T}(k_{t})v_{t}, \n$\n\n$\nn_{t} = \\sum_{i=1}^{t}\\phi^{T}(k_{i}) = n_{t-1} + \\phi(k_{t}), \n$\n\n$\no_{t} = \\frac{\\phi(q_{t})S_{t}}{\\phi(q_{t})n^{T}_{t}}.\n$\n\nIn this iterative form, each operation only needs to maintain a fixed state size $S_{t}\\in\\mathcal{R}^{d \\times d}, n_{t}\\in\\mathcal{R}^d$, resulting in a storage complexity of $\\mathcal{O}(1)$ (occupying $d^{2}+d$ space, where $d$ is a constant). Furthermore, each token computation involves $d^{2}+d$ multiplications and additions, ensuring the time complexity is also $\\mathcal{O}(1)$.\n\n**About the chunk-wise algorithm.** The Chunk-Wise algorithm combines the low complexity of recurrent operations with the high throughput of parallel processing during training. It achieves this by using serial processing between chunks and parallel processing within chunks, ensuring (1) full utilization of GPU computational units within a chunk and (2) reduced redundant calculations of historical information between chunks.\n\nFor Linear Attention, normalization by the denominator is typically omitted. We modify the earlier formulas for further discussion (a more general form is provided in Section 4.2 of [1]):\n\n$\nS_{t} = S_{t-1} + k_{t}^{T}v_{t}, \n$\n\n$\no_{t} = q_{t}S_{t},\n$\n\nDuring training, if using a parallel form of computation, i.e., $O = (QK^{T}\\odot M)V$, these matrices must be computed from left to right, resulting in a computational complexity of $\\mathcal{O}(L^{2})$. To reduce the time complexity of the linear model during training, researchers use a chunk-wise algorithm, which consists of two parts: intra-chunk and inter-chunk. We divide a sequence of length $L$ into $\\frac{L}{C}$ chunks (with padding if not divisible), each chunk of size $C$. The computation method is as follows(The first term on the right side of the equation is the inter-chunk, and the second term is the intra-chunk.):\n\n$\nO_{[i]} = Q_{[i]}S_{[i-1]} + (Q_{[i]}K_{[i]}^{T}\\odot M)V_{[i]}, \n$\n\n$\nS_{[i]} = S_{[i-1]} + K_{[i]}^{T}V_{[i]}.\n$\n\nThus, each chunk will only generate a parallel-like computation, where the resulting $C \\times C$ matrix is a submatrix of the attention map. So the computational cost is $\\frac{L}{C} [(Cd^{2}+C^{2}d)+Cd^{2}]$. If $d$ and $C$ are constants, the time complexity of the computation becomes $\\mathcal{O}(L)$.\n\n>*Q3:* Why self-attention requires infinite states.\n\n**A:** The computation for the $t$-th output of self-attention during inference is given by:\n\n$\no_{t} = \\text{softmax}(q_{t}K_{t}^{T})V_{t} = \\frac{\\sum_{i=1}^{t}\\langle q_{t}, k_{i}\\rangle v_{i}}{\\sum_{j=1}^{t}\\langle q_{t},k_{j}\\rangle}\n$\n\nwhere $K_{t}^{T} = [k_{1}, \\cdots, k_{t}]$ and $V_{t}^{T} = [v_{1}, \\cdots, v_{t}]$.\n\nIn the Transformer model, a KV-Cache is maintained as a state, updated after processing each input to avoid redundant computations. Upon receiving the $ (t+1) $ -th input, the state is updated by concatenating $ k_t+1$ and $ v_t+1$ to $ K_t$ and $ V_t$, respectively, which increases memory requirements over time.\n\n>*Q4:* \"good enough\" in line 177.\n\n**A:** In the theoretical analysis, we solved for the $Q,K,\\Lambda$ matrices given any attention scores distribution. This implies a prerequisite assumption: the parameter functions $f_{q/k/\\alpha}$ can project the input $X$ into the corresponding matrix space. This prerequisite imposes constraints on $f_{q/k/\\alpha}$ (expressive enough) and $X$ (good enough).\n\n>*Q5:* The explanation of \"fewest parameters\"\n\n**A:** The minimum parameters refer to the smallest parameter group that achieves optimal approximation. If the model dimension $d, d_k, d_v$ are fixed, the minimum parameters equates to the smallest group. Please refer to line 185.\n\n---\n[1] Gated Linear Attention Transformers with Hardware-Efficient Training, In ICML2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rtK2ItWCn5", "reviewer": "Reviewer_2RWL", "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.\n\n1. They remove the K matrice redundant parameters and achieve better training efficiency\n2. Add self-augmentation. This allows the input value to immediately affect the output depending on the query to prevent forgetting current token information.\n3. Adding a convolutional layer at input X. This enhances the local interaction further, motivating Mamba and Griffin.", "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?\n2. Can you provide the code?", "limitations": "It feels minor improvement on top of previous linear attention alternative studies.", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "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.\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful feedback. Your questions are crucial and something we have been thinking about.\n\n>*Weakness1:* The experiment is not done with various scales. It is hard to know the effect of the scaling model and training dataset with MetaLA.\n\n**A:** Based on your suggestions, we conduct additional experiments. We further evaluate our model ranging in size from 380M to 3B, trained with 300B tokens, on the *CommonSense Reasoning* benchmark. Below are experimental results:\n\n|Models|PS|T| BOOLQ| PIQA |HS| WG| ARC-E| ARC-C|ОBQA|Avg|\n|-|-|-|-|-|-|-|-|-|-|-|\n|LLaMA2 [3]|0.41|300|54.04|67.19|38.75|52.17|49.24|23.72|30.00|45.02|\n|Cosformer2 [4]|0.38|300|57.40|66.27|36.65|50.59|51.81|23.72|29.00|45.06|\n|MetaLA (ours)|0.38|300|60.09|67.79|38.51|50.99|52.19|25.60|30.00|**46.45**|\n|-|-|-|-|-|-|-|-|-|-|-|\n|LLaMA2|1|300|56.42|69.97|47.04|52.72|57.07|28.16|32.60|49.14|\n|Cosformer2|1|300|44.28|70.73|45.55|50.51|55.22|27.30|31.00|46.37|\n|MetaLA (ours)|1|300|59.05|69.37|46.43|54.38|57.41|26.96|33.00|**49.52**|\n|-|-|-|-|-|-|-|-|-|-|-|\n|LLaMA2|3|300|61.31|73.18|57.88|59.59|63.93|31.40|34.00|54.47|\n|Cosformer2|3|300|50.92|74.27|57.38|57.30|63.22|31.40|35.20|52.81|\n|MetaLA (ours)|3|300|62.84|74.16|59.25|58.80|64.52|33.28|35.80|**55.52**|\n\nThe experimental results corroborate the scalability and performance of MetaLA.\n\n>*Weakness2:* 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.\n\n**A:** Thanks for your suggestion, we report the latency/throughput results of 1.3 B models in the general response to all reviewers, please refer to the PDF file. The report indicates that:\n- Our model exhibits good linearity, maintaining processing speed with increasing sequence length, in contrast to the Transformer, which shows a sharp drop in token processing speed as sequence length increases.\n- Our model is computationally as efficient as linear models like Gated Linear Attention [1], while being significantly faster than the Mamba model [2], which also has linear complexity.\n\n>*Question1:* 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? \n\n**A:** You raise a very profound and interesting question, one that we have been thinking about and we would love to discuss with you.\n\n**The AI (performance) perspective:** Although there is no strict theoretical proof of optimality for softmax attention, its excellent performance has been empirically validated by many large models, making it undoubtedly the foundation of existing large models.\n\n**The neuroscience (bio-plausibility) perspective:** It is almost certain that softmax attention is **NOT** the optimal solution for sequence modeling. To achieve AGI (Artificial General Intelligence), there are two constraints we cannot bypass.\n\n- **Power and Latency:** In terms of power and latency, the human brain operates very stable with bounded power consumption and trascient reaction. In contrast to Transformer models, if we assume the sequence input and output are interactions over time with the environment, then it faces a severe issue. Over time, its power consumption and latency will grow, a limitation brought by quadratic complexity.\n\n- **Physical Space Limitation:** As we pointed out in the paper, softmax attention is an infinite-state mechanism, leading to an increasing demand for physical space (for storage) as the model runs. This contradicts the human brain's ability to maintain a large and efficient memory capacity despite its finite size.\n\nIt can be seen that the perspective from which we view softmax attention determines whether it is optimal, and there are trade-offs involved. Our approximation strives to restore its performance optimality (which has been empirically validated), but for aspects where softmax is not optimal, we still possess unique advantages.\n\n>*Question2:* Can you provide the code?\n\n**A:** We provide the code link here: https://anonymous.4open.science/r/MetaLA-1BB3/README.md\n\n---\n[1] Gated Linear Attention Transformers with Hardware-Efficient Training, In ICML2024.\n\n[2] Mamba: Linear-time Sequence Modeling with Selective State Spaces, In Arxiv 2023.\n\n[3]  Llama 2: Open Foundation and Fine-tuned Chat Models, In Arxiv 2023.\n\n[4]  You Only Scan Once: Efficient Multi-dimension Sequential Modeling with LightNet, In Arxiv 2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "I64huVzqw6", "reviewer": "Reviewer_XAnv", "summary": "The paper presents a theoretical analysis of existing linear attention methods such as LinFormer, SSM, and LinRNN. Building on this analysis, the authors propose a unified framework that combines the strengths of these methods. Utilizing this framework, authors develop a novel linear attention model called MetaLA. The key innovations of MetaLA include the elimination of the Key matrix, the introduction of dynamic decay for achieving dynamic memory and static approximation, and the integration of self-augmentation with short convolution. \nThese enhancements aim to improve the approximation accuracy of softmax attention and outperform existing linear attention methods.", "questions": "- How does MetaLA perform on tasks with longer sequences compared to standard softmax attention and other linear attention mechanisms? For example, on NarrativeQA or other tasks from LongBench.\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", "limitations": "N/A", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- 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.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your insightful feedback and your time in reading our paper.\n \n >*Weakness1:* More long sequence validation besides the LRA benchmark. This is not sufficient to justify the scalability of the method for real LLMs.\n\n**A:** Thank you for the suggestion. In addition to the LRA benchmark, we address longer token sequences in LLM pre-training. Specifically, we pre-train a MetaLA-based language model with token counts ranging from 15B to 300B and a sequence length of **2048. To demonstrate MetaLA’s scalability, we also extended the 3B scale model experiment to 300B tokens on the CommonSense Benchmark as follows**, where MetaLA showed strong performance.\n\n|Models|PS|T| BOOLQ| PIQA |HS| WG| ARC-E| ARC-C|ОBQA|Avg|\n|-|-|-|-|-|-|-|-|-|-|-|\n|LLaMA2 [8]|0.41|300|54.04|67.19|38.75|52.17|49.24|23.72|30.00|45.02|\n|Cosformer2 [10]|0.38|300|57.40|66.27|36.65|50.59|51.81|23.72|29.00|45.06|\n|MetaLA (ours)|0.38|300|60.09|67.79|38.51|50.99|52.19|25.60|30.00|**46.45**|\n|-|-|-|-|-|-|-|-|-|-|-|\n|LLaMA2|1|300|56.42|69.97|47.04|52.72|57.07|28.16|32.60|49.14|\n|Cosformer2|1|300|44.28|70.73|45.55|50.51|55.22|27.30|31.00|46.37|\n|MetaLA (ours)|1|300|59.05|69.37|46.43|54.38|57.41|26.96|33.00|**49.52**|\n|-|-|-|-|-|-|-|-|-|-|-|\n|LLaMA2|3|300|61.31|73.18|57.88|59.59|63.93|31.40|34.00|54.47|\n|Cosformer2|3|300|50.92|74.27|57.38|57.30|63.22|31.40|35.20|52.81|\n|MetaLA (ours)|3|300|62.84|74.16|59.25|58.80|64.52|33.28|35.80|**55.52**|\n\n>*Weakness3:* About the low-precision training of FP16 and BF16.\n\n**A:** Our model is trained using fp16, and most linear models, such as [2,3,4], can be trained with half-precision, including bf16 and fp16. The numerical precision in our training shares the same characteristics as theirs. \n\n>*Weakness4:* Computational efficiency and memory requirements compared to softmax attention and other linear methods.\n\n**A:**  We report the throughput and memory footprint of 1.3B models **in the general response (please refer to the PDF)**. The findings indicate: (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.\n\n>*Weakness5:* Code Link.\n\n**A:** https://anonymous.4open.science/r/MetaLA-1BB3/README.md\n\n>*Weakness2 and Q1:* It is not clear how their model performs in the text generation scenario. 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\n**A:** We have added some additional experiments based on your suggestions. We validate MetaLA in **NarrativeQA**, a natural language processing task that challenges LLMs to comprehend and answer questions about long-form narratives. The results show that MetaLA performs well on this long-sequence tasks.\n\n|Models|PS|NarrativeQA(F1)|\n|-|-|-|\n|LLaMA2 [8] |3.0|20.8|\n|Mamba [5]|3.0|17.9|\n|MetaLA (ours)|3.0|18.2|\n\n>*Q2:* About the hyperparameters.\n\n**A:** Our pre-training hyperparameter settings were based on the experience summarized by predecessors [6]. The majority of the settings are identical to those used in Pythia. Our hyperparameter settings are provided in the code link as a YML file, based on the parameter passing format used in gpt-neox [7].\n\n>*Q3:* Relathionship between the model in paper \"Linear Log-Normal Attention with Unbiased Concentration\" and the general linear attention model in this work.\n\n**A:** This paper and our work are complementary. In lines 633-645, we discuss that the approximation of softmax attention can be divided into two parts: value approximation and functional approximation. Our work primarily focuses on functional approximation, while the paper 'Linear Log-Normal Attention with Unbiased Concentration' that you mentioned belongs to value approximation. \n\nThe paper 'Linear Log-Normal Attention with Unbiased Concentration' is a commendable work, offering extensive mathematical derivations and identifying the Log-Normal method as suitable for linear models to approximate the softmax attention map from the perspectives of entropy and spectral angles. This method has demonstrated strong performance in language models, particularly on the GLUE benchmark, and highlighted the linear complexity advantage of Linear Attention in long-sequence modeling. Our paper, on the other hand, unifies the token mixing mechanisms of Linear Attention, Linear RNN, and State Space Model, linking them to the softmax attention map from the viewpoint of attention maps. We analyze under what circumstances these models can functionally approximate softmax attention.\n\nThus, both papers offer complementary insights into the theoretical analysis of the softmax attention map from different perspectives, indicating great potential for their combination. We are genuinely inspired by this work and will cite the paper in our own.\n\n---\n[1] Scrolls: Standardized Comparison over Long Language Sequences, In EMNLP 2022.\n\n[2] Gated Linear Attention Transformers with Hardware-Efficient Training, In ICML2024.\n\n[3] RWKV: Reinventing RNNs for the Transformer Era, In EMNLP 2023.\n\n[4] TransNormerLLM: A Faster and Better Large Language Model with Improved TransNormer, In Arxiv 2023.\n\n[5] Mamba: Linear-time Sequence Modeling with Selective State Spaces, In COLM 2024.\n\n[6] Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling, In ICML 2023.\n\n[7] https://github.com/EleutherAI/gpt-neox\n\n[8] Llama 2: Open foundation and fine-tuned chat models, In Arxiv 2023.\n\n[9] HGRN2: Gated Linear RNNs with State Expansion, In COLM 2024.\n\n[10]  You Only Scan Once: Efficient Multi-dimension Sequential Modeling with LightNet, In Arxiv 2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "r9uVTlp4kL", "reviewer": "Reviewer_LW2x", "summary": "Proposes a unifying view of recent linear attention/SSM/linear RNN methods and compares these models to Softmax Attention. Proposes MetaLA to address the shortcomings of the prior methods in approximating Softmax attention. Performs experiments on MQAR, language modeling, image classification and LRA.", "questions": "My main questions and issues are listed above.\n\n- Line 287: typo, \"Value\" is misspelled in the text", "limitations": "Sufficient", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "- 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", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your insightful feedback. We will outline your suggestions and questions, followed by our detailed responses. We hope our answers address your concerns.\n\n>*Weakness1:* The general presentation of the \"optimal linear approximation\" and theoretical analysis is dense and confusing and obscures the contributions. I would recommend moving the Proposition statements...\n\n**A:** We sincerely apologize for any confusion in the organization of this section. We carefully considered and discussed your suggestions, and thank you for your kind advice. We find ourselves in a bit of a dilemma. On one hand, we strive to maintain the completeness and rigor of the paper by including the proposition statement. On the other hand, we have to carefully consider space limitations, which may increase the difficulty for readers. We hope you can understand the balance we are striving to achieve. While we have not yet identified a perfect solution, we will make every effort to enhance the clarity of our presentation while taking your concerns into account. We greatly appreciate your understanding and patience.\n\n> *Weakness2:* The MQAR analysis is weak. The results of more difficult versions of this task and an attention baseline will be conducive to understanding.\n\n**A:**\nBased on your suggestion, we evaluate the performance of several models on more challenging settings. The attention baseline benefits from global modeling capabilities, achieving optimal results under both conditions discussed in the paper (>99.0). The additional experiments show that MetaLA outperforms Mamba (does not converge under the same training conditions), and there is still a significant gap compared to transformers:\n|Models|Seq_len|Key-Value Pairs|Model dimension| Acc|\n|-|-|-|-|-|\n|Transformer [2] |512|80|64|>99.0|\n|Mamba [4] |512|80|64|0.0|\n|MetaLA (ours)|512|80|64|28.5|\n|Transformer|512|80|128|>99.0|\n|Mamba|512|80|128|0.0|\n|MetaLA (ours)|512|80|128|90.4|\n\nFor the results of longer sequences such as 1024, since it takes a long time for the linear model to converge, it is difficult for us to get the results during the rebuttal period. We sincerely ask for your understanding. As a supplement, we test another long sequence recall benchmark NIAH. Details are in the response to weakness3.\n\n>*Weakness3:* More downstream tasks that require the recall/retrieval abilities should be included.\n\n**A:**  According to your suggestions, we evaluate several models on **MAD** [1], a collection of six synthetic tasks predictive of scaling laws, including recall, memorization, and compression. As shown in the following table, MetaLA achieves the best results across various linear complexity models.\n|Models|Compression|Fuzzy recall|In-context recall|Memorization|Noisy recall|Selective Copy|Avg|\n|-|-|-|-|-|-|-|-|\n|Multihead Hyena [1]|47.79|18.01|97.46|89.48|98.74|90.81|73.72|\n|GLA [6]| 37.70|12.45|91.58|57.05|92.58|88.63|63.33|\n|Mamba [4]|43.95|9.60|87.92|89.45|90.91|81.79 |67.27|\n|MetaLA (ours)|45.55|15.18|99.87|85.83|99.73|97.71|**73.98**|\n\nThree distinct versions of the recall task and the selective copy task serve to underscore the effectiveness of our model in the domain of information retrieval. And we will cite the relevant articles.  We also present experimental results on the **Needle in a Haystack (NIAH)** task, which is designed to evaluate the in-context retrieval capabilities of LLMs:\n|Models|PS|Acc@2K|Acc@4K|Acc@8K|Acc@16K|Acc<=4K|Acc<=8K|Acc_Avg|Weighted_Acc_Avg|\n|-|-|-|-|-|-|-|-|-|-|\n|LLaMA2 [5]|0.4|100.0|97.1|97.8|0.0|99.3|99.5|56.4|52.3|\n|HGRN2 [2]|0.4|8.6|6.3|1.3|0.0|17.0|9.3|4.9|4.8|\n| MetaLA (ours)|0.4|25.7|2.9|11.1|4.4|11.3|8.7|8.5|9.0|\n|-|-|-|-|-|-|-|-|-|-|\n|LLaMA2|1.0|100.0|71.4|73.3|0.0|92.5|90.9|47.8|44.1|\n|HGRN2|1.0|17.1|5.7|2.9|3.5|18.3|13.4|9.7|10.0|\n|MetaLA (ours)|1.0|7.9|8.3|13.7|17.8|14.3|14.9|12.0|12.6|\n|-|-|-|-|-|-|-|-|-|-|\n|LLaMA2|3.0|97.1|100.0|82.9|0.6|95.4|93.9|48.8|45.1|\n|HGRN2|3.0|58.4|11.4|2.9|7.3|46.4|28.9|18.0|17.9|\n|MetaLA (ours)|3.0|48.3|7.0|4.1|18.4|34.8|22.2|19.1|19.5|\n\nRetrieval ability in long texts is a significant challenge for linear models, as all current linear models lack good solutions to this problem. Nonetheless, 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 will address next, following the unification of linear model forms.\n\n>*Weakness4:* Some of these questions should be better explored in this work .\n\n**A:** Thank you for your valuable suggestions. During the rebuttal period, we conducted validation for MQAR, MAD, and long-text tasks based on your feedback. Additionally, we supplemented pre-training from 370M to 3B model scales on 300B tokens. We conducted more experiments around the existing work in this paper, verifying the advantages and shortcomings of MetaLA. Regarding other directions mentioned in the discussion, such as the issue of sufficient training [3] and benchmark evaluation, these are based on different fields of foundational models. We fully acknowledge the research value of these directions, which are essential for improving linear foundational models and will be part of our future work.\n\n---\n[1] Mechanistic Design and Scaling of Hybrid Architectures, In  Arxiv 2024.\n\n[2] HGRN2: Gated Linear RNNs with State Expansion, In COLM 2024.\n\n[3] Never Train from Scratch: FAIR COMPARISON OF LONGSEQUENCE MODELS REQUIRES DATA-DRIVEN PRIORS In Arxiv 2024\n\n[4] Mamba: Linear-time Sequence Modeling with Selective State Spaces, In COLM 2024.\n\n[5] Llama 2: Open foundation and fine-tuned chat models, In Arxiv 2023.\n\n[6] Gated Linear Attention Transformers with Hardware-Efficient Training, In ICML2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Dear ACs and Reviewers,

We would like to extend our sincere gratitude to all the reviewers for taking the time to read our paper and offering insightful suggestions. Linear models have emerged as a promising alternative to transformers, garnering significant interest within the foundational model research community. However, much of the existing work on linear models has primarily focused on application-level adaptations, such as customizing Mamba for various visual task scenarios. In contrast, our work seeks to advance the theoretical understanding of these models by contributing the following:

- **Unified linear attention.** We unify the LinFormer/SSM/LinRNN models with different origins into the form of linear attention.

- **Understanding linear attention.** We interpret the linear models as an approximation of softmax attention, and identify three necessary conditions for the optimal linear approximation.

- **Model design.** Based on the above theoretical understanding, we made three improvements to linear attention: removes the Key matrices, employs self-augmentation and exploits short convolutions.

We hope that our contributions to the unified linear attention framework and our insights into optimal approximation will inspire further principled exploration in this domain and foster the development of more effective linear models.

In response to the reviewers' suggestions, we have also conducted additional experiments:

- **Scalability.** We have extended the model to a 3B parameter scale and a 300B data scale for preliminary validation and validated it on Commensense Reasoning benchmarks. The results confirm MetaLA's scalability and efficacy.

- **Retrieval and long context abilities.** We evaluated MetaLA's retrieval performance on the MAD and MQAR tasks, and its effectiveness in handling long contexts on NarrativeQA and the Needle in a Haystack (NIAH) task.

- **Training efficiency.**  In the accompanying PDF, we provide comparative figures on training throughput and GPU memory usage across various 1.3B-sized models.

We have addressed each reviewer’s comments in detail and hope that our responses clarify any concerns. Should there be any further questions, please do not hesitate to contact us.

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

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

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

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