Linear Attention for Efficient Bidirectional Sequence Modeling
Authors: Arshia Afzal
Elias Abad Rocamora
Leyla Naz Candogan
Pol Puigdemont
Francesco Tonin
Yongtao Wu
Mahsa Shoaran
Volkan Cevher
Venue: NeurIPS 2025 poster ()
Keywords: Transformer
RNN
Bidirectionality
SSM
Linear Attention
Inference
Training
Primary area: deep_learning
TL;DR: We propose LION, a framework for extending Linear Transformers to the bidirectional setting by providing three theoretically equivalent representations: full attention, bidirectional RNN, and chunkwise parallel form.

Abstract: Linear Transformers and State Space Models have emerged as efficient alternatives
to softmax Transformers for causal sequence modeling, enabling parallel training
via matrix multiplication and efficient RNN-style inference. However, despite their
success in causal tasks, no unified framework exists for applying Linear Transformers to bidirectional sequence modeling. We introduce LION, the first framework to
systematically extend Linear Transformers to the bidirectional setting. LION generalizes three core representations commonly used in the causal case—full Linear
Attention , bidirectional RNN, and chunkwise parallel form—to the bidirectional
setting. These forms are theoretically equivalent and enable models to exploit the
strengths of each during training and inference. We prove that a broad class of
Linear Transformers can be extended using LION and validate our framework via
three core examples based on the choice of decay type: LION-LIT, the bidirectional
extension of [25]; LION-D, based on [44]; and LION-S, a variant using selective
decay [34, 13]. Across standard bidirectional tasks, LION enables models to match
or exceed the performance of softmax Transformers, while offering significantly
faster training and more efficient inference than existing State Space Models.

Linear Attention for Efficient Bidirectional Sequence Modeling
Arshia Afzal', Elias Abad Rocamora, Leyla Naz Candogan, Pol Puigdemont, Francesco Tonin, Yongtao Wu, Mahsa Shoaran, Volkan Cevher
EPFL
Abstract
Linear Transformers and State Space Models have emerged as efficient alternatives to softmax Transformers for causal sequence modeling, enabling parallel training via matrix multiplication and efficient RNN-style inference. However, despite their success in causal tasks, unified framework exists for applying Linear Transform- ers to bidirectional sequence modeling. We introduce LION, the first framework systematically extend Linear Transformers to the bidirectional setting. LION gen- eralizes three core representations commonly used in the causal case-full Linear Attention bidirectional RNN, and chunkwise parallel form-to the bidirectional setting. These forms are theoretically equivalent and enable models to exploit the strengths of each during training and inference. We prove that broad class of Linear Transformers can be extended using LION and validate our framework via three examples based the choice of decay type: LION-LIT, the bidirectional extension of [25]; LION-D, based on [44]; and LION-S, a variant using selective decay [34, 13]. Across standard bidirectional tasks, LION enables models to match or exceed the performance of softmax Transformers, while offering significantly faster training and more efficient inference than existing State Space Models.
LION Code
LION Blog
1 Introduction
Softmax Transformers [48] are widely used sequence modeling tasks such as causal language modeling [4, 11] due to their high performance and parallelized training. However, their quadratic computational cost is often limiting [45]. increasing interest in Recurrent Neural Network (RNN)-like models for inference. Causal Linear Transformers addresses this by replacing softmax attention with linear attention, which is equivalent to an RNN with a two-dimensional hidden state [25]. Causal Linear Transformers enjoy fast training via matrix multiplication and efficient RNN-style inference.
To improve the performance of Linear Transformers, several variants of linear attention and state space models (SSMs) have been proposed by incorporating fixed [44, 34, 15] or nepeddde [13, 6, 51] decay factors into the RNN recurrence. These recent Linear Transformers match the performance of softmax Transformers in causal tasks, while retaining efficient and fast inference.
Several real-world tasks are inherently bidirectional and benefit from bidirectional sequence modeling. Examples include DNA and protein modeling [17]. computer vision [39], and biological and chemical sequence modeling [42]. In these domains, bidirectional models often outperform causal ones, for
*Corresponding author: arshia.afalleepf.ch. Arshia Afzal, Elias Abad Rocamora, Leyla Naz Candogan, Pol Puigdemont, Francesco Tonin, Yongtao Wu, and Volkan Cevher with LIONS Mahsa Shoaran with Integrated Neurotechnologies Laboratory (INL)@EPFL.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

instance, Bi-Mamba outperforms Mamba in DNA modeling [39]. This motivates the development of architectures specifically designed for bidirectional sequence modeling
Unlike causal sequence modeling, transformers with linear attention remain largely unexplored for bidirectional sequence modeling. Current bidirectional SSMs are primarily based Mamba [53, 21] and are mostly designed for vision tasks [28, 21] and have not been generalized to the broader class of Linear Transformers. These models typically apply their causal forms in both forward and backward directions (e.g., dual scans in Vim or two separate SSDs in Hydra), which fails to leverage the natural priors of bidirectional modeling namely, the availability of the entire seeueneeuee sequence during both training and inference. As result, their training speed lags behind that of softmax Transformers [9, 18].
To show how Linear Transformers can be applied bidirectional modeling and explore their ca- pabilities, we present the LION framework, a general approach for extending Linear Transformers to bidirectional sequence modeling. LION supports three theoretically equivalent representations-fii attention (for maximum training speed), bidirectional RNN (for memory-efficiee inference), and chunkwise parallel form (balancing speed and memory)-mirorinn the efficiency and speed advantages of Linear Transformers in the causal tasks.
We show that wide range of Linear Transformers can be extended to bidirectional sequence modeling via the LION framework (Table 1). For experiments, we focus on three representative
running examples that cover the different types of decay factors: (a) LION-LIT extension of
vanilla Linear Transformer without decay [25]; (b) LION-D extension of RetNet with learnable, non-selective decay [44]; and (c) LION-S which uses selective, input-dependent decay [34].
All LION models are trained using full attention, enabling the highest training speed among existing SSMs, and comparable to softmax Transformers. For inference, LION offers three representations: (1) RNN for maximum efficiency, (2) full attention for maximum speed, and (3) a chunkwise form to balance memory and speed. We validate these advantages through extensive experiments on standard bidirectional tasks across multiple model scales. LION enables training up to 10x faster than SSMs (c.g., Vim [53]) while matching the performance of SSMs and softmax Transformers. It also achieves superior inference efficiency through its RNN form, shown in Figure 3.
2 Background
Notation. Matrices (vectors) are denoted by uppercase (lowercase) boldface letters, such as for matrices and x for vectors. Scalars represented by lowercase letters, c.g., and @ is Hadamard product.
Causal Linear Transformers: Transformers with Linear Attention
Given sequence X = X2,........ Xx] E RLdxd a single-head softmax attention is defined as:
exp(q, k,)
(q,, k; v.) = (x: x. x W y: E exp(q kp)
VJ+
(I)
where x,, qi, k, Y: E Fdd and the weights Wq. Wk. W. Rddd with d being the projection dimension. With Q == [q.] K := [k,...,ki,, enable parallel training, the attention output can be expressed in the following matrix form:
V = i vt] E R4xd To
Y softmax (QK + +M²)
where Mc E (-000,0) is a causal mask for preventing future tokens to attend to past. In contrast, equation is used during inference for generating or processing tokens. However, for causal Transformers [27], employing equation requires storing the previous L tokens to attend to the latest token during inference (i.e., the KV cache). This approach is less efficient than RNNs, where only the state is stored regardless of the previous sequence (cf., [32]).
Causal Linear Attention replaces the exponential kernel exp(q, k,) with feature map function o(q:) T okkk) where o6.) _R" + gdd maps the input to a higher-dimensional space [25]. For simplic- ity of notation, we use q o( W x:) and k, := o( <x,;) in the sequel. This formulation allows the linear transformer to be expressed an RNN with linear recurrence 2, which eliminates the need
"This models with 2D hidden state also known as "fast weights" [20, 41] and their connection to transformers were explored in [40].

---

to store previous tokens in inference, while enabling parallelized training via matrix multiplication:
9. St
Y = SCALE (QK"
MOM V, S. 1-S,i+kky
7 ki, y;
q: 2,
Here, the causal mask M E {0, E is a lower triangular matrix that enforces causal constraints. S-ALB(-) denotes the scaling of the attention matrix across its rows (SCALE( A)
A, and
Ep=i Aip
S; Rdxd and 2, Rd the hidden state matrix and the hidden state scaling vector. Eq. (3) can
also written as Y = PV with P E RLX2 known as sequence mixer [21].
To enhance the performance of linear attention and incorporate relative position encod- ings into the recurrence, decay factors (A₂) have been introduced into the state update:
Y = SCALE (QK O M) V. A, +
(4)
The mask M E RLXL is lower-triangular mask generated based on the input-dependent (selective) parameter x as bellow:
i2 j < j.
M
(5)
0,
Several Linear Transformers have been developed based on equation 4 [5, 34, 3, 44, 51], and can be unified under the form:
q: S
S. - 2, = A, - Y;k; SCALED y; q:
NON-SCALED : y, qq S,. (6)
Here, A; and 7/s denote model-specific decay factors and stabilizers. In this study, we focus scalar or diagonal forms A,, corresponding to the TC class as stated in Merrill et al. [31], which encompasses most Linear Transformers. This unified view is also referred Linear Recurrent Models [52].
Chunkwise parallel form of causal Linear Transformers
Causal Linear Transformers balance the memory-speed tradeoff during training by employing chunk- wise parallelization [51, 52, 6]. this approach, the input sequence X L and the corresponding query, key, and value vectors split into & č non-overlapping chunks, each of size C.
Let Qie (j: ViJ E Cxd represent the chunked query, key, and value matrices, and Sy S define the chunk-level hidden state after pro- cessing chunks as S E Rdxd, The causal sequence mixer which maps the input to out- put Y PV is expressed as:
KC+1) Σ 4-0+4
kj v3
KWy"
VH
(q aM))
Y PV
inter
intra
The above form is the chunkwise form of Linear Transformer without decay factor and the mask Mc {0, 1 is the causal mask. Chunking is essential for training Linear Transformers and SSMs with decay factors, such as Mamba-2 [6] and GLA [51] even more, since treating the entire sequence mixer P as single intra-chunk block, is numerically unstable. This instability arises from computing the cumulative product of decay factors I A, and their inverse I A, over the entire sequence, which can overflow or underflow even logarithmic space". Consequently, SSMs like Mamba-2 [6] need chunking for training stability, which limits their train speed, particularly for short sequences [26]. Chunkwise form has a complexity of O(Ldd) and requires O (t) sequential steps, as opposed RNNs with a complexity of (T.) and (LL) steps, and attention with (II2) complexity and O(1) steps [52].
3 LION: Unified framework for bidirectional Linear Transformers
In this section, we introduce the full linear attention (Eq, (8)), its equivalent bidirectional RNN form (Eq. (19)), and the chunkwise formulation (Eq. (21)). These are presented together to highlight their
In log space, I A, = exp (- E'n log( log(A:)) yet instability persists, as discussed in the official Mamba2 blog post here. Chunking is therefore required for stable training.
3

---

equivalence and demonstrate the range of inference strategies supported by our framework. We begin with the scalar decay case due its simplicity and effectiveness.
3.1 LION Full Linear Attention
In bidirectional tasks, the entire sequence is available during both training and inference. For softmax-bas based attention, this results in the form:
Y softmax (QK
similar to Eq. (2) but without causal masking. Inspired by the natural formulation of full softmax attention and casual Linear Attention Eq. (4), we formulate the Full Linear Attention as:
SCALE (QK') C M) V
The key construction part to define the bidircctional mask M in way that mirrors its properties in the causal setting. To motivate this, we begin with remark the structure of M in causal setting:
Remark. Causal Linear Transformers encode relative positional information through their decay factors, which reflected in their causal masks M (Eq. (4)) [44, 6]. Specifically, the causal mask between tokens and is given by M. Aj+1d ... Xi representing the product of all decays between positions j and (i 2 j).
Inspired by the above remark, we define the mask M for bidirectional sequence modeling such that Mi, equals the product of all decay factors between tokens and j. Formally, the selective M. learnalle fixed decay M. and all-ones M masks can be written as:
IT
i>
Mij
i= j
M2
X|=
Mij 1.
IT
R= AEe
i <<>
Therefore, in the most general form (selective decay), the output of Full Linear Attention is:
1
A
Asda
Ae...Ar
9 ki
4: ka
A
A
V1 VJJ
ki
9 kk
kg.
Y = SCALE
A,Aa
(10)
kg
q/.ks
k.kk
A2-1 Au
A Aa
Ag
QK
where we use diagonal of the attention matrix and mask. The selective and fixed learnalle masks in Eq. (9) have structured forms enabling efficient implementation via matrix multiplications. For the vanilla Linear Transformer [25], due to the absence of decay factor, the mask simplifies an all-ones matrix and can be omitted.
for upper triangular elements, for lower triangular elements, and for the
Selective Mask:
To build the selective mask M, decompose it into lower and upper triangular components (including the diagonal), denoted by M and " respectively. Since both masks rank-1 semi-separable (see proof in Appendix B.8) [21], they matrix multiplication. In the bidirectional setting, where all decay factors X; are known, we stack them into vector X E R analogous to stacking queries and keys into Q and K. The cumulative product LF cumprod( A), where L," -I dę. is used to construct the lower triangular mask MF For the upper triangular mask, we flip the decay vector be reversing the order of the sequence. and compute cumprod (FLIP((A), The two masks are then given by:
be constructed efficiently via
M = [L
M Triu
(11)
I
I
where Tril(-) and Triu(-) extract the lower and upper triangular parts, respectively.

---

For numerical stability, we compute the masks in log-space by defining a = log(A), where E R originates from Zero-Order Hold (ZOH) discretization (see Appendix B.11). The cumulative products L and L can then be computed as cumulative sums in log-space, followed by exponentiation:
D cumsum (a), D cumsum (FLI((a))
L L exp ) L exp(DF)
The full selective mask (Eq. (9)) is then assembled
M M + M I, as shown Eq. (13).
Learnable Fixed Mask; For learnable fixed decay factor, the mask M forms a Bere-MMudocc--aae Szegö (KMS) matrix [23], and stable when 0 1. Our codes for generating the masks are provided in Appendix C.17).
All-ones Mask: As for the case of vanilla Linear Transformer, due to the absence of decay (A = 1),
the mask is omitted and simplified to M 1.
3.2 LION RNN: Equivalent bi-directional RNN for Full Linear Attention
We first show that the naive summation of two separate linear transformers suffers from double counting and imbalanced attention. Considering two linear transformers, we have:
Sf Σ k
S = kky y. y=
+ + - = (k.v, Σ k,,
This summation causes imbalanced attention, i.e. ( QK) due to double counting of the diagonal and underperforms balanced version as shown in Table 16. To ad- dress this, we construct a bidirectional RNN formulation for balanced attention (Eq. (55)), by decomposing the attention matrix A and mask M into lower and upper triangular parts:
4: k:
q. k2
qi ke
19p ko
qike
q. ke.
q k.
q kk
q ke
qf ki
ka
Aaj ka
q k.
(12)
:
k
kg
qikk kg
k
9 kg
t
QK'
A P
A 13
Apho
Au Au
Au
Au Ao
A
Ae
A,
Au A.
ArAo
(13)
Au A.
A
M
M
M
As in Eq. (12) and Eq. (13), the attention matrix and mask are split into lower (A M and upper triangular (A B MB) matrices. The scaling operator divides each row of the attention matrix to its summed value, and hence equals to diagonal matrix C multiplied by the attention:
Y (SKA M)) (QK" o M))V, C. q MM,k
(14)
Decomposing C into causal and non-causal parts as C, = q. Σ Mijky + qf E Mysk, q. k;, we can similarly split the scaling matrix into two parts as:
C
q: Σ
k; MyJK, - ja
q
Σ
1 - k; M k,
j=
OF
Since A = A + A and M M + M
I, we can rewrite the attention output as:
Y QUK M)) V = (C [QKT MM)V
= CC C - + A ) M
+ M
D)V
(15)
= (C" + C M + A O M
A
о M
+ A O R A A evvv.
Since the forward and backward recurrence matrices (AF A
B for attention and M M for mask)
5

---

only share the diagonal with each other, and the diagonal of both forward and backward recurrence masks consists entirely of ones, we can simplify the above equation follows:
= (C cC
M
A
M
O M IA M OI
AoI-A
A
A
C
C
(A M
+ 3
(16)
FORWARD
BACKWARD
As from Eq. (4), the FORWARD part above can be expressed an RNN. We demonstrate that the BACKWARD recurrence term can also be represented by the same RNN in reverse. We re-write the Eq. (16) by flipping the vector V as:
Kg 2 9... q. ke.
0l_ Kr.i ML
A
RD
O
ALAL-1
AL-1
(17)
9: kg
9 Ki-1
9. k1
A A2 A
2
F(M
F(A B)
Above is the exact representations for the forward pass, as shown in Eq. (16), but with the tokens in reverse order. The matrices B and M are also modified match the final flipped output using flipped input values V using functions F(X) = JzXJ L and FLIP(X) = Jr where Jı is an L-dimensional exchange matrix, as detailed in Appendix B.4. Thus, the outputs of the forward and backward recurrences can be expressed as follows:
C Y + Y ). where
(18)
Y PA O M Y (A LABOME O FL.P (F(AB) O F(M
Theorem 3.1. (LION-RNN) Since Eq. (10) is the parallel form of the recurrence presented in Eq. (6).
we can therefore express the equivalent recurrence for full attention Eq. (10) as follows:
SPIn 1/S + K
(19a)
F/B
q.
T
k
(20a)
z. F/B \,2,-1
T/B +K,
2
(19b)
y + y!
F/B G qi 2,
F/B qi k.
(19c)
OUTPUT: y;
of +
(20b)
The terms 1. k and g. k; subtracted to avoid double counting.
Several Lincar Transformers can be generalized within our framework as shown in Table where many causal recurrent models adapt to the bidirectional setting. Since evaluating all models at all scales is infeasible, we focus on three representative examples based on the choice of decay A;:
LION-LIT for A which bi-directional form of Vanilla Linear Transformer [25].
LION-D for X; À fixed decay, and bi-directional form of RetNet (with scaling) [44].
LION-S where >; = o Wx, +)) the bidircctional extension of GRFA [34] (with shifted
SiLU activation) and also inspired by the selectivity of Mamba2.
For all three variants: LION-LIT, LION-D, and LION-S, we use shifted and normalized SiLU activa- +00.5
tion defined as Φ(x) = [[SLL [19]. To ensure stability in the recurrence (0 X; S 1) [32], we set A; oWW] x; + b) for LION-S and À a(a) for LION-D, with o(-) as the sigmoid and learnable scalar. An ablation study on different choices is provided at Appendix C.3.
3.3 LION-chunk: Chunkwise parallel form of Full Linear Attention
Having established the Full Linear Attention (highest speed) and RNN (highest efficiency) represen- tations, in this section we build the chunkwise parallel form form, which balances memory and speed. In the bidirectional case, chunkwise full attention only contains inter chunk leading to:

---

Table I: Mapping Linear Transformers to their bidirectional forms using LION with both Full Lin- Attention and equivalent bidirectional RNN representations. Proofs provided in Appendix B. While RetNet does not originally include scaling, we include it for completeness. LION-S closely follows GRFA, differing only in the non non-linearity o(-). Even non-diagnal models such DeltaNet can exploit LION framework as detailed in Appendix B.7.
Mlnid
Cansal Recurrmee
LON Bedlrvetionll RNN
LIox Fall Linar Altention
LinAn [15]
R. kum! y, q. St
(s) (s
1k.sch of typ
Y= QK'V
201 u (n) (k)
Sealing
q] S./n.
yf (s)
Hkor)) y. Pi LION-UT
SCAER JV
RetNet (1i)
y, q. S.
18... K.
87 x + kiw!
Hkor)) yi- of ty
Y= [QK' o MIV,
M. AP-1
Scaling
y. qi Si/a,
Ae, +l
v, q. (sf fi Hhrre N1
24 1/8 Ki. of qi
LION-a
Y- SE MIV,
M. 1h-1
8 WW kew
Y- SCALE a M.VV
tts
Gated RFA [34]
y. 3. n. S./9. x.
v, (s) (s Akove). SES LIONS
2/10 oWWxin 1
47 a T0
[k)
M
ff
iimj.
Mamba
y. q.
87T X. kww
2 Tis Hkrn), of tP
(QK MMV
XWKV-0ES]
8/10 Diap + Key
Y Trili a .V
Mamba [13] (S1
q.
y Ts! (krn)) M. of toP
Y2 TriniQ LPHK J
Y +
8. Dug(Ajj + (1 Ajv/
8 Di (1 AdVT. A1-
Y7 TrIiQ o (LF)
[16
x. of (s) fhrnh IN oe-
Y Triul 1
8 ASS HAA
SCALE QQU MIV.
8. (. + &
m A dE>
(.STM[]]
q Sa/maol 1)
1421 Ale 6717 - 14 th0
M. Hiasfa i
1 o. (8) Hkos))
ni-A
OK yTV
DelnnNer [6]]
8. (1 Akkk 18 Ahre q. S.
sn Akkl sh A.
(om) y wo
T I+ milidag & KK
HT T
Dy dingt
T
T minldngs A K.K K
digg
Pill Alizztinn
Cuirking
RNN
Pull Anaran
Chinking
RNN
Pull Amenion
Chonking
OQM
O0M
OpM
OOM
OOM OOM
60
Infermue Taace (us))
GPt: Manuuy (GBB)
Informac Time (mn))
CPC SSemuy (GB)
Infereuse Tims (mx)
Momory
Beitee
10"
20t
40x
13
15
10
101
1360
115e
Resolution
Reselttion
Resolition
Resclution
Resolation
Resoltion
Figure I: Memory-speed tradeoff's for LION in three representations on Imagenet classification task: Full attention, RNN, and chunkwise from for Left) LION-LIT, Middle) LION-D, and Right) LION-S. Theorem 3.2. (LION-Chunk) Considering the chunks for queries, keys and values as
Qga Kig, VIJ E Cxd with chunk size being C and lotal number of N chunks, we can
chunk the full Linear Attention as:
Spaj CjNn]
A [ul
K
Mjj
Cju Ciy-1l Sum(AAu))
Spo
S + Ang Vuj-
(21)
where Sum operations applies summation over the row of the input matrix.
The chunk hidden states Cjoj and S(zj] iterate over j. with the final output for chunk i computed using their last values at j = N. The chunk mask M corresponds to a submatrix of the full attention mask from Eq. (13), defined as:
M [s = Mc++--C--,,J----HCH E CxC
Below, we construct both fixed and selective chunked masks Maj)- For the fixed mask, we have:
L
ifi>j, j,
M
ifi j.
=, Тg -li-
(22)
if==j.
7

---

Table 2: Image classification results. Left) We report Top-1 accuracy and relative training time (L) on ImageNet. & denotes models with mutliple scans. Best and second-bes best (excluding ViT) are in bold and underline, respectively. Right) Accuracy vs. Training Speed scatter plot in small scale.
Acc vs. Train Speed
Model
#Params
Top-1
Acc. (%)
Train Time (L)
LIGNG
Small
Base
Small
Base
Small
Base
LIONS
VIT
22M
86M
72.2
77.9
x
x1
Hlydra
78
Hydra
DeiT
22M
22M
91M
86M
79.8
78.6
81.8
81.0
x2.50
x
x2.51
LIONS
(%6
Vim
26M
98M
80.3
81.9
x 14.95
x x10.86
76
LION-LIT
22M
86M
72.4
74.7
x0.74
x0.73
LION-D
22M
86M
73.5
77.8
x..49
X x1.39
LOND
74
LION-D
22M
86M
79.9
80.2
X 1.666
X x1.48
LIONur
LION-S
22M
86M
74.0
76.3
X2.03
X 1,.6
72
LION-S²
22M
86M
79.6
79.9
X2.72
X 1.68
10
10
Relative Training Speed
with L E nº and I E RCC being the vector and matrix used for creating the mask M,yl and they only depending the decay parameter and the chunk size C. For the selective mask have:
L L
ifi
L cumprod(A,ccuu(tncc
M
ifi J.
L cumpro((flip/))ccc-mmtt (23)
Tril
Triu
[
I if = J.
The chunkwise form in Eq. (21) be parallelized over index i, reducing sequential operations. As bidirectional tasks require per-token outputs, memory remains subquadratic at OLCC). We use this the default chunkwise mode for LION. Further visualizations and detailed proofs of LION-chunk are presented at Appendix B.9. LION full framework is shown at Figure and tradeoff between different representation of LION is shown Figure
4 Experiments
To validate the speed and efficiency of LION we focused on well-known bidirectional sequence modeling tasks: Image Classification on ImageNet-IK [38] and Masked Language Modeling (MLM) on the C4 dataset [8]. We also conduct experiments on the LRA dataset to ensure the stability of LION. For Causal Language Modeling and additional experiments, we refer to Appendix A and C.
4.1 Detail of LION and baseline for MLL and image classification tasks
We evaluate the LION framework through its bidirectional variants, LION-S, LION-D, and LION-LIT, on image classification and MLM task. For vision tasks, we compare against softmax- Trans- formers (VIT [9], DeiT [47]) and state-of-the-ar bidirectional SSMs (Vision Mamba [53], Hydra [21]). For Masked Language Modeling, we benchmark against BERT [10] and Hydra. We replace the atlention in the original DeiT and BERT codebases with LION as a drop-in substitution, with- out modifying any other parts such as configurations, optimizer settings, or data augmentation parameters.
Moreover, Linear Transformers do not use external positional embeddings in vision tasks, they often capture spatial structure through different scans by traversing the image along different paths, resulting different token orderings [35]. Each sean induces a different decay mask M. We show that LION naturally supports such variations through alternate token sequences, leading models like LION-DE and LION-S which capture spatial information similarly to existing vision Linear Transformers and SSMs [1]. Details of this strategy and its implementation for LION are provided in Appendix C.14.
4.2 Training speed and inference efficiency for image classification
Training speed and accuracy To support our claims, we pre-train three LION variants (LION-LIT, LION-D, and LION-S) across Tiny, Small, and Base scales for image classification. As shown in Table 2, all LION models outperform the vanilla ViT and perform close to DeiT. Notably, LION-DE

---

a) Transformer Block
b) LION Attention Parallel
c) LION RNN Recurrence
Norm
Scale
lav
Sequence
Mixer
A
M
Norm
Forward Recurrence
Backward Recurrence
M
M
Channel Mixer
Proj KQVA
Proj KQVA
X
Proly
Proj
Proj
X
Forward Sequence Modeling
Backward Sequence Modeling
Non-Lineer Operations
Input Projections
Channel Mixing Block
Sequence Mixing Block
Figure 2: (Left) Standard Transformer block. (Middle) Training mode of LION with full linear
attention. (Right) Inference mode of LION in the equivalent bidirectional RNN.
and Vim are the only models surpass DeiT the Small scale, with LION-D' being 10x faster to train, shown in Table 2 (Right). Tiny scale results are in Appendix C.8, and training details in Appendix C.11. The LION-S? and LION-D variants lie the optimal region of the accuracy-speed trade-off (top right corner of Figure in Table 2), achieving softmax-level accuracy and outperforming saateoofthhe-an SSMs in training speed due their full-attention. Notably, LION achieves its fast training entirely through PyTorch matrix operations, while still outperforming SSMs like Vim and Hydra despite those relying on custom kernel implementations for speedups.
4.3 Inference memory-speed trade off for LION in RNN, Attention and Chunk form
As we also show theoretically, LION-chunk balances the tradcoff between memory and speed. To illustrate this tradcoff, we evaluated LION using three inference strategies: (1) full linear attention (Eq. (8)). (2) bidirectional RNN (Eq. (19)), and (3) chunked attention (Eq. (21)),, across all three decay variants: LION-S, LION-LIT, and LION-D. As shown in Figure 2, the RNN form is the most memory-efficient, while full attention consumes the most memory. Chunked attention strikes balance, with memory usage between the two depending chunkwise provide significantly faster inference than the RNN form, specifically for LION-LIT and LION-D. Chunkwise form enables faster inference than full attention and is more memory-efficient, making a practical choice for balancing speed and memory for LION-D and LION-LIT. The effect of chunk size is further explored in Figure 4. Note that LION-D (LION-D*) achieves the highest accuracy in image classification with its three representations shown Figure (Middle).
type of mask. Both full attention and
We study the effect of chunk size on LION inference with 1248-resolution images, observing trade-off where larger chunks reduce time but increase memory, with OOM beyond size 1024, and we typically use 8-16 chunks balance as shown in Figure
4.4 Masked Language Modeling (MLM)
We evaluate BERT, Hydra, and LION variants using the standard pre-training recipe widely adopted in prior work [7, 10]. We pre-train models from the LARGE family and fine-tune them on the GLUE benchmark [49] to assess LION performance on large-scale masked language modeling. As shown in Table 3, LION variants with 334M parameters closely match the performance of softmax-based BERT and state-of-the-are SSMs like Hydra, while being 3x faster train. Thanks their RNN formulation, LION models are also significantly more memory-effictt than BERT during inference

---

Table 3: C4 MLM and GLUE results for the LARGE scale (334M). Best and second-best re- sults are in bold and underline, respectively.
OOM
COM
COM
COM
5 Samenn
Model
MLM Acc.
GLUE
Train. time
GPU Memory (GB)
BERT
69.88
82.95
x1
50
Hydra
71.18
81.77
x3.13
40
LION-LIT
67.11
80.76
x0.95
LION-D
68.6
81.34
x1.10
30-
LION-S
69.16
81.58
x..3
20
Table 4: LRA Ablation. PathX and average sults for LRA benchmark (best in bold).
10-
VIE VIM
LION WT Chank
Hydra
Model
PathX
Avg.
224
512
738
1024
1248
2496
LION-LIT
X
50.41
Resolution
LION-D (w/ HIPPO)
97.28
85.63
Figure 3: Inference efficiency of LION. GPU mem- ory usage of LION in RNN form and baselines vs the resolution of images for in base scale.
LION-S (w/ HIPPO)
97.99
86.07
Inference memory efficiency While LION matches the accuracy and training speed of softmax-basee Transformers, scales linearly during inference via its RNN formulation, unlike softmax at- tention which scales quadratically with sequence length (Figure 3). Among all baselines, LION achieves the most effi- cient inference scaling. This highlights the importance including both atten- tion and RNN representation for Linear Transformers bi-directional modeling to achieve the best of both worlds: fast training and efficient inference. More- over, inference times for all models provided at Appendix C.10.
80
Memory (Left)
GPU Memory (GB)
Time (Right)
10
Time s
40
OOM
16
32
64
128
256
512
1024
Chunk size
Figure 4: Effect of chunksize in GPU memory and Speed of LION-chunk for LION-D.
(Figure 9), consistent with results from image classification. LION achieves significantly faster train- ing than Hydra while maintaining comparable performance, marginally lower at larger scales and higher smaller scales. For detailed results on BASE scale and ablations please check Appendix C.6.
4.5
Long Range Arena stability ablation
To assess the stability of LION with diagonal decay (described Theorem B.1). specially for long sequences, we evaluated it on the LRA benchmark. The diagonal decay factors of both LION-S and LION-D were initialized using well-known and standard HIPPO-based diagonal initialization [14, 16]. For these tasks, the selectivity of LION-S is defined as A; exp( (a(((( A,) + B;) and for LION-D as A; = exp(p B.), where B, follows the HIPPO. With this setup, both LION-S and LION-D successfully solved the LRA tasks (Table 3). In contrast, LION-LIT in the absence of decay unable to solve LRA, consistent with prior findings [45, 32]. For random initializations ablation and configurations please check Table 5 and Appendix C.3.
Conclusion
We introduce LION, as the first unified framework to adapt Linear Transformers into their bidirec- tional counterparts. LION supports three main representations: full linear attention, bidirectional RNN, and chunkwise parallel form. These forms are theoretically equivalent and allow the model to benefit from high-throughput training using full attention, and efficient inference using either RNN or chunkwise computation. We also provide theoretical mappings several existing Linear Transformers into their bidirectional versions using LION (Table 1). We evaluate three representative LION-D, and LION-S -bsed - on different decay types, on masked language modeling and image classification tasks.
10

---

Acknowledgments
We would like to thank the reviewers for their valuable feedback, which significantly improved both our paper and its presentation. We thank Dr. Igor Krawezuk for insightful discussions on LION models for the masked language modeling experiment. This work was partially supported by Hasler Foundation Program: Hasler Responsible Al (project number 2 21043), partially sponsored by the Army Research Office and was accomplished under Grant Number W911NF-2 24-1-0048, and partially funded by the Swiss National Science Foundation (SNSF) under grant number 200021- 205011. Additionally, this work was supported under project ID #37 as part of the Swiss AI Initiative. through a grant from the ETH Domain and computational resources provided by the Swiss National Supercomputing Centre (CSCS) under the Alps infrastructure.
References
[1] Benedikt Alkin, Maximilian Beck, Korbinian Pöppel, Sepp Hochreiter, and Johannes Brand- stetter. Vision-LSTM: xLSTM as generic vision backbone. arXiv preprint arXiv:2406.04303, 2024.
[2] Jimmy Lei Ba. Layer normalization. arXiv preprint arXiv: 2016.
[3] Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp. Günter Klambauer, Johannes Brandstetter. and Sepp Hochreiter. xlstm: Extended long short-term memory, arXiv preprint arXiv:2405.0455 045177, 2024.
[4] Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. 2020.
[5] Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song. Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy Colwell, and Adrian Weller. Rethinking attention with performers. In International Conference on Learning Representations, 2021.
[6] Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. Forty-first International Conference on Machine Learning, 2024.
[7] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), 2019.
[8] Jesse Dodge, Maarten Sap, Ana Marasović, William Agnew, Gabriel Ilharco, Dirk Groeneveld, Margaret Mitchell, and Matt Gardner. Documenting large webtext corpora: A case study on the colossal clean crawled corpus. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, 2021.
[9] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representattons, 2021.
[10] Daniel Y Fu, Simran Arora, Jessica Grogan, Isys Johnson, Sabri Eyuboglu, Armin W Thomas, Benjamin Spector, Michael Poli, Atri Rudra, and Christopher Ré. Monarch mixer: A simple sub-quadratic gemm-based architecture. In Advances in Neural Information Processing Systems, 2023.
[11] Team Gemini, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean-Baptinn Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. Gemini: family of highly capable multimodal models. arXiv preprint arXiv:2312.11180 2023.
[12] Aaron Gokaslan and Vanya Cohen. Openwebtext corpus. http: http:////ylloooo www.00....ttuu..... io/ OpenNeeTeetCCrppus, 2019 2019.
11

---

[13] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, In Conference on Learning and Modeling (COLM 2024), 2024.
[14] Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. Advances in neural information processing systems, 33: 1474-1487, 2020.
[15] Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations (ICLR 2022), 2022.
[16] Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are as effective as structured state spaces. In Advances in Neural Information Processing Systems (NeurIPS 2022), 2022.
[17] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners In Proceedings of the IEEE/CVE conference computer vision and pattern recognition, pages 16000-16009, 2022.
[18] Pengcheng He, Xiaodong Liu, Jianfeng Gao, and Weizhu Chen. Deberta: Decoding- eechaing-enhanced bert with disentangled attention. arXiv preprint arXiv:2006.0366 2020.
[19] Alex Henry, Prudhvi Raj Dachapally, Shubham Pawar, and Yuxuan Chen. Query-key normal- ization for transformers. arXiv preprint arXiv:2010.0424:3 2020.
[20] Geoffrey E Hinton and David C Plaut. Using fast weights to deblur old memories. In Proceedings of the ninth annual conference of the Cognitive Science Society, pages 177-186, 1987.
[21] Sukjun Hwang, Aakash Lahoti, Tri Dao, and Albert Gu. Hydra: Bidirectional state space models through generalized matrix mixers, 2024.
[22] Peter Izsak, Moshe Berchansky, and Omer Levy. How to train BERT with an academic budget. In Proceedings of the 2021 Conference Empirical Methods in Natural Language Processing, 2021.
[23] Marek Kac, WL Murdock, and Gabor Szegö. On the eigen-values of certain hermitian forms. Journal of Rational Mechanics and Analysis, 2:767-800, 1953.
[24] Rudolph Emil Kalman. new approach linear filtering and prediction problems. 1960.
[25] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. In International conference on machine learning, pages 5156-5165. PMLR, 2020.
[26] Tobias Katsch. Gateloop: Fully data-controlled linear recurrence for sequence modeling. arXiv preprint arXiv:2311.01927, 2023.
[27] Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. Large language models are zero-shot reasoners. Advances in neural information processing systems, 35:22199 222133, 2022.
[28] Kunchang Li, Xinhao Li, Yi Wang, Yinan He, Yali Wang, Limin Wang, and Yu Qiao. Video- mamba: State space model for efficient video understanding. In European Conference on Computer Vision pages 237-255. Springer, 2024.
[29] Yue Liu, Yunjie Tian, Yuzhong Zhao, Hongtian Yu, Lingxi Xie, Yaowei Wang, Qixiang Ye, Jianbin Jiao, and Yunfan Liu. Vmamba: Visual state space model. Advances in neural information processing systems, 37:103031-103063, 2024.
[30] Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv: 1608. 03983, 2016.
[31] William Merrill, Jackson Petty, and Ashish Sabharwal. The illusion of state in state-space models. arXiv preprint arXiv:2404.08879, 2024.
12

---

[32] Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences. In International Conference on Machine Learning, pages 26670-26698. PMLR, 2023.
[33] Bo Peng, Daniel Goldstein, Quentin Anthony, Alon Albalak, Eric Alcaide, Stella Biderman, Eugene Cheah, Teddy Ferdinan, Haowen Hou, Przemystaw Kazienko, et al. Bagle and finch: Rwkv with matrix-valee valuux states and dynamic recurrence. arXiv preprint arXiv:2404. 05892, 2024.
[34] Hao Peng. Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. arXiv preprint arXiv:2103.02 2021.
[35] Zhen Qin, Yuxin Mao, Xuyang Shen, Dong Li, Jing Zhang. Yuchao Dai, and Yiran Zhong. You only scan once: Efficient multi-dimension sequential modeling with lightnet. arXiv preprint
arXiv:2405.21002 2024.
[36] Zhen Qin, Songlin Yang, Weixuan Sun, Xuyang Shen, Dong Li, Weigao Sun, and Yiran Zhong. Hgrn2: Gated linear rans with state expansion. arXiv preprint arXiv:2404007004, 2024.
[37] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
[38] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma. Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. Imagenet large scale visual recognition challenge. 11(())211-2533) 2015.
[39] Yair Schiff, Chia-Hsiang Kao, Aaron Gokaslan, Tri Dao, Albert Gu, and Volodymyr Kuleshov. Caduceus: Bi-directional equivariant long-range dna sequence modeling. Proceedings of machine learning research, 235:43632, 2024.
[40] 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.
[41] Jürgen Schmidhuber. Learning to control fast-weight memories: An alternative to dynamic recurrent networks. Neural Computation, 4(1):131-139, 1992.
[42] Niklas Schmidinger, Lisa Schneckenreiter, Philipp Seidl, Johannes Schimunek, Pieter-Jan Hoedt, Johannes Brandstetter, Andreas Mayr, Sohvi Luukkonen, Sepp Hochreiter, and Gün- ter Klambauer. Bio-xlstm: Generative modeling, representation and in-context learning of biological and chemical sequences. arXiv preprint arXiv:241 1.04165, 2024.
[43] Jimmy T. H. Smith, Andrew Warrington, and Scott W. Linderman. Simplified state space layers for sequence modeling. In International Conference on Learning Representations (ICLR 2023), 2023.
[44] 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 2rXiv:2307.08621,, 2023.
[45] 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. arXiv preprint +=Xiv:2011.04000 2020.
[46] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. CoRR, abs/2012. 18877. 2020.
[47] Hugo Touvron, Matthicu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International conference machine learning, pages 10347-10357. PMLR, 2021.
[48] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is you need. Advances neural information processing systems, 30, 2017.
13

---

[49] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In Proceedings of the 2018 EMNLP Workshop BlackboxNI.P Analyzing and Interpreting Neural Networks for NLP, 2018.
[50] Ross Wightman. pytorch- -mage--mm------- models, 2019.
Pytorch image models.
https://github.
[51] Songlin Yang. Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficient training. arXiv preprinl arXiv:2312.06635, 2023,
[52] Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transformers with the delta rule over sequence length. arXiv preprint arXiv:2406.. 06484, 2024.
[53] Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, and Xinggang Wang. Vision mamba: Efficient visual representation learning with bidirectional state space model, 2024.
14

---

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:
We have theoretically and empirically evaluated all claims.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation as 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:
We included our conclusion.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect 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 the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
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?
Answer: [Yes]
15

---

We have proved all the theoretical parts of the paper in main body and Appendix.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch 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 properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of 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 or not)?
Question:
Answer: [Yes]
Justification: Details of reproduction are in the Appendix. Our code is also publicly available. Guidelines:
The answer NA means that the paper does include experiments.
If the paper includes experiments, 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 is a specific model and empirical evaluation, it may be necessary to either make 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 can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all 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 a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe 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., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
16

---

Answer: [Yes]
Justification: Code is available.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" 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: /CodeSubmie
/ //nips.com wttpccccpppccccceeececceccccccct / swwisssssnnoossoeeeie onPolicy))) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code 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: Appendix we have explained all details even snapshots of our codes clearly. Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes].
Justification: We have followed the original baselines.
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, statistical significance tests, at least for the experiments that support 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, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
17

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of enrors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes].
Justification:
In Appendix all details provided.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well 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 NeurIPS Code of Ethics https: : /ttp///eerrpseee/ee//ee//e cc cttppbiicc/thchccccteeeeeeeeeeee
didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the
Answer: [Yes],
Justification: does.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations 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: No social impact. Work is mainly theory.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to
18

---

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 train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative socictal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
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)?
Question:
Answer: [NA].
Justification: No high risk models.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have 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 challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators 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: Yes.
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 is used and, if 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 (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's is creators.
13. New assets
Question: Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
19

---

Answer: [Yes]
Justification: Yes.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer:[No
Justification:
No human subject.
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 the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer:[NG
Justification: No human subject.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
20

---

Justification:
No LLM used in this study only for writing and revision and gammer.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to LLM policy (https: / http//uurrssccccccceee /ttrr/ppp..c//ooofeeeeee///////2 for what should or should not be described.
21

---

Appendix
Table of Contents
A Causal Language Modeling
23
A.1 Memory allocation in LION during Forward and Backward recurrences
24
B Proofs
24
B.1
Duality between Linear Recurrence and Attention
24
B.2
Forward and Backward Recurrences Theoretical Details
25
B.3
Efficient and Simple Method for Scaling Attention During Inference
26
B.4
Flipping Operation in Backward recurrence
28
B.5
LION framework for diagonal decay..
28
B.6
Mapping Existing autoregressive models into LION
30
Discussion on DeltaNet
B.8 Mask M &M are Semiseperable with rank-1
33
B.9 Details for LION of Chunkwise Parallel
33
B.10 Parallel Chunkwise Materialization for Parallel Output Computation
35
B.11 Zero-Order Hold Discretization
35
Additional experimental validation
37
IRA full dataset results for LION
37
LRA Configurations LION
37
C.3 Ablation Studies LRA dataset
37
C.4 LRA full Results ................................
38
C.5 Experimental details for the MLM/GLUE tasks
38
C.6 Ablation studies in the MLM/GLUE tasks
40
C.7 Additional experimental results for the MLM/GLUE tasks
40
C.8 ImageNet classification results for Tiny scale Inference time comparison image classification tasks
41
41
C.10 Ablation studies with image classification
41
C.11 Hyperparameters for Training Image Classifiers
43
C.12 Calculation of Number of FLOPS
43
C.13 Distillation Results of LION-S
43
C.14 Different scans vision Linear Transformers
43
C.15 Ablation studies on importance of didrccctiooality on image classification
44
C.16 LION-D Decay factor distribution
45
C.17 Generation of the Mask
45
D Future Direction & Limitation
46
22

---

A Causal Language Modeling
We also evaluated the performance of LION-S with decay factor of X, = oWW,, + b) in causal sequence modeling by replacing the softmax attention in the GPT-2 architecture with LION-S attention, following the setup of Radford et al. [37]. Note that LION-S does not use absolute positional encoding. We train models in the OpenWebTerr corpus [12]. We evaluate the architectures in the 124 M parameter setup. Our implementation is based on nanoGPT*. We use the default GPT-2 hyperparameters and train our models for days in 4 NVIDIA A100 SXM4 40 GB GPUs.
35
LION-S (ID)
GPT-2
30
LinAit
Model
Perplexity
Perplesiiy
GPT-2 Sequence Length
15
GPT-2
(7.42(+1.11)
LinAtt
21.07
20
LION-S (ID)
(8116))
18.25
Avg.
10
1000
2000
3000
4000
Sequence Length
(a) OpenWebTent PPL
(b) Perplexity vs. sequence length
Figure Causal Language Modelling results the GPT-2 128M size. (a) Perplexity in the OpenWebText dataset. (b) Perplexity vs. sequence length in Open WebText. LION-S improve over the LinAtt baseline [25] while obtaining similar performance to the GPT baseline and being able to extrapolate to larger context lengths than the one used during training.
In Figure 5 we can observe LION-S (ID) significantly improve over the LinAtt baseline, while obtain perplexity close to GPT-2. The lack of absolute positional encodings allows LION-S (1D) to scale to larger sequence lengths than the one used during training.
In Figure we evaluate the latency and memory of LION-S (ID) in three modes: Attention, Attention + KV cache and RNN. While the three modes have the same output, the RNN formulation allows to save computation from previous token generations to require constant memory and latency for generating the next token. Our results align with the findings of Sun et al. [44], showing that efficient
hhttpa://githut hwwp:/aaaaaaaaaayynnoorrrr
1.25
Recurrecce
Recurrence
16
Next taion latency (6]
1.00
Atteation
Atsention
Atteation KV Cache
Memory (GB)
14
Attention
KV Cache
12
0.50
10
0.25
7.90
0.00
6.95
0
20010
6000
8000
2000
4000)
6000
8000
Sequenee length
Soquence Length
(a) Latency
(b) Memory
Figure 6: Efficiency of the LION-S framework in the next-token generation task. In (a) and (b) we measure respectively the latency and memory to generate the next token in the sentence. We compare three generation modes: Attention, Attention with KV cache and the Recurrence formulation. While all three produce the same output, the Recurrence formulation is the most efficient, requiring constant memory and latency to generate the next token.
23

---

models in training and inference, with strong performance (up to a small degradation) can be obtained.
A.1
Memory allocation in LION during Forward and Backward recurrences
Backward
Backward
y?
y
y
cf
cf
yf
yf
of
yf
Re-wiitten on same Memory cell
cf
Re-wiitten on same Meroory cell
Forward
Forward
y1
Used for computing the output
F
Figure 7: Memory allocation in LION during Forward and Backward recurrences. The efficient way of re-using the memory during inference is explained.
B Proofs
B.1 Duality between Linear Recurrence and Attention
Considering the following recurrence:
S, = A,S + k 7, + k,,
(24)
(25)
9: S:
SCALED Yi
(26)
qi zi
We calculate each output y; recursively below:
kivi, z1 ki, yı v1
(27)
kav ArkivT, ne = k2 Aski, y== 9: [k.v + (gkiv])
(28)
(k2 + (1k))
S, = kavs Ankavd Dadskivi, zs = ka Arka + AudikI, ya 9. (kav A dedikiv
o (ks + d: k2 )ıd
(29)
of ( Mykk))
M [ i2j i<j
(30)
9. [E Mōk,)
This can be shown in vectorized form as:
Y SCALE QK @
M
(31)
Where SCALE is the scaling function which scaled the attention matrix with respect to each row or can also be written as:
Ay
SCALE
(32)
ELA Ag
Similarly if the SCALE is applied before masking we have:
Y =
(SCALE(QK) McAUSAI. 2 M)V
(33)
With McAUSAL being the causal mask used in autoregressive models [27]. This vectorized form equivalent to:
24

---

y; q: (E== M)
m
R=
Au
i i2) i<j
My
(34)
())=k))) kj)
And the recurrence for this vectorized form can be written as:
S. = A,S + k. = 2 + k;,
(35)
(36)
9: T S:
SCALED Yv
(37)
qi 7j
B.2 Forward and Backward Recurrences Theoretical Details Considering the following recurrence:
S. A + k
(38)
21
k;
(39)
i=1
qi
S
Y1
(40)
ZT.
This recurrence is the same as recurrence equation 35 but with zy being fixed to the summation of all keys in the sequence, therefor the output y; simply be written
9. (Ej= Mykyy)) y:
I 0
i2j i i<j
M.
(41)
9, Z
By replacing the z, ンj ky in the denominator of equation equation with 2t. Therefore in vectorized form, it will become:
Y Y M)V
(42)
With A being:
q. Ri.
9, ki
ni kg
xL
A
ki
ko
ks
qi ki
a ko
9. mL
m
Importantly this equation can be written as:
Y (SCALEIQQK))) M
(43)
which despite equation equation 33 scaling is applied over the whole sequence not for the causal part of the sequence. The matrix A is helpful for driving the recurrent version of LION for Forward and Backward recurrences and the mask here M equal to LION'S forward mask M in equation equation 16. As shown in equation 16 the forward recurrence for the causal part of the attention can be presented as Y A @ F the matrix A can be created simply by using matrix A as bellow:
25

---

1
9
q; Kg x..
4: li
ka
95 kg
2
ki
9:
Ka
ko
2
4: kg
ki
kg
9.
91.5
9, Xy
9..5
4
Or equivalently:
YF= = = (A--DD O
(44)
Since the diagonal values of the mask M are all ones and the matrix D is diagonal, we have:
Y = (AC - D o M = A OM DF
(45)
As A C @ M corresponds to linear recurrence shown at equation 41. The vectorized form equation 45 can be presented as linear recurrence:
= = 1
1 Miskyy))
q q zr.
i2 j
Mi
(46)
q 2r.
i i<j
This is equivalent to the linear recurrence presented in equation equation 40. The same theoretical approach applies to the backward recurrence, leading the following linear recurrence for both recurrences:
S. = X + k.V
(47a)
S = At. + kr.-i+1V
(48a)
S
1 q:
k;
91-1+1 SS 1 1di-it1 KL-it1 B yL-i-1
y.
(47b)
(48b)
qi 2t, 2
q:
Zt.
qr-i41 2r 2 qr-it1 Zr
However, the above equation requires access the summation of scaling values 2t. A naive approach would involve adding an additional scaling recurrence alongside the forward and backward recurrences to compute the summation of all keys in the sequence. This approach, however, is inefficient, as it complicates the process. While the forward and backward recurrences can traverse the sequence in parallel to obtain the forward and backward recurrences outputs Y and B the scaling recurrence must be computed prior to these recurrences because both the forward and backward recurrences computations rely on the final scaling value zz to generate their outputs.
B.3 Efficient and Simple Method for Scaling Attention During Inference
As shown in previous section scaled attention matrix can be formulated as two recurrences equation 47 and equation 48 with an additional recurrence sum all the keys (zz). This section will proof how to avoid an extra scaling recurrence by simple modifications to equation equation and equation 48.
Considering having scaling recurrence part of forward and backward recurrence we will have:
SF = A +
K.V,
(49a)
S
i
kr-
(50a)
F 2 2 + k;
(49b)
2
B
2-1
B
+ ki,iitt1
(50b)
cf = qi T F
2 9i
1
K; (49c)
c
1
91. Ky-141
(50c)
2
y = - TsF
-
qi
T
ki
y yL-141 = q
Ts b
(50d)
(49d)
26

---

The equations above are similar to the previous ones, with the addition of scalar states c" and c² for the backward and forward recurrences, respectively. During each recurrence, the outputs y and y along with the scalars of and cf are saved for each token to construct the final output of each layer. It is also important to note that there is need to save z" and zB for each token; these states can simply be overwritten in memory. The final output of each layer is equal to:
y y Yi
(51)
C + c
Where yf and y can be written as:
yf = q. M kv, 29 kvv
y = q. 2
Myk,v,)
29 T k
(52)
j=z
So the addition yf + y is equal to:
y y B =
M
ky
+ q 1
Mākjy]) - T ki
(53)
j=1
4
y y
kyy
+ Σ
M qi k
(54)
j=i
Where by considering the mask M bellow:
Xa
A2d3
d d
m+
Av
}>j
A
As
As d
Myy-
i<j
A1
1
(55)
i=j
A ... X
---1---A
X
The above mask is equal to M + M allowing equation equation 53 to be rewritten as:
y +yp
q
M + +- MM qi k,
(56)
1=1
L
= q.
Mykyy,) + q. k,v, - q. k.
(57)
j=1
q.
MykyvJ)
(58)
So finally find the output of each layer y; as:
qi (E/ MorkyvJ)
y
equation 58
yi
y;
cf +
(59)
It can casily be shown that:
27

---

1
cf = q.
kj)
24 ki
di = q. 2
kj)
2
k;
(60)
==1
5=1
9 + oj
F
qi
+ qb
29
1 T ki
7
T ki
(61)
j=1
1.
F + C
q
kj) + q ki q. k;
q,
T <
kj)
q 22
(62)
==1
So the final output of the layer is:
yf y af + q ()
q Mc,k,,]
Yi
(63)
Alternatively, in vectorized form, it can be expressed as:
Y Y Y = (SCALE QK O M)
(64)
with M being the attention mask created by A,S as in equation 55.
B.4 Flipping Operation in Backward recurrence
Here define the operation which flip the matrices B M for the reverse reccurence th goal is find the F() such that:
ta ki
q
ke
kL. q
qikkr
a. Rz
1 k2
th ky
3 Si-kk-1
A
FIA
(65)
ka
9y mr.
A
Azds
A Ar As
Ay
Ar.
R
Ae
7 F(MM))
=
Aedema
A
(66)
Ae A Au d As du da
The above can be achieved by:
F(A) JJAJ, JI
(67)
B.5 LION framework for diagonal decay
The recurrence in the causal form for a diagonal case presented as:
S, = A,S + k.V
(68)
28

---

which followed by Dao and Gu [6], Hwang et al. [21] it can be shown in matrix form as:
Y = PV
(69)
P.j Q
A A..1
A K An =1 i>j
(70)
II Ak.
A
(71)
i<j
with P being the sequence mixer of the Linear Transformer. We can build Full Attention starting by decomposing the decay factros A, as followes:
Miy = Q Att K
(72)
Q AoA A (Aa) (Ao) (A1)
(Aj)- (A.))-...(AA) K
Mij - J- K,
(73)
In the above decomposition blue parts are only depending on the index i where the magenta depending on j this implies by defining Q A X Q and Kj- (A ) K, the Equation equation 72 can be written exactly as attention in parallel form M K*T Note that since > all terms with indices < cancel out, leaving only the matrices A₁ with j 1 S k S i to be multiplied in expression A (Ax ) However it is important that the matrices Q* and K be fast and memory efficient to built we show that indeed they are starting by following remark:
Remark. For diagonal matrices Ao; A A,...., A;; the product A = Ao A.... A, is given by
Ao:: ag aj a2 ag
(74)
where a, E RR a vector containing the diagonal elements of the matrix A, E NNN or a, = (A.. Utilizing Remark, we can rewrite equation equation 72 as follows:
Y= MV
My Q A
K,=
Q (Ao o A 2 .2....................... A @ A O o K
A2....A))
(Q, A 2 X O A
(Au A A,) 3 K,
o
(75)
where in the above, (An e o A X,) indicates the clement-wise inverse (Inv) of the vector
Ao @ A A2... Ay. During training, we have access to all tokens. Thus, we can compute the
vectors A; for all and construct matrix D RLLN where the i-th row D; X; DIAG(A;).
Our goal is to find a matrix L E A whose i-th row L, is given by Ao O Y ① A2 0 ... A, This allows us simplify the term Q @Ao C A O ...... C A.) in equation 72 to Q L,
Ii Akp. it can be viewed the cumulative product of the vectors A; Consequently, we can obtain L by applying a cumprod operation along the sequence dimension of D:
Since Lip
L cumprod (D, dim = 0).
(76)
A similar approach applies to (Ao C X ..... @ 11) Instead of computing the cumulative product of X, directly, we compute cumulative product of the element-wise inverse of these vectors, leading to the definition of the matrix U:
L, Ao O A O
(77)
(78)
29

---

Note that U Inv(L) is the element-wii wise inverse of L. and thus L@U- an all-ones matrix.
= 1, where 1 E R L.X is
By analogy with Transformers [48], we can store the states Q K, A, for all tokens in matrix form as Q. K, RLLN, Under these conditions, selective diagonal SSM (Equation equation 69) can be simplified to:
Y TRIL
(QOL) K (L)
V
(79)
Q'
L cumprod(D), with D, = DIAG(A,)
(80)
Theorem B.1. LION-Diagonal Full Linear Attention: For bi-directional case the above alike Eq. (11) can be expressed as:
Y = (TRiL [QQ (K (noInv(L')) TRIU (QOL (K OKo (KoInv(()))
(81)
L cu))rod(D), with D, = -DAAG(A,)
(82)
L cumprod (FFi(((D) (D)). with D. (I,
(83)
The TRIL fanction the above formulation indicating that the output Y computed only from previous tokens, thereby preserving the causal structure analogous to causal Linear Transformers, where Y = (RIL(QQ)) V.
Note that this parallel form is less straightforward than the scalar case, but the masks L, LB can still be decomposed from the queries and keys as shown in Eq. (11). Theorem B. extends LION's Full Linear Attention to the diagonal decay case. While scaling can be applied directly to Y, we omit it for simplicity, as most Linear Transformers with diagonal decay do not use attention scaling
[13, 16, 33].
Theorem B.2. (LION-RNN) The equivalent bi-directional RNN for Theorem B. I is same as LION- RNN presented at Theorem 3.1 without scaling:
q k V.s 2
yf + y
S//B A/B K.V,, / = q T
OUTPUT: y,
(84)
c + c
B.6 Mapping Existing autoregressive models into LION
As noted, other autoregressive recurrent models can also be integrated into our bidirectional frame- work, benefiting from parallelization during training and fast bidirectional inference. Here, we demonstrate how to map several well-known Linear Transformers into the bidirectional form of LION, along with their corresponding masked attention matrix and inference linear recurrence.
Linear Transformer (LION-LIT). According to [25] the linear transformer has a recurrence:
S: S + K.V., of = _-_ + ki,
(85)
(86)
S qi z
SCALED yf
(87)
NON-SCALED of q. S
(88)
As observed, this is special case our bidirectional recurrence defined in equation with X; as LION resembles the scaled masked attention. In the case of the linear transformer, we require attention without scaling for the recurrence. The vectorized form for the scaled version can then be derived easily follows:
30

---

S/1B S.1B k.
(89)
2. F/B 2 11 +
F/B
(90)
Y = SCALE(QK
(93)
9 6/1 qi 2
F/B
qi
k;
(91)
F/B yi
S F/B
1 2 kiV;
(92)
For the non-scaled variant, we simply remove the scaling state as well as the scaling parameter Consequently, the bidirectional linear transformer, which is equivalent to and parallelizable with attention without scaling, can be expressed as follows:
SPIB - F/BB KVV F/B = qi S F/B
(94)
Y QK
(96)
k.
(95)
The final output for scaled version can be extracted as y; y. ty for scaled and y, y. + y +c
for non-scaled version. Variations of linear transformers, such as Performer [5], which employ different non- non-lincarities 6(.) for keys and queries, can be adapted to bidirectional format using the framework established for linear transformers.
Retentive Network (LION-D). According to [44] the forward equation for retentive network can be written as:
S = XS + k yf qi S
(97)
(98)
This architecture can also be expanded to bi-directional setting simply by not scaling the attention in our framework and only using the mask with non input-dependent X; values:
SFIB -S +
(99)
Y (QK MR
(101)
B qq: T q: k,v;
(100)
Note that: M Xاا and LION-D uses scaled version of above.
XLSTM (LION-LSTM)) According to [3] the recurrence for forward recurrence of xLSTM can be written as:
S. = J.S + +is.V
(102)
2. = fiz Hiki
(103)
y q: S
(104)
qi 2i F
The above recurrence is equivalent to equation 19 by considering i,k, new key. The term i,k; can easily vectorized by aggregating all i; values for each token into a vector Thus, we can express the vectorized form of the bidirectional xLSTM and its equivalence to attention as follows:
31

---

s/ = f.S 1-1 + i
(105)
2 F/B fzz F/B iilk
1-1
(106)
F/B
T F/B
2
2 41
k. (107)
Yi
F/B
qi T S
F/B
2 qi ki
Y SCALE Qii ④ V
(110)
(108)
y y B
Output: y: max +
(109)
where the mask M is equal to the LION mask equation 55 just by replacing X; = fi. And where operation SCALEmaz consider the maximum of operation in the denominator
Aij ij max 2==1 A.j.
SCALEmoz (A)y
(111)
Gated REA (LION-S) Gated REA [51] in autoregressive mode exhibits recurrence similar to that of XLSTMM with only minor differences:
S 9, + (1 -9,,k,
(112)
zf 9,2 -1-10 -i
(113)
y. qi S
(114)
zi
Thus, the bidirectional version of the model retains similar output, achieved by replacing the vector i equation 110 with where g represents the vectorized form of all scalar values gr.
F/B + (1 g.)k.v, S S/1B g.S,
(115)
2, F/B F/B
g: 1-1
9:)
(116)
Y SCALE(Q((1 -g) (ooK) @K)
(119)
=
6/1 q: z:
F/B
1
ki (117)
2
y F/B = q, S F/B
29 9i k,v;
(118)
Note that for LION-S we just use different non-linearty compared to original GRFA and the term (1 9.) can be considered as part of k,.
Mamba2 Mamba2 [6] is the selective SSM with scalar decay factor X; SoftPlus( (x;) we can observe that Mamba simply can be considered as variant of LION-S with different activation and without scaling therefor
extend it to bi-directional setting via:
F/B + K.V. S.P B/B = A/S 1-1
(120)
Y = (QK O MMV
(122)
y F/B = qq T
1 2 q k,
(121)
Here, M is the selective mask used in LION-S. This variant is closely related to Hydra [21], but differs architecturally. Hydra incorporates convolutions and gating, and uses SSDs for training, whereas LION-S trains using full attention without any gating convolutions.
RWKV-6: RWKV-6 [33] Linear Transformer with diagonal decay causal recurrence can be written as:
32

---

S: A + k.V; yi - TS
(123)
(124)
as seen this choice is exactly aligned with our proof on Theorem B.1 which allows the bi-directional form of:
S/ = -,S 1-1 + kv
(125)
Y = (TRIL (Q (Q)L") (K IK)lv(LL)))
yi qi s
F/B B//B__
1
k;
TRIU (Q @ L (K (KoInv(L
2
(126)
(127)
HGRN-2: HGRN-2 [36] has similar recurrence as RWKV-6 just by replacing the k; = Di, therefor its equal bi-directional form via LION is alike Appendix B.6.
B.7
Discussion DeltaNet
For the special case DeltaNet [52], our theorem supports an equivalent bidirectional RNN formula- tion with LION-style correction terms avoid double counting:
S//B = S (I k + 3kkVV.
yi - F/BB -
The chunkwise parallel form of DeltaNet (Eqs. 8-11 in the original paper) also be extended to the bidirectional setting by removing causal masks and applying Theorem 3.2 for the chunkwise form of LION. However, key bottleneck in DeltaNet arises in its full attention formulation, where the attention computed as
QK
with T' (Eq. 10 in the original paper) requiring a matrix inverse that scales cubically with sequence length. This inverse undermines the fast-training advantage of LION, so we avoid using the fully parallel form for training DeltaNet.
B.8 Mask MF &M are Semiseperable with rank-1
For the lower triangular part of the selective mask M the upper triangular part can be filled such that it creates full matrix with rank 1, which aligns with the definition of a semi-separable matrix with rank 1, as below:
M
TRIL
TRIL
TRIn(u))
where TRIL is the function which masks the upper part of the matrix and set to zero. Same is applied for the upper triangular part of the matrix M as well. Also since decay mask is the specific case selective mask by setting A; = À all the proofs above also holds for the fixed decay mask used in RetNet.
B.9 Details for LION of Chunkwise Parallel
As the full linear attention is written as:
Y = SCALE 6 M
(128)
by apply chunking queries/keys/values and defining the:
Q1g. Kij Vya Q,CHLCCH,,KKCHHCCAH ViCHLMC+4 Cxd
33

---

dh
JAA
dh
64
d4
44
44.
db
64
d4
Figure 8: Visualization of attention (Left), LION-D (Middle), and LION-S (Right) chunkwise forms.
we can simply rewrite the Eq. (128) in chunkwise form as:
Aps Q K [
M
Ciyy C(tj-1) Sum
(129)
S Span]
Sty) Sp +
M
Y []]
(130)
C[EN]
where N is the number of total chunks and N = & and Sum is the summation over the rows of the matrix. Since the full attention matrix needs to be scaled according to the full row of attention we need to update the scaling value for each chunk stored in Cı and the final output for chunk i is computed by using the last chunkwise hidden state Sy divided by the scaling for that chunk Cj j] which are equal to S jin: C(iN]-
To construct the chunkwise mask Mij. we define the chunk-level selective parameters as:
L [
cumprod(A (A )1C+11+11 L cumprod(A" )C++1(++1CC
Since the full mask is composed of lower and upper triangular components:
M F = TRIL(L ) M = TRIU(L B
L
we determine the appropriate chunkwise form based on relative chunk positions:
If > j, the chunk falls entirely within the lower triangular part, requiring only M which can be efficiently computed as L
If i j.j, chunk is fully in the upper triangular region, needing only M which follows
from [ L
If j. the chunk lies along the diagonal and requires both the lower triangular part of M and the upper triangular part of S expressed as:
The full matrix is:
FF 1 (1)
if > j. if ifi<j,
LB LI
M [J]
Tril
L 1
T
Triu
LB [1] [R]
if j
For the fixed decay mask, a simpler case of the general selective mask, LF and L 13 are identical and simplify to L, = Xs. Since the full mask follows M.j = X|k-s|, the chunkwise mask for i, can be written as:
34

---

M M
L Lij
A -
L [0] [o]
Similarly, for the upper triangular part:
Mayl = Ali 1
Lio
[0]
For diagonal chunks, the mask remains fixed matrix P E RCxC where Fij X representing smaller version of the full fixed decay mask M RRL with M.j Al-51, The visualization of chunking is presented in Appendix B.9..
B.10 Parallel Chunkwise Materialization for Parallel Output Computation
As mentioned in Section 3.3, Eq. (129) can be further parallelized by processing all i tokens simulta- neously, leading the following matrix multiplications in the parallel chunkwise form:
A j QK [
O MG),
Cbi = C
Sum ( Ap);
(131)
S
S = S + Apol Db
Y =
(132)
[w]
and the mask M3 is created based on:
M TI. 1 L'
diagonal -)) + Tri
diagonal = j]
(133)
L B
Consider real matrix X E Rux" The operator Tril(X, d) returns the lower triangular region of X, including all elements on and below the diagonal shifted by d. other words, tril(X, d), Xij whenever < + d and is 0 otherwise. Similarly, Triu(X, d) returns the upper-trianee prpee-triangular region, keeping elements on and above the diagonal shifted by d. Formally, Triu(X, d)is X., if j 2 zit and () otherwise.
B.11 Zero-Order Hold Discretization
Below we explain the zero-order hold discretization derived by [24]. An LTI system can be represented with the equation:
h1(1) Ah(l) Bx((),
(134)
which can be rearranged to isolate h(t):
ht(t): - Ah(() = Bx(1).
(135)
By multiplying the equation by
At, we get
M hh((t) - At Ah(t) Bx((t
(136)
Since ee AEAl = - A, Eq. (136) can be written as:
ot (e (h())))))))))) = -c-AA
(137)
After integrating both sides and simplifications, get
e -nh(()) = 1/ c - Bx
(138)
35

---

By multiplying both sides by et to isolate h(t) and performing further simplifications, at the end we get
h(t)= - Ar S[ - Bx(T) e (h()).
(139)
To discretize this solution, we can assume sampling the system at even intervals, i.e. each sample is at kT for some time step T, and that the input x(t) is constant between samples. To simplify the notation, we can define hş in terms of hl?T) such that
h& h(kT).
(140)
Using the new notation, Eq. (139) becomes
h&s = e AKT hAkTh(()) h(0) + AKT S
pKT e e -x(T) dr.
(141)
Now we want to express the system in the form:
h Åhe + Bxi.
(142)
To start, let's write out the equation for Xx+1
h&+1 A(k+1))T(O)) + A(k+IjT (k+1)T e - -AT)x(T)dr dr.
(143)
After multiplying by AT and rearranging we get
A(++IT h(0) = AT has A(tt1)T for pht
e (-ATBx(T)
(144)
Plugging this expression for xk+1 in Eq. (143) yields to
Bx J e -AABx(r)drr) h AT A(K+IJT (%ee
KT
(K+1)T
(145)
which can be further simplified
(++1yT
h
AT hE A(++1)T
e e-ATBx(T dr.
(146)
KT
Now, assuming that x(I) is constant on the interval [KT, (k: + 1)T), which allows us to take Bx(() outside the integral. Moreover, by bringing the Ak++1)T term inside the integral we have
(k+1)T
h = AT h4
A((++1)T-1) dr Bx
(147)
KT
Using change of variables v: = (K: + 1)T - T, with dr = dd. and reversing the integration bounds results in
/-"
T A" eA" du Bx
h = AT +
(148)
Finally, if we evaluate the integral by noting that de
At = Ae and assuming A invertible, we
get
h
AT The A AT 1) Bx
(149)
Thus, we find the discrete-time state and input matrices:
À AT
(150)
B A A AT -1) 1)BB B.
(151)
And the final desecrate state space representation is:
h& c h + +A-1 (EAT B
(152)
36

---

As in case of LION-S (similar to choice of Mamba2 [6]) the matrix A is identity while the time step is selective and equal aş- And simply for LION-S scenario the term Bx(t) will change into k, therefor considering Linear Transformer continuous system like:
(/1) = (()) + (
(153)
[1) (1) + (())
(154)
(155)
By applying the ZOH discritization the final descreate LION-S will be equal to:
DISCRETE
S: S + (6 1l)kvv
(156)
2; = 21-1 - 1)ki
(157)
And applies both directions forward and backward.
C
Additional experimental validation
C.1
LRA full dataset results for LION
Table 5: Training Stability on Long Range Arena Task. Our family of models can solve the LRA
benchmark with the appropriate initialization using full attention.
Model
ListOps
Text
Retrieval
Image
Pathfinder
PathX 16K
Avg.
(imput length)
2048
4096
4000
1024
1024
LION-LIT
16.78
65.21
54.00
43.29
72.78
X
50.41
LION-D (w/o HIPPO)
32.5
64.5
50.0
47.4
73.6
X
53.6
LION-S (w/o HIPPO)
36.34
60.87
55.0
42.6
74.2
X
53
LION-D (w/ HIPPO)
62.0
88.78
90.12
85.66
90.25
97.28
85.63
LION-S (w/ HIPPO)
62.25
88.10
90.35
86.
91.30
97.99
86.07
C.2 LRA Configurations for LION
For the LRA task, we utilized the same model dimensions specified in the S5 [43] paper, following the guidelines from the S5 GitHub repository Our state matrix was represented as vector A, A; where each element contains a scalar non-input dependent value e The value was initialized based on HIPPO theory, alongside the input-dependent ú; as described main body.
We employed the ADAMW optimizer with an initial learning rate of 5 x 10-1 and cosine learning rate scheduler [30]. The weights for the queries and keys, as well as the selective component of A. were initialized using Gaussian distribution with a standard deviation of 0. For the values v, we initialized Wy using zero-order hold discretization, represented Wiiiit (A 11-1... (A 1)). The non-selective parts of A were initialized based on the HIPPO [43] matrix.
C.3 Ablation Studies on LRA dataset
We have did an ablation study for choosing the activation functions and using non-scalar decay factor for LION.
We have observed that bounding the keys and queries significantly enhances the model's ability to solve tasks. This finding is consistent with the observations in [52]. As demonstrated in variation [1], it can successfully tackle the LRA task even without scaling, while the RELU activation fails to do so. Additionally, we found that scaling plays a crucial role, particularly when it comes to scaling the masked attention. The approach used LION, which scales the attention before applying the mask expressed as Y = SCALE(( QK K @ M has proven ineffective in addressing the challenging PathX task, as shown in Table (5).
'https://githunn com http:///tthhh.ccco//oooaaaaaaaaa
37

---

Table Effects of different parameter choices and non-linearities in LION-S on LRA tasks. Codes:
[1] Sigmoid non-lineat non-linearity was applied to the k and q values with unscaled masked attention; [2]
ReL.U non-linearity was utilized, and the masked attention was scaled; [3] The parameter o; was
selected as scalar instead of vector; [4] LION-S model parameters were used without scaling; [5]
The attention matrix of LION-S was scaled, but attention values were adjusted without the factor of
As [6] The selective component of
removed; [7] SoftPlus activation function was employed
for the a, values. We used the HIPPO [14] initialisation for LRA task since random initalisation of
LION-S and LION-D can not solve LRA.
Model
ListOps
Text
Retrieval 4000
Image 1024
Pathfinder 1024
PathX 16K
Avg.
(input length)
2048
2048
6(x) = (xx) w.o scaling
61.02
88.02
89.10
91.06
97.1
85.41
((x) (ELU(z) w. scaling
36.37
65.24
58.88
42.21
69.40
54.42
3
a, only scalar
36.23
60.33
60.45
58.89
70.00
X
57.17
LION w.o scaling
58.76
67.22
59.90
60.0
65.51
X
62.27
scaled attention mask
60.12
87.67
87.42
88.01
89.23
82.49
a, From HIPPO w.o selectivity
60.12
88.00
89.22
83.21
91.0
96,30
84.64
7 a; (OFTPLUS(())
16.23
59.90
60.00
45.12
70.07
50.26
LION-S (w/ HIPPO)
62.25
88.10
90.35
86.14
91.30
97.99
86.07
C.4 LRA full Results
We have evaluated LION variants against benchmarks for long-ram Long-range sequence modeling across different categories, including softmax-based Transformers, RNNs, SSMs, and Linear Transformers.
Table 7: Performance Long Range Arena Tasks. For each column (dataset), the best and the second best results are highlighted with bold and underline respectively. Note that the MEGA architecture has roughly 10x the number of parameters as the other architectures.
Category
Model
ListOps 2048
Text
Retrieval
Image
Pathfinder 1024
PathX 16K
Avg.
(imput length)
4096
4000
1024
Transformer
36.37
64.27
57.46
42.44
71.40
X
54.39
Transformer
MEGA (O(L2))
63.14
90.43
91.25
90.44
96.01
97.98
88.21
MEGA-chunk (O(L)
58.76
90.19
90.97
85.80
94.41
93.81
85.66
DSS
57.60
76,60
87.60
85.80
84.10
85.00
79.45
S4 (original)
58.35
86.82
89.46
88.19
93.06
96.30
85.36
SSM
SS
62.15
89.31
91.40
88.00
95.33
98.58
87.46
Mamba (From
32.5
N/A
90.2
68.9
99.2
N/A
N/A
RNN
LRU
60.2
89.4
89.9
89.0
95.1
94.2
86.3
xI.TTM (From [3D)
41.1
N/A
90.6
69.5
91.9
N/A
N/A
Local Att.
15.82
52.98
53.39
41.46
66.63
X
46.06
Sparse Transformer
17.07
59.59
44.24
71.71
X
$1.24
Longformer
35.63
62.85
56.89
69.71
X
53.46
Linformer
16.13
$3.09
42.34
75.30
X 50.55
Reformer
37.27
56.10
53.40
38.07
68,50
50.67
Sinkhorn Trans.
33.67
61.20
53.83
67.45
$1.48
Linear Transformer
BigBird
36.05
64.02
59.29
74.87
X 55.01
Linear Trans.
16.13
65.90
53.09
75.30
X
50.55
Performer
18.01
65.40
53.82
77.05
$1.41
FNet
35.33
65.11
59.61
77.80
X $5.30
Nyströmformer
37.15
65.52
79.56
41.58
70.94
58.95
Luna-2256
37.25
64.57
79.29
47.38
77.72
X
61.24
H- H-Transformer-ID
49.53
78.69
63.99
46.05
68.78
X
61.41
LION-LIT
16.78
65.21
54,00
43.29
72.78
50.41
LION-D HIPPO)
62.0
88.78
90.12
85.66
90.25
97.28
85.63
LION-S (wl HIPPO)
62.25
88.10
90.35
86.14
91.30
97.99
86.07
C.5
Experimental details for the MLM/GL.UE tasks
Architectures We train the BASE (110M parameters) and LARGE (334M parameters) model families from the original BERT paper [7]. For the LION models, we replace the standard self-attention blocks with LION-LIT/LLONNDDLLON blocks while keeping all hyperparameters the
38

---

Table 8: GLUE finetuning recipes employed in this work. All recipes finetune on RTE, STSB and MRPC from the weights finetuned in MNLI and the rest from the C4-pretrained weights. All recipes sequence length of 128 tokens except BERT24 and Hydra, that use 256. D. AdamW stands for decoupled AdamW.
Dataset
Recipe
Param.
QNLI
QQP
MRPC
COLA
STSB
LR
5- 6-10-5
1-10-3
3 3-10-5
1-10-5
3- 3-10-3
8. 8-10-5
5 5-10-3
3-10
BERTZ4
WD
6-10-6
1 1-10-5
3-10 10-6
1-10-6
3 1-10-
8- 8-10-5
5 100-
3. 1000
[22]
Epochs
10
10
10
Optimizer
D. Adamw
D. Adamw
D. Adamw
D. AdamW
D. AdamW
D. Adamw
D. Adamw
D. AdamW
6-10
5 10
3 10-
10
30
8. 10
10
100
M2-BASB
WD
5 10-0
1 10-3
3 10-0
1. 10-
10-
8 8.10-3
5 10-4
3- 10-6
[10]
Epochs
10
10
10
10
Optimiscr
AdomW
Adamw
D. AdamW
AdumW
AdamW
Adamw
D. AdamW
AdamW
M2-1 ARGGR
WD
LR
5- 5-10-5
5-10-6
5 10
1 1-10-
3- 1005
3- 10-6
5- 10
1- 10-6
3- 3-10-6
3- 3-10-5
8- 8-10-5
8-10-6
5- 5-10-3
1-10-6
3- 10-3
8- 8-10
[10]
Epochs
Optimizee
D. Adami
D. Adamw
D. Adami
AdamW
D. Adami
D. Adamw
D. Adamw
D. AdamW
10-
5- 5-10-5
5- 1005
10-5
5-10-
8- 10-0
10-1
10-3
[2]]
Hydra
Epochs
WD
5- 5-10-6
10-6
7
3 10--
3
10-6
3
3 3-10-6
8-10-6
10
8-10-6
10
3- 3-10-6
Optimizer
AdamW
Adamw
D. Adamw
AdamW
AdamW
Adamw
D. AdamW
Adamw
LR
10-5
10-5
10-5
10-5
10-3
10-5
10-3
10-5
Modified (Ours)
WD
5 1010
1 1-10-
3- 10-5
1. 100
3 3-10--
8 8-10-6
1-10-6
3-10-
Epochs
10
10
Optimizer
D. AdamW
D. AdamW
D. AdamW
AdamW
D. AdamW
D. AdamW
D. AdamW
D. AdamW
same. For LION-LIT, we incorporate LayerNorm [2] after the attention block to enhance stability. For Hydra, we take the default hyperparameters in [21] and increase the number of layers to 45 to match the number of parameters in the LARGE scale. Our implementation is based on the M2
repository [10], i.c.,
com/ HazyHesearcc/mn
Pretraining All our pretraining hyperparameters follow Fu et al. [10]: We employ the C4 dataset [8], a maximum sequence length during pretraining of 128 and masking probability of 0.3 and 0.15 for the training and validation sets respectively. We train our model for 70, 000 steps with a batch size of 4096. We employ the decoupled AdamW optimizer with a learning rate of 8 10- 31 = 0.9, 32 = 0.98, = 10 6 and weight decay 10 10-5 As a scheduler, we perform a linear warm-up for 6% of the training steps and linear decay for the rest of training until reaching 20% of the maximum learning rate.
Our only change in the pretraining hyperparameters is setting the learning rate to 2 - 10-4 for the LARGE model family. In our preliminary experiments, we found that training diverged when using a learning rate of 8 10-4 for BERT-LARGE.
For completeness, we present the results with the BERT pretraining and BERT 24 finetun- ing" recipes available in the M2 repository.
Finetuning For the GLUE finetuning experiments, we employ five different configurations:
BERT24: Available in Izsak et al. [22] and the file this link.
M2-BASE: Available in Fu al. [10], Section C. and the file this link.
M2-LARGE; Available in Fu al. [10], Section C.1 and the file this link.
Hydra: Available in Hwang al. [21], Section D.2 and the file this link.
Modified: Same M2-LARGE but all learning rates are set to 10-5
The recipes are summarized in Appendix C.5. The Modified hyperparameter set was devised as M2-LARGE was found to diverge for BERT-LARGE.
hf cransforerr--reeraii-eeeee bert uncase, yaml
hf attttransforeer- www.:oormmrr-iiieneeeeee geeeeegluueeeeeeee bert base- uncased, yaml
39

---

Table 9: Combining positional embeddings with LION-D and LION-S. Both pretrained models improve in the validation MLM acc. when employing positional embeddings.
Model
Pos. Emh.
MLM Acc.
MNLI
RTE
QQP
QNLI
SST2
STSB
MRPC
COLA
Avg.
66.62
82.85
52.49
89.63
88.43
91.86
85.96
83.94
53.58
78.59
LION-D
/
66.97
83.37
54.08
89.52
88.32
92.35
83.58
79.40
54.53
78.15
67.05
83.17
53.50
89.35
88.89
93.00
37.73
77.87
53.18
72.09
LION-s
/
67.35
83.26
52.42
89.82
88.38
92.58
83.87
79.54
55.25
18.14
C.6 Ablation studies in the MLM/GLUE tasks
Combining positional embeddings with LION. We compare the GLUE performance of LION-D and LION-S when including positional embeddings. We pretrain the BASE models and finetune them with the M2-BASE recipe.
In Table 9 we can observe that adding positional embeddings increased the MLM acc. in around 0.3 percentage points. In the GLUE benchmark, we observe that for LION-D performance degraded in 0.44 percentage points, while for LION-S, performance improved in 6.05 percentage points. We attribute this behavior in GLUE to the dependence on the finetuning recipe.
Recipe selection. In this section, we select the best finetuning recipe for each model family and size. For the BASE models, we test the M2-BASE and Modified recipes. For the LARGE models, we test the M2-LARG-A and Modified recipes.
In Table 10, firstly, we observe that the M2-BASE recipe generally provides a higher GLUE score than the Modified recipe for the BASE models, e-g. 82.25 v.s. 80.26 for the BERT model. Secondly, we observe that for the LARGE model family, the M2-LARGE recipe fails, providing poor performances between 60.96 and 72.41 GLUE points. When reducing the learning rate to 10-8 (Modified recipe), training is more stable and performance reaches between 80.76 and 82.95 GLUE points. We find that small changes in the finetuning recipe have a large effect the performance. Our results in standard recipes show that the LION family of models can obtain a high performance without extensive tuning and closely follow the performance of the BERT family models, at 80.31 82.25 for the BASE model size and 81.58 v.s. 82.95 for the LARGE model size.
Table 10: Recipe selection for the GLUE benchmark.
Model
MLM Acc.
Recipe
MNLI
RTE
QQP
QNLI
SST2
STSB
MRPC
COLA
Avg.
M2-BASE
84.63
64.33
89.99
80
92.51
86.69
89.62
60.42
82.25
BERT
67.70
Mod.
83.09
58.27
89.35
89.88
92.16
86.56
87.78
55.02
80.26
M2-BASE
82.50
63.47
89.72
89.27
91.74
87.18
89.37
LION-LIT
65.47
Mod.
80.88
54,95
88.80
88.83
91.32
85.42
87.07
46.98
78.03
M2-BASE
82.85
52.49
89.63
88. 13.43
91.86
85,96
83,
53.58
78.59
LION-D
66.62
Mod.
80.52
52.85
88.93
88.36
91.55
82.05
84.48
49.13
77.23
M2-BASE
83.17
53.50
89.35
88.88
93.00
37.73
77.87
$3.18
LION-S
67.05
Mod.
78.14
56.39
88.68
88.52
92.39
51.22
77.60
49.75
72.84
M2-LARGE
84.97
69.10
31.59
49.15
91.93
53.61
87.87
$1.16
64.92
BERTLARGE
69.88
Mod.
85.68
67.44
89.90
91.89
93.04
88.63
90.89
56.14
82.95
HydRIARRGE
71.18
Hydra
Mod.
84.39
84.24
60.44
59.42
90.38
91.31
89.73
91.70
93.44
87.19
88.21
88.57
88.99
59.46
81.77
M2-LARGE
83,20
54,51
89.08
84,90
90.44
68.57
85.25
23.35
72.41
LION-LIT LARGE
67.11
Mod.
83.73
57.18
89.85
89.93
91.86
88.02
90.18
55.36
80.76
M2-LARGE
52.85
41.48
53.67
91.13
36.87
LION-D LARGE
68.64
Mod.
83,82
60.72
89.72
89.79
92.93
87.29
89.66
56.83
81.34
M2-LARGE
83.71
50.04
38.81 81
53.98
91.59
36.98
82, 29
50.27
60.96
LION-S LARGE
69.16
Mod.
84.38
57.69
89.57
90.30
92.93
87.68
90.57
59.54
81.58
C.7 Additional experimental results for the MLMIGLUE tasks
In this section, we present our bidirectional MLM results in the BASE scale using the BERT pretraining recipe described in Appendix C.5 and BERT24 [22] finetuning recipes (Table 11),
40

---

OOM
Enoe 1111-
BERT
Hydra
70
LION-S
GPU Memory (GB)
LION
60
60
Trinnin Sequerss Length
320 38 sectracy
BERT
LION-LIT
Accuracy
50
40
Bette
40
20
30
20
Positional Encodings
01
200
400
600
800
1000
10000
Sequence Length
Context Length
Figure 9: (Left) MLM Acc. for different sequence lengths. LION-S is able generalize larger sequence lengths and does not require positional encoddings. (Right) GPU Memory for different sequence lengths. Results for the LARGE (334M) scale. The memory employed by LION and Hydra scales linearly with the sequence length, being able to process a sequence length of ~ 16.000 tokens with less than 20GB. Contrarily, the memory employed by BERT scales quadratically, going of memory (80GB) at sequence length of ~ 12.000 tokens.
present the per-task GLUE results omitted in the main text (Table 11) and present the length scaling capabilities of the LION-S model (Figure 9).
Table 11: C4 Masked Language Modeling and GLUE results for the BASE scale (110M) and LARGE scale (334M), For each column (dataset), the best and the second best results are highlighted with bold and underline respectively.
Scale
Model
MLM Acc.
MNLI
RTE
QQP
QNLI
SS12
STSB
MRPC
COLA
Avg.
BERT
67.23
84.26
59.21
89.87
90.24
92.35
88.12
90.24
56.76
81.38
BASE
Hydra"
69.10
84.50
57.20
91.30
90.00
93.50
91.20
88.90 28.90
77.50
84,30
LION-LIT
65.08
82.37
55.81
89.49
89.57
91.74
86.27
88.25
44.46
78.50
LION-D
66.62
82.85
52.49
89.63
88.43
91.86
85.96
83.94
53.58
78.59
LION-S
66.19
82.50
57.47
89.38
87.88
92.70
82.42
82. 26.46
53.39
78.40
BERT
69.88
85.68
67.44
89.90
91.89
93.04
90.89
56.14
82.95
LARGE
Hydra
71.18
84.39
59.42
90.38
93.43
88.57
59.46
81.77
LION-LIT
67.11
83.73
57.18
89.85
89.93
91.86
88.02
90.18
55.36
80.76
LION-D
68.64
83.82
60.72
89.72
89.79
92.93
87.29
89.66 66
56.83
81.34
LION-8
69.16
84.38
57.69
89.57
90.30
92.93
87.68
90.57
59.54
81.58
Results extracted from the original paper [21].
C.8 ImageNet classification results for Tiny scale
In Table 12, present the image classification results of LION models the tiny scale models an compare them against the baseline models. Results indicate that the high training speed and competitive performance of LION models is also applicable in the tiny scaled models.
C.9 Inference time comparison for image classification tasks
In Figure 10, we compare the baseline models on the image classification task. lower resolutions, all models exhibit similar inference times, with Vim and Hydra being slightly slower than ViT and LION-D chunking. However, at higher resolutions, vision SSMs (Vim and Hydra) become faster. This trend arises because vision SSMs leverage specialized kernels, whereas VIT and LION-D rely plain Python implementations, leading to increasing overhead as resolution grows.
C.10 Ablation studies with image classification
Choice of X; values. In this section, we study the properties of the selectivity parameter on CIFAR- dataset. We tested, three cases: (1) fixed mask with scalars a; a², (ti) vector, input- dependent E IRd (cf., Appendix B.5) and iii) input dependent scalar a, R. The results, presented
41

---

1000
VIT
Hydra
Table 12: Image classification task in Tiny scale. Top-1 accuracy on the validation data. repre- sents changes in patch orders. Best and second best results are in bold and underline.
Inference Time (s)
800
LION-D
VIM
600
400
Imagenet Top-1 Acc.
Train time
Model
Param
200
ViT
5M
70.2
x1
DeiT
5M
72.2
224
512
738
1024
1248
Vim
7M
76.1
x9.48
Resolution
LION-LIT
5M
68.9
x0.69
LION-D
5M
72.4
x1.48
LION-D
5M
74.2
x1.73
Figure 10: Inference times comparison. The inference time LION-D with chunking, Vim, Hydra and ViT models presented for different resolutions.
LION-S
5M
72.4
x2.05
LION-S
5M
73.5
x3.83
in Table 13, show that while the input dependency is beneficial, the expansion of a, is not necessary for image tasks. As a result, we employ option three in all image classification tasks, and the end model is called LION-S.
Table 13: Ablation studies on image classification. Additional ablations with CIFARIO0 dataset to determine the size and input dependency of the selectivity parameter of the model LION-S.
Models
Top-1 I Acc.
Fixed mask a; a' Vector a, E IRdd
75.66
67.55
Scalar, input dependent a; E IR (LION-S)
77.56
Understanding the power of non-linearity, softmax, and positional embeddings. In Table 14, we present additional ablations certain design elements of Vision Transformer. We perform these experiments on CIFAR- 100 data using the same hyperparameters with LION-S. We have observed that either nonlinearity or softmax is essential for the model to converge with nice accuracy. Though positional embedding boosts the accuracy, a mask can easily replace it.
Table 14: Ablation studies on image classification. Additional ablations with the CIFAR- 100 dataset to understand the contribution of softmax, nonlincarities model is presented. Soft., PosEmb and NonLin iin expresses softmax, positional embedding, and non-linearity have been applied. X means the model did not converge. The symbol denotes the adaptation of recurrent models that achieve equivalence attention during training while utilizing recurrence during inference, established by our theorem.
Models
Top-1 I Acc.
[ Soft. PosEmb NonLin
73.88
2 Soft. PosEmb (ViT-T)
77.33
3 Soft. NonLin
4
Soft.
73.15
5 PosEmb + Non. (LION-LIT)
73.61
6 PosEmb
68.54
7 NonLin
65.28
8 Base
Non.Lin Mask (LION-S)
77.56
42

---

C.11 Hyperparameters for Training Image Classifiers
All experiments were conducted on single machine for CIFAR- and multiple machines for ImageNet, using NVIDIA A100 SXM4 80GB GPUs. For LION models and Hydra, the VIT architec- ture serves as the main structure, with Hydra following the Hydra training recipe and other models following the VIT recipe. The training and evaluation codes are adapted from [46] and [50].
C.12
Calculation of Number of FLOPS
Below we present theoretical number of FLOPS used in the attention of vision transformers and LION-S during inference where L, is the resolution/context length and D is the hidden dimension. Results show that while transformer has O + IDD2) LION-S has (IL22). Note that in this calculation, the exponentials and other nonlinearities considered FLOP whereas in reality, the Softmax introduces additional complexities. The same calculations should also apply to other bi-directional models.
The number of FLOPs in the one head of the one layer attention for vision transformer:
Calculating Q, K, V: 6LD
Attention A QK :2LDD
Softmax (assuming FLOP for exp): 2L
Calculating Y: 2L2D
TOTAL: L(6D2 + 4LD + 2L)
The number of FLOPs in the attention module for LION:
Calculating K, V, A: 6LD 2LD,
For each token in one forward/backward recurrence:
Updating F//B 3D2
Updating F/B. 2D
Calculating C: //8; 4D +
Calculating y F/B; 2D2 + 47 +1
Total: 5D2 + 10D +
L forward + backward recurrences: 2L(5D² + 10D +
Calculating Y: 2L(D+ +1
TOTAL: L(16D² + 24D + 7)
C.13 Distillation Results of LION-S
We have also used the same recipe from DeiT distillation [47] and distilled the RegNet network into LION-S. We observed that the distillation outperforms the original ViT-Tiny on the ImageNet dataset. The results are shown in the table below:
C.14 Different scans in vision Linear Transformers
When processing images, both the spatial relationships among neighboring pixels and their positions critical as the pixel values themselves. Positional embeddings provide way to incorporate these spatial relationships. The common approach in Transformers involves flattening, also called scanning, the image row-by-row, from left to right and top bottom illustrated in the left panel of Figure 1..
Linear Transformers and SSMs such XLSTM [1], VMamba [29], and Vim [53], often traverse (or scan) the image patches using a the same flattening pattern. To enhance spatial understanding, scans usually applied in both vertical and horizontal directions [l, 29]. These scans act separate
43

---

Figure 11: Reordering of patches. Left is the naive approach flatten images, also used in LION-S. Center and right figures are the new approaches applied in LION-S* to consider further spatial information.
recurrences within the model, each capturing different spatial patterns within the image through their respective decay factors.
However, we argue that this method of flattening is suboptimal and can be enhanced to include additional contextual information. To address this, we propose a new scanning scheme for patch values. In the attention module, the patch values are traversed following the patterns depicted in the center and right panels of Figure 11. Forward and backward passes are then executed based on this new ordering, adhering to established procedures. The outputs from these two passes are subsequently averaged to generate the final result.
We indicate this method with throughout the paper. This approach demonstrated a notable im- provement in accuracy for image classification tasks while maintaining the efficiency and flexibility inherent to the method. A similar concept has been previously explored in Vision-LSTM []].
C.15 Ablation studies on importance of bi-directionality on image classification
To highlight the importance of bi-direcctonnlltyy and demonstrate the versatility of the LION framework, we conducted additional experiments examining the processing directions of the blocks. We evaluated four settings: (i) all blocks process patches in the forward direction only (Forward), (ii) all blocks process patches in the backward direction only (Backward), (iii) odd-numbered blocks process backward direction
patches in the forward direction while wvennnumered blocks process them (Forward- (Forward-Backward) Backward), and (iv) all blocks process patches in both directions (Bi-directional). The results reveal that incorporating both directions improves performance by approximately 4%, while full bi-directionality achieves significant boost of up to 10%.
Table 16: Results for LION-S and LION-s with different directional settings CIFAR-100. In- corporating both directions improves performance by approximately 4%. while full bi-directionality achieves a significant boost of up to 10%.
Table 15: Distillation results of LION-S.
Models
Top-1 Acc.
Model
Top-1 Acc.
LION-S
67.95
LION-S (Forward)
71.08
VIT-Tiny
70.23
LION-S (Backward)
69.61
LION-S (Distilled)
70.44
LION-S
73.93
LION-S (Bi-directional)
77.56
LION-S (Forward)
70.24
LION-S (Backward)
70.42
L (Bi-directional)
80.07
44

---

Leurrll Dew pocuddes LARCH
Loccel Coop potrrte LION DASSS
Eouoed Ocar pascemnii LIN BASE Vaice
131
1/-
0.5
04
Layer
(a) Chunkwise Parallel LION-LIT
(b) Chunkwise Parallel LION-D
(c) Chunkwise Parallel LION-S
Figure 12: LION-D Decay factor distribution.
C.16 LION-D Decay factor distribution
visualize the distribution decay factors in LION-D across layers for both image classification and MLM tasks. Interestingly, the model learns a diverse range of decay values, highlighting their importance across different heads and layers.
C.17 Generation of the Mask
Below we present the Python code used for the creation of the bidirectional mask M as described in previous sections.
# #caption-Code for generation of the selective bidirectional mask of \lion
def create_ tensor (tensor):
cumsum torch. (torch cumsum (tensor, dim=-1)) A = torch matmul (cump.. unsqueeze 1/ ( comp.unsqueeze (-1) (- return torch, (A)
8 def Mask_ sel.eeeceeeeeeee).
vec_shape vec. shape
A_for treate_mttixx_fe tensor (vec)
A_back create_matrix_frot (elsor
return A_for A_back torch, eye (A_for shape [1]]
def Mask_Decay (a L):
idx torch. arange (L, .eviice=_i.device) torch,
I, torch. meshgrid (idx, idx, indexing-'1j
17
E atorch.abs((1) ((I-)) fllat () view (1,1,L,L)
x torch, sigmoid (a_i) view:..iiiiw.......ei 11, ++E
return M
# #caption=Code for generation of the partial selective bidirectional mask of Alion for chunking,
def Partial_Mask selective Partisl_Nakk_alleetive(ee.))
B,H,L vec. A_for vec. shape
(vecl chunk_index chunk_len)
3
A_back =
chunk_indes index chunk_len
torch diag embed (torch ones ( index_index= chunk len)
offset chunk_inden chunk_les [....
return A_for + A_back to (A_ device)
45

---

D Future Direction & Limitation
While LION maps many Linear Transformers the bidirectional setting, we empirically focused on three examples. LION does not rely on additional design choices (e.g., gating in Mamba), but could incorporate them to further boost performance and scale larger models, potentially suppressing softmax Transformers in multiple domains and tasks.
46


Reviews:
[{"review_id": "qKRSWPMHL6", "reviewer": "Reviewer_GQMH", "summary": "The authors introduce LION, a bidirectional extension of linear transformers. To do so, they go beyond the causality preserving formulation of the attention masking matrix $M$ to adapt it to the setting where entire sequences are available at both training and inference. The decay terms in $M$ allow for an elegant triangular decomposition that lends itself to efficient computation. The attention framework formulation is then translated into its recurrent equivalent (similar idea to the unidirectional causal transformer-to-RNN). Three variations of the approach can be trained, depending on the decay terms selection scheme. Extensive experiments on image classification, masked language modeling and a LRA task, present tradeoffs between training and inference time, performance, GPU and memory costs, etc.", "questions": "See Weaknesses, I have incorporated questions in points brought up to avoid repetition.", "limitations": "The authors fail to properly address limitations, notably be it on the architectures that fit in the Lion framework (line 87 mention on TC$^0$ class, what are the models outside of that? Would something like a Delta Net fail to be represented [\"Parallelizing Linear Transformers with the Delta Rule over Sequence Length\", Yang et al. NeurIPS 2024] and variants])", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "## Strengths\n\n- **Clarity:** Though dizzying at first, the colors and blocks all over end up being useful for the visualization of the ideas behind the math. Anyhow, the authors clearly state the goal of the work and its scope.\n- **Significance:** The problem raised is of interest especially to practitioners, as it allows first an introduction of a framework for attention/SSM duality in bidirectional models as well as a better understanding of advantages and tradeoffs in this regime.\n- **Technical quality:** Though I have not checked all the derivations, the work is sound. The choice of decay terms to match popular architectures of different inspirations does show the unifying setup the framework claims to be (though it opens questions on whether other design choices could in turn go beyond the initial inspirations).\n- **Rigor:** The work is quite comprehensive and self contained. I commend the authors for the effort of compiling an extensive appendix to complement most details that do not fit in the main text.\n\n---\n\n## Weaknesses\n\n- **Design of decay terms:** Following up on the last strength. Given that performance is competitive and albeit better than vanilla ViT but not convincingly superior to other baselines, the question of extending the design search arises. I also find it surprising that using fixed decay terms seems to perform no worse than parametrizing and learning the decay terms (D vs S). Do the authors have an understanding or intuition about these differences (or lack of) in performance.\n\n- **Experiment takeaways:** It is not immediately clear what we should be taking away from the experiments. More specifically:\n1) The first experiment, figure by table 2: if the argument is faster training speed, then Lion-LIT is indeed faster than DeiT, yet performs significantly worse. The latter is then the second fastest model to train while still outperforming D and S (base at least, matching on small) The case for Lion is mitigated here. \n2) Figure 1 shows tradeoffs between attention, RNN, and chunking approaches for inference. This is a generic tradeoff I suppose and not related to the bidirectional formulation as it would similarly hold in other cases (correct me if I'm wrong). I fail to see why it is relevant to the findings of the paper as a result and not as a motivation, given that the whole point of having both attention and RNN representations is to be able to move from one to the other where advantageous. \n3) All in all I fail to gather a clear conclusion from the experiments. Is it better performing, faster, more efficient than the state of the art? Any or none? Is the contribution mainly in the novel formulation?", "was_revised": true, "final_justification": "The authors have clarified the presentation points raised which ensure the claims and main contributions of the paper are clear. \n\nThe discussion on the extension of the bidirectional rnn formulation is valuable in showing that the approach is valid beyond the diagonal case, without suffering from additional compute constraints beyond those that arise with full attention models.", "rebuttal": "We thank the reviewer for their detailed feedback and for appreciating our work, particularly noting its **rigor, technical quality, and clarity**. Below, we address the main concerns raised.\n\n-------\n\n### 1) Design of decay terms\n\nWe like to clarify that, both the selective and fixed decay factors are learned during training: the fixed decay uses `non-input dependent` $\\lambda = \\sigma(a)$ with a learned scalar $a$, while the selective decay is `input-dependent`, defined as $\\lambda_i = \\sigma(Wx_i)$ with learned weights $W$.\n\nRegarding performance, we find the choice of decay factor to be **task-dependent**. The selective variant (LION-S) generally outperforms the fixed variant (LION-D), by **0.25% on masked language modeling** (Table 3) and **0.44% on LRA** (Table 4),likely due to its greater representational power and inductive bias. In contrast, LION-D performs slightly better in image classification, where fixed decay appears sufficient to capture relative token distances in signal and image modalities [1].\n\nAs the reviewer noted, different linear transformers use different decay parameterizations, as shown in **Table 1** of our paper and **Table 2** of DeltaNet [2]. For LION, we selected three diverse decay forms to demonstrate the generality of our framework. We believe each decay choice introduces different inductive biases, making them suitable for different tasks, for example, DeltaNet performs well on retrieval. We believe that evaluating different decay choices is a valuable direction for future work.\n\n-------\n\n### 2) Experiment takeaways\n\n**DeiT is second fastest while outperforming LION D and S**\n\nThe main takeaway from Table 2 is that LION, and more broadly, linear transformers are competitive with softmax-based transformers like DeiT in both training speed and accuracy. However, they offer a significant advantage in **inference efficiency**. As shown in Theorem 3.1, LION can be expressed in RNN form, enabling linear-time inference ($O(L)$) compared to the quadratic complexity ($O(L^2)$) of softmax transformers. This is illustrated in Figure 2 and 8.\n\nWe agree that Table 2 alone may not fully capture LION's advantage over softmax transformers. Since we aim to highlight three key aspects, 1) accuracy, 2) memory usage, and 3) training time, presenting all of them in a single plot would reduce clarity. Therefore, we separated them: **Table 2** focuses on performance and training speed, while **Figure 2** emphasizes inference memory efficiency and scalability. Together, they demonstrate that LION **combines the strengths of transformers in speed** and accuracy with the **inference efficiency of RNNs.**\n\n**Chunking Tradeoff is Generic; Figure 1 Serves as Motivation Rather Than a Result**\n\nWe agree with the reviewer that the tradeoff between RNNs, chunking, and attention is generic, which also applies to bidirectional models, as noted in the background section of our paper. Figure 1 is primarily intended to motivate the value of supporting multiple representations depending on resource constraints. Thanks to the reviewer’s suggestion, we will move this figure to Section 3, where it better fits as a motivation.\n\n\n**Final Conclusion from Experiment Section**\n\nOur main contribution, as noted by the reviewers (particularly in the second point), is the **novel formulation of LION** as the  framework for extending linear transformers to the bidirectional setting while supporting and analyzing tradeoffs across different regimes.\n\nOur experiments shows that LION works across multiple modalities. Key takeaways of experiments include:\n\n- LION achieves training speed and accuracy close to softmax transformers (Table 1), *while being significantly more efficient at inference*.\n\n- LION matches or exceeds the inference efficiency of SSMs (Figure 2), *while training much faster*.\n\n------\n\n### 3) Discussion on DeltaNet\n\nWe appreciate the reviewer's sharp observation on this point. Our theorem is specifically focused on linear transformers with **diagonal decay**, and models outside this class **generally cannot be expressed in the full attention formulation**. However, for the special case of **DeltaNet**, our theorem supports an equivalent **bidirectional RNN** formulation with LION-style correction terms to avoid double counting, as follows:\n\n$S_i^{F/B} = S_{i-1}^{F/B}(I - \\beta_i k_i k_i^\\top) + \\beta_i k_i v_i^\\top, \\quad y_i^{F/B} = q_i^\\top S_i^{F/B} - \\frac{1}{2} q_i^\\top k_iv_i^\\top$\n\nThe chunkwise parallel form of DeltaNet (Eqs. 8–11 of the DeltaNet paper) can also be extended to the bidirectional setting by removing causal masks and applying Theorem 3.2 for the chunkwise form of LION.\n\nHowever, a key bottleneck in DeltaNet arises in its **full attention formulation**. As described in the original paper (page 6, section “Fully Parallel Form for DeltaNet”), the attention is computed as:\n\n$A = QK^\\top T$\n\nComputing the matrix $T$ (defined in equation 10 of DeltaNet) is costly, which directly citing from paper:  **\"$T$ requires a matrix inverse (Eq. 10), which scales cubically with sequence length without further algorithmic changes.  Due to the above we avoid using the fully parallel form\nfor training DeltaNet\"**.  The same limitation applies in the bidirectional case, where the inverse undermines the fast-training advantage of LION.\n\nWe thank the reviewer for raising this point and will include the discussion, along with a note on models outside $\\text{TC}^0$ with non-diagonal decay, in the limitations section.\n\n-----\n\n### References\n\n[1] Mamba: Linear-Time Sequence Modeling with Selective State Spaces Albert Gu, Tri Dao\n\n[2] Parallelizing Linear Transformers with the Delta Rule over Sequence Length\", Yang et al. NeurIPS 2024]"}, {"review_id": "zjMlBdOH8T", "reviewer": "Reviewer_tuMu", "summary": "The authors introduce a new technique to perform, for the first time, bidirectional attention with Linear Transformers. They create a novel framework (LION) for bidirectional sequence modeling, which supports three types of representations. These different representations present advantages on different fronts: training speed, memory-efficient inference and a balanced mode between the two. Then, LION models using full attention are trained on computer vision and NLP tasks, rivaling the performance of other SSM techniques and vanilla Transformers, while having considerably faster training time and inference speeds.", "questions": "- How does LION perform on tasks requiring long-range dependencies (more than 16k tokens)?\n- Why is the reported training time of Vim so high compared to ViT (Table 2)? Although they don't report numbers, the authors of Vim affirm that its training is more efficient than ViT's training.\n- Why is also the reported training time of Hydra so high compared to the BERT baseline (Table 3)?\n- Have you explored custom kernel implementations that could further improve speed and improve memory usage? Have\n- Why does Figure 8 (left) not include the Hydra baseline, which is included on the right?", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "# Strengths\n- **Novel theoretical contribution**: the paper addresses an important gap in the literature by introducing the bilinear formulation to the linear family of transformers. The theoretical analysis is rigorous, and the authors apply their framework to various existing linear transformers and SSMs.\n- **Generality**: the framework can extend multiple existing Linear Transformers (RetNet, xLSTM, ...).\n- **Comprehensive experiments**: Evaluation across a diverse set of tasks (image classification, masked language modeling, LRA) and model sizes, with a thorough ablation studies examining various design choices.\n- **Practical efficiency gains**: the experiments demonstrate significant training speed-ups while maintaining a competitive performance. Moreover, the ability to change between full-attention (fast training), RNN (efficient inference), and chunking (memory-speed tradeoff) is valuable for practical deployment.\n- **Clear presentation**: the paper is overall well-written, with clear explanations and helpful visual cues to better understand the notation.\n# Weaknesses\n- **Limited scalability analysis**: while the authors train models of up to 334M parameters, they don't explore larger scales where the efficiency gains might be more critical. Moreover, there is no analysis on how the model performs with very long sequences (>16k tokens), where these models shine against the softmax attention of the transformer.\n- **Comparison fairness**: some baselines (Vim, Hydra) use specialized kernels while LION is implemented in PyTorch without any compilation. This makes speed and memory usage somewhat unfair and understates the potential of optimized implementations.\n- **Experimental details**: some important details are relegated to the appendices.\n- **Performance-efficiency tradeoff**: while LION achieves significant efficiency gains, it consistently underperforms softmax attention.\n- **Scalability unknowns**: the experiments only go up to 334M parameters, which does not prove the scalability of this method to billion-parameter models, widely used in current applications.\n\n\n### Minor issues\n- Line 34: \"modeling, namely\"\n- Should the mask in Eq (4) have a different name to differentiate it with the causal mask?", "was_revised": true, "final_justification": "The authors have successfully addressed my questions and suggestions.\n\nAlthough LION does not achieve state of the art results, this manuscript provides a first step into bidirectional linear attention formulations, which I believe introduces an interesting area of future research on efficient architectures.", "rebuttal": "We thank the reviewer for finding **LION framework novel, general and effective in practical scenrios** bellow we will address the questions and concerns mentioned in the review.\n\n--------\n\n### 1) Scalability of LION to longer sequences and models sizes\n\n\nDue to computational constraints and alignment with prior work [1,2,3], our experiments focus on models with ~100–300M parameters. We expect LION to scale well on non-causal tasks with larger models. Similarly, while we couldn’t train on longer sequences (>16K), **LION’s strong LRA performance** suggests effective generalization to longer inputs.\n\n--------\n\n\n### 2) Speed comparison fairness\n\nDespite using only PyTorch, **LION still outperforms Hydra and Mamba in speed**, highlighting the efficiency of its formulation. We chose PyTorch to keep LION's implementation simple and easy to integrate. We agree with the reviewer that custom kernels could offer additional speedups and represent a promising direction for future work. LION already achieves faster training even with a pure PyTorch implementation.\n\n--------\n\n\n### 3) Experimental details\n\nThank you for the pointer. If there is any specific part of the appendix you would like to see included in the main body, please let us know and we would be happy to incorporate it. Due to space constraints, we focused on key content in the main body. For the final submission, we plan to bring in additional details from the appendix, as the *NeurIPS 2025 Call for Papers* allows **one extra page**.\n\n--------\n\n\n### 4) Performance-efficiency tradeoff\n\nLION achieves performance close to Transformers on most tasks and even outperforms them at smaller scales, i.e., in image classification with Tiny models (Table 12, Appendix). It also solves LRA tasks from scratch, where Transformers are not able to.\n\nIn general SSMs and linear transformers may slightly lag behind Transformers in accuracy, but their inference efficiency and flexibility to switch between attention, RNN, and chunkwise forms make them particularly valuable in practice, this is the core motivation behind LION. We believe that more fine-tuned architectural modifications, such as convolutions over $Q, K$ and additional gates, as used in DeltaNet (ablation of Table 1) [6], could further enhance LION's performance. However, we intentionally used a **simple backbone to isolate and highlight the effect of LION’s attention and sequence mixing block to reflects more its generality across different decays.**\n\n-----\n\n### 5) Minor Issues\n\nWe thank the reviewer for pointing out the minor issues related to typos and sentence clarity. We will address these to improve the overall readability in the final version.\n\n-----\n\n### 6) Questions\n\n- **1) Performance on longer sequences** \nAs mentioned above, we believe LION can scale to longer sequences and larger model sizes; however, due to resource constraints, we were unable to explore this in our current work.\n\n- **2) Long training times of Vim and Hydra**\nThanks for the sharp observation. The slow training time of Vim, even with custom kernels, is due to the use of parallel scan in Mamba’s [4] training algorithm, which is less efficient on GPUs compared to attention. This issue has been addressed in Mamba2 [5] by using a scalar decay for $A$ and adopting a chunkwise parallel training algorithm (aka SSD), significantly improving training speed (which is also shown in details in Mamba2 blogpost). Hydra, built on Mamba2, benefits from this and trains faster. However, both still **rely on recurrent processing over the state** (also stated in related work of LION and DeltaNet [6]) and are not as fast as attention, which is a key motivation behind LION, which aims to match the training throughput of softmax attention.\n\n- **3) Custom kernel implementation of LION** \nAs we addressed this concern in details in concern 4 above, LION is implemented solely in PyTorch. Our goal is to keep LION simple and fast to use within standard PyTorch workflows.\n\n- **4)  Figure 8 (left) does not include the Hydra** \nThanks for the pointer we will add Hydra to Figure 8 (left) and it also can extrapolate alike LION.\n\n\n-----\n\n###  References\n\n[1] Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers, Sukjun Hwang, Aakash Lahoti, Ratish Puduppully, Tri Dao, Albert Gu\n\n[2] MosaicBERT: A Bidirectional Encoder Optimized for Fast Pretraining\nJacob Portes, Alex Trott, Sam Havens, Daniel King, Abhinav Venigalla, Moin Nadeem, Nikhil Sardana, Daya Khudia, Jonathan Frankle\n\n[3] Monarch Mixer: A Simple Sub-Quadratic GEMM-Based Architecture\nDaniel Y. Fu, Simran Arora, Jessica Grogan, Isys Johnson, Sabri Eyuboglu, Armin W. Thomas, Benjamin Spector, Michael Poli, Atri Rudra, Christopher Ré\n\n[4] Mamba: Linear-Time Sequence Modeling with Selective State Spaces\nAlbert Gu, Tri Dao\n\n[5] Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality, Tri Dao, Albert Gu\n\n[6] Parallelizing Linear Transformers with the Delta Rule over Sequence Length\", Yang et al. NeurIPS 2024"}, {"review_id": "DzgAyLc1VA", "reviewer": "Reviewer_1Uqp", "summary": "The authors introduce LION, a framework to extend Linear Transformers from causal to bi-directional. LION supports three forms: recurrent for efficiency and low latency, full attention for parallelization, and chunkwise for a balance between memory and parallelization. Additionally, LION incorporates three types of decay factors: LION-LIT, which uses no decay; LION-D, featuring a learnable, non-selective decay; and LION-S, which employs a selective, input-dependent decay. The authors implemented a bi-directional Linear Transformer using the LION framework and evaluated its performance on various tasks, including vision (ImageNet-1k) and natural language processing (C4, GLUE, and LRA). While LION may not surpass the vanilla Softmax Transformer in speed or effectiveness, it presents a compelling alternative to bi-directional extensions of State Space Models like Hydra.", "questions": "* Can the authors include a more challenging and modern benchmark such as MTEB? \n* Can the author replicate the experiments with a more modern setup, notably longer sequences and more training tokens?", "limitations": "Yes.", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Major Strengths:\n- The bidirectional setting has been less studied compared to the causal setting in recent years. However, many applications are inherently non-causal and benefit from the bidirectional setting. For instance, document retrieval in NLP, protein design in biology, and image classification in vision. Therefore, I believe the paper is timely and valuable.\n- While I haven’t closely followed the literature on the bidirectional extension of sub-quadratic Transformers and SSMs, I haven’t seen any previous work that explores this for linear Transformers.\n- The paper is well-written, and I appreciate the authors’ efforts in creating nice visualizations. \n\nMajor weakness:\n- The NLP experimental setup is outdated, resulting in relatively low performance. For instance, the BERT baseline achieves 82.95% on GLUE, while Hydra achieves 81.77%. In comparison, RoBERTa 340M reported 88.5% [1], and the recent NeoBERT 250M reported 90% [3]. Interestingly, the authors of Hydra 112M reported 84.3%, more than 2.5% higher [2]. I believe the discrepancy in performance is primarily due to the shorter training duration (36B tokens) and the short sequences (128 tokens instead of the standard 512 tokens). Additionally, GLUE has long been recognized as outdated and not reflective of the performance on real use cases. Unfortunately, given the lower performance, it is difficult to draw conclusions on the performance of LION compared to the baselines, as its models may be constrained by the training setup.\n\nMinor weakness:\n- Line 51, the claim that LION is up to 10 times faster than SSMs while matching the performance of Sotmax Transformers and SSMs is too strong. Except for Vim, LION is only about 2 times faster than other SSMs, and it performs almost consistently slightly worse than the vanilla Transformer and the best SSMs.\n- Figure 8 studies the generalization to longer sequences of Transformer and LION. While the original position encoding did not generalize well, as illustrated, most models now rely on rotary positional encoding (RoPE), and sometimes on expansion techniques such as YaRN, which have been shown to generalize much better [3].\n\nMinor suggestions:\n- I suggest the authors explicitly state whether they are using FlashAttention, which their code suggests.\n- I suggest the authors move Figure 5 in the main content of the paper, as it illustrates very well their approach.\n- FYI, LION also refers to an optimizer proposed by Google Brain in 2023. \n \n[1] Liu, Y., Ott, M., Goyal, N., Du, J., Joshi, M., Chen, D., Levy, O., Lewis, M., Zettlemoyer, L., & Stoyanov, V. (2019). RoBERTa: A Robustly Optimized BERT Pretraining Approach. ArXiv, abs/1907.11692.\n\n[2] Hwang, S., Lahoti, A., Dao, T., & Gu, A. (2024). Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers. ArXiv, abs/2407.09941.\n\n[3] Breton, L.L., Fournier, Q., Mezouar, M.E., Morris, J.X., & Chandar, S. (2025). NeoBERT: A Next-Generation BERT. Trans. Mach. Learn. Res., 2025.", "was_revised": true, "final_justification": "I will maintain my score because the empirical evidence is insufficient. \n\nThe NLP experimental setup is outdated, resulting in relatively low performance. For instance, the BERT baseline achieves 82.95% on GLUE, while Hydra achieves 81.77%. In comparison, RoBERTa 340M reported 88.5% [1], and the recent NeoBERT 250M reported 90% [3]. The proposed method might be effective at smaller scales, with shorter sequences, or fewer gradient steps but perform poorly at scale. The experiment provided in the rebuttal suggests that this might be the case, as Hydra improves more with longer sequences than LION.", "rebuttal": "We highly thank the reviewer for considering our paper **timely and valuable, the first to study linear transformers for bi-driectional setting**, and for describing it as well-written. Below, we address the **only major weakness** and the minor weaknesses raised by the reviewer.\n\n------\n\n### Major Weakness on NLP Experiment\n\nWe thank the reviewer for their feedback which helped us improve our results. Our setup follows **M2-BERT** (released in 2024 and relatively recent) for two key reasons:\n\n1. It also serves as the benchmark and training recipe used in **Hydra** and previous works on bi-directional sequence modeling [1,2,3].\n2. It offers a **more GPU-friendly training setup**. Due to resource constraints, we could not adopt heavier training configurations such as **NeoBERT** (2.8TB dataset and 4k sequence length) or **RoBERTa** (500K steps, 8K batch size), and thus followed the lighter M2-BERT recipe.\n\nWe chose the GLUE benchmark as a practical and commonly used evaluation for models at this scale [1,2]. While larger benchmarks may offer additional insights, GLUE remains appropriate given our scope and resources.\n\nWe thank the reviewer for the suggestion and have incorporated a longer sequence length into our experiments. **We increased the training length to 512 for base scale, as suggested**. Results are shown below:\n\n\n|Model|GLUE-Score|\n-|-\nHydra|**84.3**\nLION_LIT|81.7\nLION_D|82.5\nLION-S|*83.0*\n\nThe above results were obtained on sequences of length **512** in a single run, without any hyperparameter tuning, due to time constraints during the rebuttal period and limited GPU resources. **All LION variants showed improved performance in GLUE score** compared to the shorter **128-length** setting reported in **Table 3** of the paper. This demonstrates that **LION has strong potential for handling longer context lengths**, even without additional tuning.\n\n\n-----\n\n### Minor Weaknesses\n\n- **\"LION is up to 10× faster than SSMs\" is too strong:** We will revise this statement to clarify that LION is 10× faster than Vision Mamba and 2× faster than Hydra, ensuring the claim is precise and accurately reflects the comparisons.\n- **New models relying on RoPE generalize better in context extrapolation:**\nContext extrapolation is a general advantage of linear transformers, and LION inherits this property as well. We include Figure 8 in the Appendix to illustrate this. While we fully agree with the reviewer that modern positional embeddings like RoPE, YaRN, and ALiBi also support extrapolation, the focus of our study is not to compare positional encoding methods. Therefore, in Figure 8 we simply aimed to show that LION can extrapolate beyond trained sequence length, due to its use of  decays instead of fixed positional embeddings.\n\n\n--------\n\n### Minor Suggestions\n\n- **FlashAttention usage:** Thank you for the pointer. We will clarify in the final version that FlashAttention is used for the softmax Transformer baseline in our experiments.\n- **Bringing Figure 5 to the main content:** We appreciate the suggestion and will move Figure 5 to the main paper in the final submission.\n- **LION also refers to an optimizer:** Thanks for the pointer. We are aware of the potential confusion. However, naming models after animals is common in the linear transformer community, and there is no prior model named LION in this context. Additionally, since the domains differ (sequence models vs. optimizers), we believe confusion is unlikely.\n\n-------\n\n###  Questions\n\n- **Including MTEB benchmark:** Due to time constraints and our limited compute resources, we were unable to support this benchmark. However, we agree it is a valuable direction for future work.\n\n- **Increasing sequence length in NLP:** Thank you for the suggestion. We have ablated the 512-token sequence length for LION models in GLUE task and included the findings in the table above.\n\n\n-----\n\n### References \n\n[1] Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers, Hwang, S., Lahoti, A., Dao, T., & Gu, A. (2024). \n\n[2] Benchmarking and building long-context retrieval models with loco and m2-bert, J Saad-Falcon, DY Fu, S Arora, N Guha, C Ré\n\n[3] Monarch Mixer: A Simple Sub-Quadratic GEMM-Based Architecture\nDaniel Y. Fu, Simran Arora, Jessica Grogan, Isys Johnson, Sabri Eyuboglu, Armin W. Thomas, Benjamin Spector, Michael Poli, Atri Rudra, Christopher Ré"}, {"review_id": "aC7mDq5ciS", "reviewer": "Reviewer_zmVf", "summary": "The manuscript introduces a technique to leverage linear transformers in non-causal settings.\nThis is done by mapping the masking mechanism of the parallel (transformer-like) formulation of linear transformers to bi-directional recurrent and chunk-wise formulations.\nExperiments on vision tasks and masked language modelling compare the proposed method to transformers and bi-directional state-space models in terms of model performance and training/inference time.", "questions": "1. Would it be possible to highlight the connection between bi-directional recurrence and parallel masking a bit better (and with fewer equations)?\n 2. Why is bidirectionality important/relevant in times of causal language modelling?\n 3. Is it possible to provide a fair speed comparison between Lion and competing methods?\n 4. Why should someone use Lion instead of regular transformers if SSMs are too slow?\n 5. How does hydra perform with alternative token sequence orders in the image classification tasks?\n 6. How well do non-bidirectional models perform on the different tasks?\n 7. How well does alternate forward-backward layers [cf. 1] perform on the different tasks?\n 8. Why does the memory consumption of the Lion RNN increase (non-linearly) with resolution in Figure&nbsp;2?\n 9. What chunk-size(s) are/were used to produce Figure&nbsp;1 and how was this chosen?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n - The proposed technique has not been presented so generally or with this much detail.\n - Improving efficiency for bidirectional sequence models can have significant impact.\n   Even though most language modelling is done causally, other tasks often benefit from bidirectionality.\n\n**Weaknesses**\n - I believe that the core idea in this work could be presented in a much more concise and clear way.\n   After all, the key insight seems to be that there is a correspondence between masking and a bidirectional recurrence.\n   Concretely, starting from a bidirectional linear attention recurrence, we obtain:\n   $$\\begin{aligned}\n       \\boldsymbol{S}^\\mathrm{f}\\_t &= \\sum\\_{\\tau=1}^t \\boldsymbol{k}_\\tau \\cdot \\boldsymbol{v}\\_\\tau \\\\\\\\\n       \\boldsymbol{S}^\\mathrm{b}\\_t &= \\sum\\_{\\tau=t}^L \\boldsymbol{k}\\_\\tau \\cdot \\boldsymbol{v}\\_\\tau \\\\\\\\\n       \\boldsymbol{y}\\_t &= \\boldsymbol{q}\\_t \\cdot \\boldsymbol{S}^\\mathrm{f}\\_t + \\boldsymbol{q}\\_t \\cdot \\boldsymbol{S}^\\mathrm{b}\\_t \\\\\\\\\n       &= \\boldsymbol{q}\\_t \\cdot \\Big(\\sum\\_{\\tau=1}^t \\boldsymbol{k}\\_\\tau \\cdot \\boldsymbol{v}\\_\\tau + \\sum\\_{\\tau=t}^L \\boldsymbol{k}\\_\\tau \\cdot \\boldsymbol{v}\\_\\tau\\Big) \\\\\\\\\n       &= \\boldsymbol{q}\\_t \\cdot \\Big(\\boldsymbol{k}\\_t \\cdot \\boldsymbol{v}\\_t + \\sum\\_{\\tau=1}^T \\boldsymbol{k}\\_\\tau \\cdot \\boldsymbol{v}\\_\\tau\\Big),\n   \\end{aligned}$$\n   which should make it easy to see that $\\boldsymbol{Y} = \\big((\\boldsymbol{1} + \\boldsymbol{I}) \\odot \\boldsymbol{Q} \\boldsymbol{K}^\\mathsf{T}\\big) \\boldsymbol{V}$.\n   I found it hard to distill this key relation from the manuscript in its current form.\n   Therefore, I believe a lot of clarity gets lost in the way these ideas have been presented.\n   Also, adding weights, scaling and other technicalities should be easier to follow once this connection has been established.\n - There is little to no motivation for why bidirectional modelling is relevant.\n   To this end, it would be useful to provide references to tasks where bidirectionality is important/useful [e.g. 1, 51, 52, 53].\n - According to the introduction, one of the main arguments for Lion is improved training speed.\n   However, line 240 indicates that no efforts have been taken to provide competitive implementations.\n   As a result, it seems more reasonable to suggest to use transformers if SSMs are too slow, because they are also faster and preserve performance even better.\n   Furthermore, a lot of techniques have been developed to improve inference times of softmax transformers.\n - Comparisons between the lion models described in lines&nbsp;227-228 and hydra seem to be useless.\n   If the goal is to compare the bi-directional modelling capacity of both models,\n   the hydra model should also be tested with different token sequences.\n - Experimental results appear to be inconclusive. \n   E.g. Tables&nbsp;2 and&nbsp;3 show that Lion is faster than Hydra, but delivers worse performance.\n - On lines&nbsp;243-244, the authors claim that Figure&nbsp;2 shows linear scaling for the RNN formulation.\n   However, the curve in Figure&nbsp;2 does not look like a straight line.\n   Furthermore, I would have expected a constant memory consumption for the RNN formulation.\n   After all, the RNN formulation only requires storage of the state and is otherwise independent of the sequence length.\n - It is unclear what the chunk-size(s) are and how they were chosen for the curves in Figure&nbsp;1.\n - Naive baselines are missing from the comparisons (e.g. running base models without Lion, alternating directions [cf. 1], ...).\n\n**Additional References**\n - [51] He et al. (2022). [Masked autoencoders are scalable vision learners](). In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition (pp. 16000-16009).\n - [52] Schiff et al. (2024). [Caduceus: Bi-Directional Equivariant Long-Range DNA Sequence Modeling](https://proceedings.mlr.press/v235/schiff24a.html). Proceedings of the 41st International Conference on Machine Learning. In Proceedings of Machine Learning Research 235:43632-43648.\n - [53] Schmidinger et al. (2025). [Bio-xLSTM: Generative modeling, representation and in-context learning of biological and chemical sequences](https://openreview.net/forum?id=IjbXZdugdj). International Conference on Learning Representations 13.", "was_revised": true, "final_justification": "The authors have addressed most of my concerns, but even after the rebuttal I am unsure if I understand what the precise goal of the proposed method is.\nIn that sense, I am unsure if my issues concerning clarity will be resolved after the modifications promised by the authors.\nNevertheless, I think the core idea can be useful and therefore I decided to increase my score towards acceptance.", "rebuttal": "We thank the reviewer for their feedback and for recognizing our work as **having significant** impact and being among the **first to present** bi-directional sequence modeling for linear transformers in such detail. \n\nThe reviewer's feedback highlights the following key concerns:\n\n- 1-3) Clarity of equations and explaining more on motivation behind bidirectional modeling.\n- 4-6) Experimental results of LION compared to baselines.\n- 7-10) Questions regarding Figures and some more ablations.\n\nWe address this in detail in the responses to each point below.\n\n-----\n\n### 1) Presenting the idea in more clear way\n\nWe thank the reviewer for suggesting the inclusion of $Y = ((1 + I) \\odot QK^\\top)V$ as the parallel form of a naive bidirectional linear transformer. As suggested, we will incorporate this clarification to better highlight the limitations of naive bidirectional linear transformers and to further clarify LION’s formulation. This formulation suffers from imbalanced attention along the diagonal due to double counting introduced by the identity matrix $I$. In contrast, LION uses $Y = (QK^\\top)V$ with correction terms in its RNN formulation to avoid this issue in the forward and backward passes. \n\n-----\n\n### 2) Motivation behind bi-directional modeling \n\nWe thank the reviewer for recommending additional references on bidirectional modeling. As many tasks beyond next-token prediction (e.g., biomedical, image, and signal domains) benefit from bidirectional modeling we will highlight this with the suggested references.\n\n-----\n\n### 3) Why not using Transformers instead of LION if SSMs are too slow\n\nWe agree with the reviewer that Transformers train faster than SSMs on bidirectional tasks while maintaining strong performance. However, their inference cost and memory requirements even by considering technical efficiency improvements remain theoretically quadratic with sequence length ($O(L^2)$). In contrast, SSMs offer linear inference complexity ($O(L)$), making them valuable for constrained resources. LION addresses this gap by building the first framework for bidirectional modeling with linear transformers **combining the fast training of softmax transformers with the inference efficiency of SSMs**, as shown in Figure 2 and Table 2.\n\n-----\n\n### 4) LION's speed comparison with baselines and using custom kernels\n\nOur main goal was to keep **LION’s implementation simple and easy to integrate**, which is why we relied on a pure PyTorch implementation. Notably, even for Transformers, custom CUDA implementations for **bidirectional tasks** show more modest speedups compared to causal tasks. As also shown in Table 1 of the FlashAttention paper [1], the speedup gain for *BERT is only 1.17× compared to 3.5× for GPT*, since bidirectional attention lacks the sparsity introduced by causal masks.\n\nImportantly, **even without custom kernels, LION achieves significantly faster training than SSMs like Vim and Hydra**, which do rely on custom kernels for speedups.  While we fully agree with the reviewer that custom implementations could further improve speed, our results already show that the pure PyTorch version outperforms existing SSMs.\n\n-----\n\n### 5) Hydra's performance with different token sequence orders\n\nWe trained Hydra with different sequence orders at the Tiny and base scale (times are reported relative to a Transformer of the same size):\n\n\nModel| Accuracy (Tiny) |Time (Tiny)| Accuracy (Base) |Time (base)|\n-|-|-|-|-\nHydra |68.7|2.3x|**81.0**|2.5x\nHydra$^\\natural$ (w/ multiple scans)|71.7|4.6x| 79.5 |5.2x\nLION-D|72.4|**1.5x**|77.8|**1.4x**\nLION-D$^\\natural$ (w/ multiple scans)|**74.2**|1.7x|80.2| 1.5x\n\n\nHydra benefits from using different scans at the Tiny scale but still fails to surpass its original version, even with extensive tuning. This approach does not scale well to the base model, likely due to Hydra’s sensitivity to its decay factor. In contrast, LION-D$^\\natural$ with **multiple scans outperforms both Hydra variants** at the Tiny scale, highlighting LION’s stability during training. Plus, unlike LION, which integrates multiple scans directly into its attention mask $M$, Hydra requires an extra pass in the rotated direction, **doubling its training time.**\n\n-----\n\n### 6) LION is faster than Hydra, but delivers worse performance.\n\nIn most scenarios (except for image classification at the base scale) **LION delivers competitive performance**. For instance, LION-D$^\\natural$ outperforms Hydra at the Tiny scale (Table above) by **3% and Small scale (Table 1 of our paper) by 1% in image classification while being ~2× faster to train**. Similarly, LION-S is **3× faster to train on the MLM task**, with only a small **0.1% gap in GLUE score** compared to Hydra. We believe that even with these marginal accuracy differences, LION’s training speed gains remain highly valuable.\n\n-----\n\n### 7) Figure 2 shows linear scaling for the RNN formulation. However, the curve does not look like a straight line.\n\nIn Figure 2, we report memory usage with respect to **image resolution**, which refers to the number of pixels (or tokens) per row. Since sequence length $L \\approx \\text{res}^2$, LION-RNN scaling quadratically with resolution means **LION scales linearly with sequence length $L$**.\nIn contrast, the Transformer in Figure 2 scales **quartically ($O(res^4)$)** with resolution (i.e., **quadratically with sequence length**). We use resolution on the x-axis as it is more intuitive for images, same visualization also is used in Vim [2] (Figure 1 c). Linearity is also clear in our MLM results in **Figure 8 (right)** in Appendix, where LION shows **linear** scaling and the Transformer **quadratic** scaling. \n\n-----\n\n### 8) The RNN formulation only requires storage of the state which is independent of the sequence length.\n\n\nRNNs have **constant memory in causal tasks (particularly in autoregressive settings)**  as each token depends only on the forward hidden state $S^F_i$. In contrast, for **bidirectional tasks**, outputs depend on both $S^F_i$ and $S^B_i$, which are not available simultaneously. This requires storing all hidden states from both directions, leading to **linear memory** in sequence length $L$. While Section A.2 of our paper describes a strategy to reduce this, memory remains linear. This is also consistent with Vim [2] in Figure 1-c where its memory scales linearly.\n\n-----\n\n### 9) How well do non-bidirectional and alternating models perform?\n\nWe investigate the importance of full bidirectional modeling in **Appendix Section C.15**, where we compare LION variants using only forward, only backward, alternating forward-backward layers, and full bidirectional processing on CIFAR classification.\n\nAs shown in **Table 16**, **LION with full bidirectionality outperforms both unidirectional variants by up to 7%**, and also surpasses the **alternating forward-backward approach by 3.6%**. This highlights the **clear benefits of bidirectionality in each layer**, rather than alternating directions or using only a single directions. We copy the results from paper below:\n\n| Model| Top-1 Acc. |\n| - | -|\n| LION-S (Forward)| 71.08|\n| LION-S (Backward)| 69.61|\n| LION-S (Alternating)    | 73.93      |\n| **LION-S (Bi-directional)**  | **77.56**  |\n| LION-S$^\\natural$ (Forward)| 70.24|\n| LION-S$^\\natural$ (Backward)| 70.42|\n| **LION-S$^\\natural$ (Bi-directional)** | **80.07**  |\n\n-----\n\n### 10) What chunk-size(s) are/were used to produce Figure 1 and how was this chosen?\n\nWe demonstrate how to choose the chunk size to balance memory and speed in **Section C.16** of our Appendix, *\"Effect of Chunk Size on Speed and Memory.\"* **Figure 11** illustrates this trade-off, with the optimal range falling between **64 and 128**, as used in the LION variants (LIT,D,S) for image classification tasks in **Figure 1**. We will bring this Figure to the main body thanks to the reviewers notice.\n\n---\n\n### References \n\n[1] FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré\n\n[2] Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model"}]

Author rebuttal:

Meta-review: This paper proposes LION, a framework for generalizing (gated) linear attention models to the bidirectional setting (similar to Hydra [Hwang et al. 2024], which generalizes SSMs to the bidirectional case). Across language/vision applications, LION is generally found to be more efficient than Hydra, at the cost of slight performance degradations.

The approach is sensible, clearly explained, and supported by a suite of experiments across different domains. The efficiency improvements over Hydra is particularly strong, despite the fact that LION does not make use of highly specialized kernels. While training efficiency compared to Transformers is not always favorable, the inference efficiency gains still make the approach worthwhile.

On the negative side, the technical novelty is somewhat limited (it's a pretty straightforward generalization of unidirectional linear attention models), and the benchmarks (GLUE/LRA) are a bit outdated. But I still think this is a worthwhile contribution to the community.
Decision: Accept (poster)