Title: Linear Attention for Efficient Bidirectional Sequence Modeling
Venue: NeurIPS 2025
Accepted: yes
Keywords: Transformer

RNN

Bidirectionality

SSM

Linear Attention

Inference

Training
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

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 Li near Attenti on , 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.

LION Code

,

LION

Blog

1 Introduction

Softmax Transformers [48] are widely used in 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 input-dependent [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

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 on 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 sequence during both training and inference. As a result, their training speed lags behind that of softmax Transformers [9, 18].

To show how Linear Transformers can be applied to bidirectional modeling and explore their capabilities, we present the LION framework, a general approach for extending Linear Transformers to bidirectional sequence modeling. LION supports three theoretically equivalent representations-full attention (for maximum training speed), bidirectional RNN (for memory-efficient inference), and chunkwise parallel form (balancing speed and memory)-mirroring the efficiency and speed advantages of Linear Transformers in the causal tasks.

We show that a 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 10 × faster than SSMs (e.g., Vim [53]) while matching the performance of SSMs and softmax Transformers. It also achieves superior inference efficiency through its RNN form, as shown in Figure 3.

2 Background

Notation. Matrices (vectors) are denoted by uppercase (lowercase) boldface letters, such as X for matrices and x for vectors. Scalars are represented by lowercase letters, e.g., x and ⊙ is Hadamard product.

Causal Linear Transformers: Transformers with Linear Attention

Given a sequence X = [ x 1 , x 2 , . . . , x L ] ⊤ ∈ R L × d , a single-head softmax attention is defined as:

where x i , q i , k i , v i , y i ∈ R d and the weights W q , W k , W v ∈ R d × d with d being the projection dimension. With Q := [ q 1 , . . . , q L ] ⊤ , K := [ k 1 , . . . , k L ] ⊤ , V := [ v 1 , . . . , v L ] ⊤ ∈ R L × d , To enable parallel training, the attention output can be expressed in the following matrix form:

where M C ∈ {-∞ , 0 } L × L is a causal mask for preventing future tokens to attend to past. In contrast, equation 1 is used during inference for generating or processing tokens. However, for causal Transformers [27], employing equation 1 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 ⊤ i k j ) with feature map function ϕ ( q i ) ⊤ ϕ ( k j ) where ϕ ( · ) : R n → R d maps the input to a higher-dimensional space [25]. For simplicity of notation, we use q i := ϕ ( W q x i ) and k i := ϕ ( W k x i ) in the sequel. This formulation allows the linear transformer to be expressed as an RNN with linear recurrence 2 , which eliminates the need

to store previous tokens in inference, while enabling parallelized training via matrix multiplication:

Here, the causal mask M C ∈ { 0 , 1 } L × L is a lower triangular matrix that enforces causal constraints. SCALE ( · ) denotes the scaling of the attention matrix across its rows (SCALE ( A ) ij = A ij ∑ L p =1 A ip ) and S i ∈ R d × d and z i ∈ R d are the hidden state matrix and the hidden state scaling vector. Eq. (3) can also written as Y = PV with P ∈ R L × L known as sequence mixer [21].

To enhance the performance of linear attention and incorporate relative position encodings into the recurrence, decay factors ( λ i ) have been introduced into the state update:

The mask M C ∈ R L × L is a lower-triangular mask generated based on the input-dependent (selective) parameter λ i as bellow:

Several Linear Transformers have been developed based on equation 4 [5, 34, 3, 44, 51], and can be unified under the form:

Here, Λ i and γ i denote model-specific decay factors and stabilizers. In this study, we focus on scalar or diagonal forms of Λ i , corresponding to the TC 0 class as stated in Merrill et al. [31], which encompasses most Linear Transformers. This unified view is also referred to as Linear Recurrent Models [52].

Chunkwise parallel form of causal Linear Transformers

Causal Linear Transformers balance the memory-speed tradeoff during training by employing chunkwise parallelization [51, 52, 6]. In this approach, the input sequence X ∈ R L × d and the corresponding query, key, and value vectors are split into L C non-overlapping chunks, each of size C .

Let Q [ i ] , K [ i ] , V [ i ] ∈ R C × d represent the chunked query, key, and value matrices, and define the chunk-level hidden state after processing i chunks as S [ i ] ∈ R d × d . The causal sequence mixer which maps the input to output Y = PV is expressed as:

S

[

i

]

=

S

[

i

-

1]

+

i

(

C

+1)

∑

j

=

+1

k

⊤

j

v

j

︸

︷︷

︸

K

[

i

]

⊤

V

[

i

]

(7)

Y

[

i

]

=

Q

[

i

]

S

[

i

]

︸

︷︷

︸

inter

+

(

Q

[

i

]

K

⊤

[

i

]

⊙

M

C

)

V

[

i

]

︸

︷︷

︸







































︸

︷︷

︸

Y

=

PV

intra

The above form is the chunkwise form of Linear Transformer without decay factor and the mask M C ∈ { 0 , 1 } L × L 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 a single intra-chunk block, is numerically unstable. This instability arises from computing the cumulative product of decay factors ∏ L t =1 Λ t and their inverse ∏ L t =1 Λ -1 t over the entire sequence, which can overflow or underflow even in logarithmic space 3 . 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 ( LCd ) and requires O ( L C ) sequential steps, as opposed to RNNs with a complexity of O ( L ) and O ( L ) steps, and attention with O ( L 2 ) 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

equivalence and demonstrate the range of inference strategies supported by our framework. We begin with the scalar decay case due to 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-based attention, this results in the form:

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:

The key construction part is to define the bidirectional mask M in a way that mirrors its properties in the causal setting. To motivate this, we begin with a remark on the structure of M C in causal setting:

Remark. Causal Linear Transformers encode relative positional information through their decay factors, which is reflected in their causal masks M C (Eq. (4)) [44, 6]. Specifically, the causal mask between tokens i and j is given by M C ij = λ j +1 λ j +2 . . . λ i , representing the product of all decays between positions j and i ( i ≥ j ).

Inspired by the above remark, we define the mask M for bidirectional sequence modeling such that M ij equals the product of all decay factors between tokens i and j . Formally, the selective M , learnable fixed decay M , and all-ones M masks can be written as:

Therefore, in the most general form (selective decay), the output of Full Linear Attention is:

where we use for upper triangular elements, for lower triangular elements, and for the diagonal of the attention matrix and mask. The selective and fixed learnable 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 to an all-ones matrix and can be omitted.

Selective Mask: To build the selective mask M , we decompose it into lower and upper triangular components (including the diagonal), denoted by M F and M B , respectively. Since both masks are rank-1 semi-separable (see proof in Appendix B.8) [21], they can be constructed efficiently via matrix multiplication. In the bidirectional setting, where all decay factors λ i are known, we stack them into a vector λ ∈ R L , analogous to stacking queries and keys into Q and K . The cumulative product L F = cumprod ( λ ) , where L F i = ∏ i k =0 λ k , is used to construct the lower triangular mask M F . For the upper triangular mask, we flip the decay vector be reversing the order of the sequence, and compute L B = cumprod ( FLIP ( λ )) . The two masks are then given by:

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( λ ) , where a ∈ R L originates from Zero-Order Hold (ZOH) discretization (see Appendix B.11). The cumulative products L F and L B can then be computed as cumulative sums in log-space, followed by exponentiation:

The full selective mask (Eq. (9)) is then assembled as M = M F + M B -I , as shown in Eq. (13).

Learnable Fixed Mask: For learnable fixed decay factor, the mask M forms a Kac-Murdock-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 ( λ = 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:

This summation causes imbalanced attention , i.e., Y = (( I + 1 ) ⊙ QK ⊤ ) V , due to double counting of the diagonal and underperforms balanced version as shown in Table 16. To address 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:





























1

λ

2

λ

2

λ

3

· · ·

λ

2

· · ·

λ

L

λ

1

1

λ

3

· · ·

λ

3

· · ·

λ

L

λ

1

λ

2

λ

2

1

· · ·

λ

4

· · ·

λ

L

.

.

.

.

.

.

.

.

.

.

.

.

.

.

.

λ

L

-

1

· · ·

λ

1

λ

L

-

1

· · ·

λ

2

λ

L

-

1

· · ·

λ

3

· · ·

1



























︸

︷︷

︸

M

=





















1

λ

1

1

λ

1

λ

2

λ

2

1

.

.

.

.

.

.

.

.

.

.

.

.

λ

L

-

1

· · ·

λ

1

λ

L

-

1

· · ·

λ

2

λ

L

-

1

· · ·

λ

3

· · ·

1



















︸

︷︷

︸

M

F

+





















1

λ

2

λ

2

λ

3

· · ·

λ

2

· · ·

λ

L

1

λ

3

· · ·

λ

3

· · ·

λ

L

1

· · ·

λ

4

· · ·

λ

L

.

.

.

.

.

.

1



















︸

︷︷

︸

M

B

-

I

(13)

As in Eq. (12) and Eq. (13), the attention matrix and mask are split into lower ( A F , M F ) and upper triangular ( A B , M B ) matrices. The scaling operator divides each row of the attention matrix to its summed value, and hence equals to a diagonal matrix C -1 multiplied by the attention:

Decomposing C into causal and non-causal parts as C i = q ⊤ i ∑ i j =1 M ij k j + q ⊤ i ∑ L j = i M ij k j -q ⊤ i k i , we can similarly split the scaling matrix into two parts as:

Since A = A F + A B and M = M F + M B -I , we can rewrite the attention output as:

Since the forward and backward recurrence matrices ( A F , A B for attention and M F , M B for mask)

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 as follows:

As seen from Eq. (4), the FORWARD part above can be expressed as an RNN. We now 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:

Above is the exact representations for the forward pass, as shown in Eq. (16), but with the tokens in reverse order. The matrices A B and M B are also modified to match the final flipped output using flipped input values V using functions F ( X ) = J L XJ L and FLIP ( X ) = J L X , where J L 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:

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:

The terms 1 2 q i ⊤ k i v i and 1 2 q i ⊤ k i are subtracted to avoid double counting. Several Linear Transformers can be generalized within our framework as shown in Table 1, where many causal recurrent models we 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 λ i :

LION-LIT for λ i = 1 which is bi-directional form of Vanilla Linear Transformer [25].

LION-D for λ i = λ fixed decay, and bi-directional form of RetNet (with scaling) [44].

LION-S , where λ i = σ ( Wx i + b ) is the bidirectional 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 activation defined as ϕ ( x ) = SILU ( x )+0 . 5 ∥ SILU ( x )+0 . 5 ∥ [19]. To ensure stability in the recurrence ( 0 < λ i ≤ 1 ) [32], we set λ i = σ ( W ⊤ a x i + b ) for LION-S and λ = σ ( a ) for LION-D, with σ ( · ) as the sigmoid and a a 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) representations, 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:

Model

LION

Bi-directional RNN

LION

Full Linear Attention

LinAtt [25]

S

i

=

S

i

-

1

+

k

i

v

⊤

i

y

i

=

q

⊤

i

S

i

S

F/B

i

=

S

F/B

i

-

1

+

k

i

v

⊤

i

y

F/B

i

=

q

i

⊤

(

S

F/B

i

-

1

2

k

i

v

i

)

,

y

i

=

y

F

i

+

y

B

i

Y

=

QK

⊤

V

+ Scaling

z

i

=

z

i

-

1

+

k

i

y

i

=

q

⊤

i

S

i

/

q

⊤

i

z

i

z

F/B

i

=

z

F/B

i

-

1

+

k

i

,

c

F/B

i

=

q

i

⊤

(

z

F/B

i

-

1

2

k

i

)

y

F/B

i

=

q

i

⊤

(

S

F/B

i

-

1

2

k

i

v

i

)

,

y

i

=

y

F

i

+

y

B

i

c

F

i

+

c

B

i

LION-LIT

Y

=

SCALE

(

QK

⊤

)

V

RetNet [44]

S

i

=

λ

S

i

-

1

+

k

i

v

⊤

i

y

i

=

q

⊤

i

S

i

S

F/B

i

=

λ

S

F/B

i

-

1

+

k

i

v

⊤

i

y

F/B

i

=

q

i

⊤

(

S

F/B

i

-

1

2

k

i

v

i

)

,

y

i

=

y

F

i

+

y

B

i

Y

= (

QK

⊤

⊙

M

)

V

,

M

ij

=

λ

|

i

-

j

|

+ Scaling

z

i

=

λ

z

i

-

1

+

k

i

y

i

=

q

⊤

i

S

i

/

q

⊤

i

z

i

z

F/B

i

=

λ

z

F/B

i

-

1

+

k

i

,

c

F/B

i

=

q

i

⊤

(

z

F/B

i

-

1

2

k

i

)

y

F/B

i

=

q

i

⊤

(

S

F/B

i

-

1

2

k

i

v

i

)

,

y

i

=

y

F

i

+

y

B

i

c

F

i

+

c

B

i

LION-D

Y

=

SCALE

(

QK

⊤

⊙

M

)

V

,

M

ij

=

λ

|

i

-

j

|

Gated RFA [34]

S

i

=

σ

(

Wx

i

)

S

i

-

1

+

k

i

v

⊤

i

z

i

=

σ

(

Wx

i

)

z

i

-

1

+

k

i

y

i

=

q

⊤

i

S

i

/

q

⊤

i

z

i

S

F/B

i

=

σ

(

Wx

i

)

S

F/B

i

-

1

+

k

i

v

⊤

i

y

i

=

q

⊤

i

S

i

z

F/B

i

=

σ

(

Wx

i

)

z

F/B

i

-

1

+

k

i

,

c

F/B

i

=

q

i

⊤

(

z

F/B

i

-

1

2

k

i

)

y

F/B

i

=

q

i

⊤

(

S

F/B

i

-

1

2

k

i

v

i

)

,

y

i

=

y

F

i

+

y

B

i

c

F

i

+

c

B

i

LION-S

Y

=

SCALE

(

QK

⊤

⊙

M

∗

)

V

,

M

∗

ij

=





Π

j

+1

i

λ

k

if

i > j,

Π

j

i

+1

λ

k

if

i < j,

1

if

i

=

j,

Mamba-2 [6]

S

i

=

λ

i

S

i

-

1

+

k

i

v

⊤

i

y

i

=

q

⊤

i

S

i

S

F/B

i

=

λ

i

S

F/B

i

-

1

+

k

i

v

⊤

i

y

F/B

i

=

q

i

⊤

(

S

F/B

i

-

1

2

k

i

v

i

)

,

y

i

=

y

F

i

+

y

B

i

Y

= (

QK

⊤

⊙

M

)

V

M

=

M

∗

RWKV-6[33]

Mamba [13]

GLA [51]

S

i

=

Diag

(

λ

i

)

S

i

-

1

+

k

i

v

⊤

i

y

i

=

q

⊤

i

S

i

S

F/B

i

=

Diag

(

λ

i

)

S

F/B

i

-

1

+

k

i

v

⊤

i

y

F/B

i

=

q

i

⊤

(

S

F/B

i

-

1

2

k

i

v

i

)

,

y

i

=

y

F

i

+

y

B

i

Y

F

=

Tril

(

Q

⊙

L

F

)(

K

⊙

(

L

F

)

-

1

)

V

Y

B

=

Triu

(

Q

⊙

L

B

)(

K

⊙

(

L

B

)

-

1

)

V

Y

=

Y

F

+

Y

B

HGRN-2 [36]

S

i

=

Diag

(

λ

i

)

S

i

-

1

+(1

-

λ

i

)

v

⊤

i

y

i

=

q

⊤

i

S

i

S

F/B

i

=

Diag

(

λ

i

)

S

F/B

i

-

1

+(1

-

λ

i

)

v

⊤

i

,

k

i

= (1

-

λ

i

)

y

F/B

i

=

q

i

⊤

(

S

F/B

i

-

1

2

k

i

v

i

)

,

y

i

=

y

F

i

+

y

B

i

Y

F

=

Tril

(

Q

⊙

L

F

)(

K

⊙

(

L

F

)

-

1

)

V

Y

B

=

Triu

(

Q

⊙

L

B

)(

K

⊙

(

L

B

)

-

1

)

V

Y

=

Y

F

+

Y

B

xLSTM [3]

S

i

=

f

i

S

i

-

1

+

i

i

k

i

v

⊤

i

z

i

=

f

i

z

i

-

1

+

i

i

k

i

y

i

=

q

⊤

i

S

i

/

|

q

⊤

i

z

i

|

,

1)

S

F/B

i

=

f

i

S

F/B

i

-

1

+

i

i

k

i

v

⊤

i

,

y

i

=

q

⊤

i

S

i

z

F/B

i

=

f

i

z

F/B

i

-

1

+

i

i

k

i

,

c

F/B

i

=

q

i

⊤

(

z

F/B

i

-

1

2

k

i

)

y

F/B

i

=

q

i

⊤

(

S

F/B

i

-

1

2

k

i

v

i

)

,

y

i

=

y

F

i

+

y

B

i

c

F

i

+

c

B

i

,

1)

Y

=

SCALE

(

QU

⊤

)

⊙

M

)

V

,

M

ij

=





Π

j

+1

i

f

k

if

i > j,

Π

j

i

+1

f

k

if

i < j,

1

if

i

=

j,

DeltaNet [52]

S

i

= (1

-

β

i

k

i

k

⊤

i

)

S

i

-

1

+

β

i

k

i

v

⊤

i

y

i

=

q

⊤

i

S

i

S

F/B

i

= (1

-

β

i

k

i

k

⊤

i

)

S

F/B

i

-

1

+

β

i

k

i

v

⊤

i

y

F/B

i

=

q

i

⊤

(

S

F/B

i

-

1

2

k

i

v

i

)

,

y

i

=

y

F

i

+

y

B

i

Y

= (

QK

⊤

)

TV

,

T

=

1

2

(

T

F

+

T

B

)

,

T

F

= (

I

+

tril

(

diag

(

β

i

)

K

i

K

⊤

i

,

-

1))

-

1

diag

(

β

i

)

,

T

B

= (

I

+

triu

(

diag

(

β

i

)

K

i

K

⊤

i

,

-

1))

-

1

diag

(

β

i

)

512

738

1024

1248

1536

2496

Resolution

0

20

40

60

80

Better

OOM

OOM

512

738

1024

1248

Resolution

10 3

10 4

Inference Time (ms)

512

738

1024

1248

1536

2496

Resolution

0

20

40

60

80

Better

OOM

OOM

512

738

1024

1248

Resolution

10 3

10 4

Inference Time (ms)

512

738

1024

1248

1536

Resolution

0

20

40

60

80

Better

OOM OOM

512

738

1024

Resolution

10 3

10 4

Inference Time (ms)

Full Attention

Chunking

RNN

Full Attention

Chunking

RNN

Full Attention

Chunking

RNN

where Sum operations applies summation over the row of the input matrix.

The chunk hidden states C [ ij ] and S [ ij ] iterate over j , with the final output for chunk i computed using their last values at j = N . The chunk mask M [ ij ] corresponds to a submatrix of the full attention mask from Eq. (13), defined as:

Below, we construct both fixed and selective chunked masks M [ ij ] . For the fixed mask, we have:

Model

#Params

Top-1 Acc. (%)

Train Time (

↓

)

Small

Base

Small

Base

Small

Base

ViT

22M

86M

72.2

77.9

×

1

×

1

DeiT

22M

86M

79.8

81.8

×

1

×

1

22M

91M

78.6

81.0

×

2.50

×

2.51

Vim

26M

98M

80.3

81.9

×

14.95

×

10.86

LION-LIT

22M

86M

72.4

74.7

×

0.74

×

0.73

LION-D

22M

86M

73.5

77.8

×

1.49

×

1.39

LION-D

♮

22M

86M

79.9

80.2

×

1.66

×

1.48

LION-S

22M

86M

74.0

76.3

×

2.03

×

1.46

LION-S

♮

22M

86M

79.6

79.9

×

2.72

×

1.68

with L ∈ R C and Γ ∈ R C × C being the vector and matrix used for creating the mask M [ ij ] and they are only depending on the decay parameter λ and the chunk size C . For the selective mask we have:

The chunkwise form in Eq. (21) can be parallelized over index i , reducing sequential operations. As bidirectional tasks require per-token outputs, memory remains subquadratic at O ( LC ) . We use this as 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 2 and tradeoff between different representation of LION is shown at Figure 1.

4 Experiments

To validate the speed and efficiency of LION we focused on well-known bidirectional sequence modeling tasks: Image Classification on ImageNet-1K [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 MLM 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-based Transformers (ViT [9], DeiT [47]) and state-of-the-art bidirectional SSMs (Vision Mamba [53], Hydra [21]). For Masked Language Modeling, we benchmark against BERT [10] and Hydra. We replace the attention in the original DeiT and BERT codebases with LION as a drop-in substitution, without 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 in different token orderings [35]. Each scan induces a different decay mask M . We show that LION naturally supports such variations through alternate token sequences, leading to models like LION-D ♮ 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-D ♮

a)TransformerBlock

b)LiONAttentionParallel

c) LION RNN Recurrence

Norm

Scale

Inv

Sequence

Mixer

A

M

B

B

Norm

Forward

Backward

MF

MB

Recurrence

Recurrence

Channel

Mixer

Proj KQV入

X

Proj V

Proj KQ

Proj 入

T

X

X

Forward Sequence Modeling

Non-Linear Operations

InputProjections

Channel Mixing Block

SequenceMixingBlock

and Vim are the only models to surpass DeiT in the Small scale, with LION-D ♮ being 10 × faster to train, as 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 in the optimal region of the accuracy-speed trade-off ( top right corner of Figure in Table 2), achieving softmax-level accuracy and outperforming state-of-the-art SSMs in training speed due to 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 tradeoff between memory and speed. To illustrate this tradeoff, 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 a balance, with memory usage between the two depending on type of mask. Both full attention and 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 it 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 in Figure 1 ( Middle ).

We study the effect of chunk size on LION inference with 1248 -resolution images, observing a trade-off where larger chunks reduce time but increase memory, with OOM beyond size 1024 , and we typically use 8 -16 chunks as a balance as shown in Figure 4.

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 's 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-art SSMs like Hydra, while being 3 × faster to train. Thanks to their RNN formulation, LION models are also significantly more memory-efficient than BERT during inference

Inference memory efficiency While LION matches the accuracy and training speed of softmax-based Transformers, it scales linearly during inference via its RNN formulation, unlike softmax attention which scales quadratically with sequence length (Figure 3). Among all baselines, LION achieves the most efficient inference scaling. This highlights the importance of including both attention and RNN representation for Linear Transformers in bi-directional modeling to achieve the best of both worlds: fast training and efficient inference. Moreover, inference times for all models ar provided at Appendix C.10.

224

512

738

1024 1248

2496

0

10

20

30

40

50

60

70

80

OOM

OOM

OOM

OOM

Better

- 82.0 % memory

- 52.5 % memory

- 38.7 % memory

- 64.8 % memory

ViT

LION

VIM

ViT-Chunk

Resolution

80

Memory (Left)

(GB)

Time (Right)

60

104

GPUMemory

40

Time

OOM

20

4

1632641282565121024

Chunk size

(Figure 9), consistent with results from image classification. LION achieves significantly faster training than Hydra while maintaining comparable performance, marginally lower at larger scales and higher at 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 in 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 Λ i = exp( σ ( A i ) + B i ) and for LION-D as Λ i = exp( B i ) , where B i 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 was unable to solve LRA, consistent with prior findings [45, 32]. For random initializations ablation and configurations please check Table 5 and Appendix C.3.

5 Conclusion

We introduce LION, as the first unified framework to adapt Linear Transformers into their bidirectional 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 of several existing Linear Transformers into their bidirectional versions using LION (Table 1). We evaluate three representative models-LION-LIT, LION-D, and LION-S -based on different decay types, on masked language modeling and image classification tasks.

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 Krawczuk for insightful discussions on LION models for the masked language modeling experiment. This work was partially supported by Hasler Foundation Program: Hasler Responsible AI (project number 21043), partially sponsored by the Army Research Office and was accomplished under Grant Number W911NF-24-1-0048, and partially funded by the Swiss National Science Foundation (SNSF) under grant number 200021205011. 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

Benedikt Alkin, Maximilian Beck, Korbinian Pöppel, Sepp Hochreiter, and Johannes Brandstetter. Vision-LSTM: xLSTM as generic vision backbone. arXiv preprint arXiv:2406.04303 , 2024.

Jimmy Lei Ba. Layer normalization. arXiv preprint arXiv:1607.06450 , 2016.

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.04517 , 2024.

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 are few-shot learners. 2020.

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 J Colwell, and Adrian Weller. Rethinking attention with performers. In International Conference on Learning Representations , 2021.

Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. In Forty-first International Conference on Machine Learning , 2024.

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 1 (Long and Short Papers) , 2019.

Jesse Dodge, Maarten Sap, Ana Marasovi´ c, 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.

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 Representations , 2021.

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.

Team Gemini, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. Gemini: a family of highly capable multimodal models. arXiv preprint arXiv:2312.11805 , 2023.

Aaron Gokaslan and Vanya Cohen. Openwebtext corpus. http://Skylion007.github.io/ OpenWebTextCorpus , 2019.

Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. In Conference on Learning and Modeling (COLM 2024) , 2024.

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.

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

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.

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/CVF conference on computer vision and pattern recognition , pages 16000-16009, 2022.

Pengcheng He, Xiaodong Liu, Jianfeng Gao, and Weizhu Chen. Deberta: Decoding-enhanced bert with disentangled attention. arXiv preprint arXiv:2006.03654 , 2020.

Alex Henry, Prudhvi Raj Dachapally, Shubham Pawar, and Yuxuan Chen. Query-key normalization for transformers. arXiv preprint arXiv:2010.04245 , 2020.

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.

Sukjun Hwang, Aakash Lahoti, Tri Dao, and Albert Gu. Hydra: Bidirectional state space models through generalized matrix mixers, 2024.

Peter Izsak, Moshe Berchansky, and Omer Levy. How to train BERT with an academic budget. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing , 2021.

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.

Rudolph Emil Kalman. A new approach to linear filtering and prediction problems. 1960.

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

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

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-22213, 2022.

Kunchang Li, Xinhao Li, Yi Wang, Yinan He, Yali Wang, Limin Wang, and Yu Qiao. Videomamba: State space model for efficient video understanding. In European Conference on Computer Vision , pages 237-255. Springer, 2024.

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.

Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv:1608.03983 , 2016.

William Merrill, Jackson Petty, and Ashish Sabharwal. The illusion of state in state-space models. arXiv preprint arXiv:2404.08819 , 2024.

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.

Bo Peng, Daniel Goldstein, Quentin Anthony, Alon Albalak, Eric Alcaide, Stella Biderman, Eugene Cheah, Teddy Ferdinan, Haowen Hou, Przemysław Kazienko, et al. Eagle and finch: Rwkv with matrix-valued states and dynamic recurrence. arXiv preprint arXiv:2404.05892 , 2024.

Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. arXiv preprint arXiv:2103.02143 , 2021.

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.21022 , 2024.

Zhen Qin, Songlin Yang, Weixuan Sun, Xuyang Shen, Dong Li, Weigao Sun, and Yiran Zhong. Hgrn2: Gated linear rnns with state expansion. arXiv preprint arXiv:2404.07904 , 2024.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

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. 115(3):211-252, 2015.

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.

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.

Jürgen Schmidhuber. Learning to control fast-weight memories: An alternative to dynamic recurrent networks. Neural Computation , 4(1):131-139, 1992.

Niklas Schmidinger, Lisa Schneckenreiter, Philipp Seidl, Johannes Schimunek, Pieter-Jan Hoedt, Johannes Brandstetter, Andreas Mayr, Sohvi Luukkonen, Sepp Hochreiter, and Günter Klambauer. Bio-xlstm: Generative modeling, representation and in-context learning of biological and chemical sequences. arXiv preprint arXiv:2411.04165 , 2024.

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.

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

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 arXiv:2011.04006 , 2020.

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.12877, 2020.

Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International conference on machine learning , pages 10347-10357. PMLR, 2021.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, \Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

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 BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP , 2018.

Ross Wightman. Pytorch image models. https://github.com/rwightman/ pytorch-image-models , 2019.

Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficient training. arXiv preprint arXiv:2312.06635 , 2023.

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.

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.

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 be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We included in 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 are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory assumptions and proofs

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

Answer: [Yes]

Justification: We have proved all the theoretical parts of the paper in main body and Appendix.

Guidelines:

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

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

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

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

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

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

4. Experimental result reproducibility

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

Answer: [Yes]

Justification: Details of reproduction are in the Appendix. Our code is also publicly available.

Guidelines:

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

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

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

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

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

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

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

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

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

5. Open access to data and code

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

Answer: [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/ public/guides/CodeSubmissionPolicy ) for more details.

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

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

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

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

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

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

6. Experimental setting/details

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

Answer: [Yes] ,

Justification: In 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 a level of detail that is necessary to appreciate the results and make sense of them.

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

7. Experiment statistical significance

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

Answer: [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, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

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

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

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

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

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

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

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

8. Experiments compute resources

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

Answer: [Yes] ,

Justification: In Appendix all details are 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 of compute required for each of the individual experimental runs as well as estimate the total compute.

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

9. Code of ethics

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

Answer: [Yes] ,

Justification: It 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 circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA] ,

Justification: 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 a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to

generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA] .

Justification: No high risk models.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: 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 or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: Yes.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and research with human subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer:[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 in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional review board (IRB) approvals or equivalent for research with human subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer:[No]

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, or originality of the research, declaration is not required.

Answer: [NA]

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, or non-standard components.

Please refer to our LLM policy ( https://neurips.cc/Conferences/2025/LLM ) for what should or should not be described.

Appendix

Table of Contents

A Causal Language Modeling

We also evaluated the performance of LION-S with decay factor of λ i = σ ( W a x i + 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 our models in the OpenWebText corpus [12]. We evaluate the architectures in the 124 M parameter setup. Our implementation is based on nanoGPT 4 . We use the default GPT-2 hyperparameters and train our models for 8 days in 4 NVIDIA A100 SXM4 40 GB GPUs.

Model

Perplexity

GPT-2

17

.

42

(

±

1

.

11)

LinAtt

21

.

07

(

±

1

.

32)

LION-S (1D)

18

.

16

(

±

1

.

16)

0

1000

2000

3000

4000

Sequence Length

10

15

20

25

30

35

Perplexity

18.25 Avg.

LION-S (1D)

GPT-2

LinAtt

(a) OpenWebText PPL

(b) Perplexity vs. sequence length

In Figure 5 we can observe LION-S (1D) 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 6 we evaluate the latency and memory of LION-S (1D) 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

0

2000

4000

6000

8000

Sequence length

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

1

.

25

Next token latency (s)

Recurrence

Attention

Attention + KV Cache

(a) Latency

0

2000

4000

6000

8000

Sequence Length

8

10

12

14

16

Memory (GB)

16.83

7.90

6.95

6.95

Recurrence

Attention

Attention + KV Cache

(b) Memory

models in training and inference, with a strong performance (up to a small degradation) can be obtained.

A.1 Memory allocation in LION during Forward and Backward recurrences

Backward

Backward

y

yi-1

B

ye-2

ci-1

B

Ci-2

B

Ci-3

B

yi-3

F

yi-2

yi-1

yE

Re-writtenonsameMemorycell

ci-3

F

Ci-2

F

Re-writtenonsameMemorycell

Forward

Forward

yi

Used for computing the output

Ci

B Proofs

B.1 Duality between Linear Recurrence and Attention

Considering the following recurrence:

We can calculate each output y i recursively as below:

This can be shown in a vectorized form as:

Where SCALE is the scaling function which scaled the attention matrix with respect to each row or can also be written as:

Similarly if the SCALE is applied before masking we have:

With M CAUSAL being the causal mask used in autoregressive models [27]. This vectorized form is equivalent to:

With A C being:

And the recurrence for this vectorized form can be written as:

B.2 Forward and Backward Recurrences Theoretical Details

Considering the following recurrence:

This recurrence is the same as recurrence equation 35 but with z L being fixed to the summation of all keys in the sequence, therefor the output y i can simply be written as:

By replacing the z i = ∑ i j =1 k j in the denominator of equation equation 37 with z L . Therefore in vectorized form, it will become:

Importantly this equation can be written as:

which despite equation equation 33 scaling is applied over the whole sequence not for the causal part of the sequence. The matrix A C is helpful for driving the recurrent version of LION for Forward and Backward recurrences and the mask here M is equal to LION's forward mask M F in equation equation 16. As shown in equation 16 the forward recurrence for the causal part of the attention can be presented as Y B = A F ⊙ M F the matrix A F can be created simply by using matrix A C as bellow:























1

2

q

⊤

1

k

1

q

⊤

1

z

L

q

⊤

2

k

1

q

⊤

2

z

L

1

2

q

⊤

2

k

2

q

⊤

2

z

L

q

⊤

3

k

1

q

⊤

3

z

L

q

⊤

3

k

2

q

⊤

3

z

L

1

2

q

⊤

3

k

3

q

⊤

3

z

L

.

.

.

.

.

.

.

.

.

.

.

.

q

⊤

L

k

1

q

⊤

L

z

L

q

⊤

L

k

2

q

⊤

L

z

L

· · ·

1

2

q

⊤

L

k

L

q

⊤

L

z

L





















︸

︷︷

︸

A

F

=























q

⊤

1

k

1

q

⊤

1

z

L

q

⊤

2

k

1

q

⊤

2

z

L

q

⊤

2

k

2

q

⊤

2

z

L

q

⊤

3

k

1

q

⊤

3

z

L

q

⊤

3

k

2

q

⊤

3

z

L

q

⊤

3

k

3

q

⊤

3

z

L

.

.

.

.

.

.

.

.

.

.

.

.

q

⊤

L

k

1

q

⊤

L

z

L

q

⊤

L

k

2

q

⊤

L

z

L

· · ·

q

⊤

L

k

L

q

⊤

L

z

L





















︸

︷︷

︸

A

C

-























1

2

q

⊤

1

k

1

q

⊤

1

z

L

1

2

q

⊤

2

k

2

q

⊤

2

z

L

1

2

q

⊤

3

k

3

q

⊤

3

z

L

.

.

.

1

2

q

⊤

L

k

L

q

⊤

L

z

L





















︸

︷︷

︸

D

F

Or equivalently:

Since the diagonal values of the mask M F are all ones and the matrix D F is diagonal, we have:

As A C ⊙ M F corresponds to linear recurrence shown at equation 41. The vectorized form equation 45 can be presented as linear recurrence:

This is equivalent to the linear recurrence presented in equation equation 40. The same theoretical approach applies to the backward recurrence, leading to the following linear recurrence for both recurrences:

However, the above equation requires access to the summation of scaling values z L . 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 F and Y 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 z L 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 to sum all the keys ( z L ). This section we will proof how to avoid an extra scaling recurrence by simple modifications to equation equation 47 and equation 48.

Considering having a scaling recurrence as part of forward and backward recurrence we will have:

The equations above are similar to the previous ones, with the addition of scalar states c F and c B for the backward and forward recurrences, respectively. During each recurrence, the outputs y F i and y B i , along with the scalars c F i and c B i , are saved for each token to construct the final output of each layer. It is also important to note that there is no need to save z F and z B for each token; these states can simply be overwritten in memory. The final output of each layer is equal to:

Where y F i and y B i can be written as:

So the addition y F i + y B i is equal to:

Where by considering the mask M as bellow:

The above mask is equal to M F + M B -I , allowing equation equation 53 to be rewritten as:

So we can finally find the output of each layer y i as:

It can easily be shown that:

So the final output of the layer is:

Alternatively, in vectorized form, it can be expressed as:

with M being the attention mask created by λ i s as in equation 55.

B.4 Flipping Operation in Backward recurrence

Here we define the operation which flip the matrices A B , M B for the reverse reccurence th goal is to find the F ( . ) such that:

The above can be achieved by:

B.5 LION framework for diagonal decay

The recurrence in the causal form for a diagonal case is presented as:

which followed by Dao and Gu [6], Hwang et al. [21] it can be shown in matrix form as:

with P being the sequence mixer of the Linear Transformer. We can build Full Attention starting by decomposing the decay factros Λ i as followes:

Q

⊤

i

Λ

0

Λ

1

. . .

Λ

i

×

(

Λ

0

)

(

Λ

1

)

. . .

(

Λ

j

)

K

j

In the above decomposition blue parts are only depending on the index i where the magenta depending on j this implies by defining Q ∗ i = Λ × 0: i Q i and K ∗ j = ( Λ × 0: j ) -1 K j the Equation equation 72 can be written exactly as attention in parallel form M = Q ∗ K ∗⊤ . Note that since i > j , all terms with indices k < j cancel out, leaving only the matrices Λ k with j +1 ≤ k ≤ i to be multiplied in the expression Λ × 0: i ( Λ × 0: j ) -1 . 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 A , A , . . . , A , the product A = A A · · · A

where a i ∈ R N is a vector containing the diagonal elements of the matrix A i ∈ R N × N or a i = DIAG ( A i ) .

Utilizing Remark , we can rewrite equation equation 72 as follows:

where in the above, ( λ 0 ⊙ λ 1 ⊙ λ 2 . . . λ j ) -1 indicates the element-wise inverse ( Inv ) of the vector λ 0 ⊙ λ 1 ⊙ λ 2 . . . λ j . During training, we have access to all L tokens. Thus, we can compute the vectors λ i for all i and construct a matrix D ∈ R L × N , where the i -th row D i = λ i = DIAG ( Λ i ) .

Our goal is to find a matrix L ∈ R L × N whose i -th row L i is given by λ 0 ⊙ λ 1 ⊙ λ 2 ⊙···⊙ λ i . This allows us to simplify the term Q i ⊙ ( λ 0 ⊙ λ 1 ⊙··· ⊙ λ i ) in equation 72 to Q i ⊙ L i .

Since L ip = ∏ i k =0 λ kp , it can be viewed as the cumulative product of the vectors λ i . Consequently, we can obtain L by applying a cumprod operation along the sequence dimension of D :

A similar approach applies to ( λ 0 ⊙ λ 1 ⊙··· ⊙ λ j ) -1 . Instead of computing the cumulative product of λ i directly, we compute the cumulative product of the element-wise inverse of these vectors, leading to the definition of the matrix U :

-

1

-

1

-

1

Note that U = Inv( L ) is the element-wise inverse of L , and thus L ⊙ U = 1 , where 1 ∈ R L × N is an all-ones matrix.

By analogy with Transformers [48], we can store the states Q i , K i , Λ i for all tokens in matrix form as Q , K , D ∈ R L × N . Under these conditions, a selective diagonal SSM (Equation equation 69) can be simplified to:

The TRIL function in the above formulation is indicating that the output Y is computed only from previous tokens, thereby preserving the causal structure analogous to causal Linear Transformers, where Y = TRIL ( QK ⊤ ) V .

Note that this parallel form is less straightforward than the scalar case, but the masks L F , L B can still be decomposed from the queries and keys as shown in Eq. (11). Theorem B.1 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.1 is same as LION -RNN presented at Theorem 3.1 without scaling:

B.6 Mapping Existing autoregressive models into LION

As noted, other autoregressive recurrent models can also be integrated into our bidirectional framework, 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:

As observed, this is a special case of our bidirectional recurrence defined in equation 19 with λ i = 1 , 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 as follows:

=

Y

=

SCALE

(

QK

⊤

V

)

(93)

For the non-scaled variant, we simply remove the scaling state z as well as the scaling parameter c . Consequently, the bidirectional linear transformer, which is equivalent to and parallelizable with attention without scaling, can be expressed as follows:

The final output for scaled version can be extracted as y i = y B i + y B i c B i + c B i for scaled and as y i = y B i + y B i for non-scaled version. Variations of linear transformers, such as Performer [5], which employ different non-linearities ϕ ( . ) for keys and queries, can be adapted to a bidirectional format using the framework established for linear transformers.

Retentive Network (LION-D) . According to [44] the forward equation for a retentive network can be written as:

S

F

i

=

λ

S

F

-

i

1

+

k

i

v

⊤

i

,

(97)

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 λ i = λ values:

Note that: M R ij = λ | i -j | , 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:

The above recurrence is equivalent to equation 19 by considering i i k i as a new key. The term i i k i can be easily vectorized by aggregating all i i values for each token into a vector i . Thus, we can express the vectorized form of the bidirectional xLSTM and its equivalence to attention as follows:

where the mask M f is equal to the LION mask equation 55 just by replacing λ i = f i . And where operation SCALE max consider the maximum of operation in the denominator as:

Gated RFA ( LION-S ) Gated RFA [51] in autoregressive mode exhibits a recurrence similar to that of xLSTM, with only minor differences:

Thus, the bidirectional version of the model retains a similar output, achieved by replacing the vector i in equation 110 with 1 -g , where g represents the vectorized form of all scalar values g i .

Note that for LION-S we just use different non-linearty compared to original GRFA and the term (1 -g i ) can be considered as part of k i .

Mamba2 Mamba2 [6] is the selective SSM with scalar decay factor λ i = SoftPlus ( Wx i ) we can observe that Mamba simply can be considered as variant of LION-S with different activation and without scaling therefor we can extend it to bi-directional setting via:

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

RWKV-6 : RWKV-6 [33] is a Linear Transformer with diagonal decay it causal recurrence can be written as:

as seen this choice is exactly aligned with our proof on Theorem B.1 which allows the bi-directional form of:

HGRN-2 : HGRN-2 [36] has similar recurrence as RWKV-6 just by replacing the k i = 1 -Diag ( Λ i ) therefor its equal bi-directional form via LION is alike Appendix B.6.

B.7 Discussion on DeltaNet

For the special case of DeltaNet [52], our theorem supports an equivalent bidirectional RNN formulation with LION-style correction terms to avoid double counting:

The chunkwise parallel form of DeltaNet (Eqs. 8-11 in the original paper) can also be extended to the bidirectional setting by removing causal masks and applying Theorem 3.2 for the chunkwise form of LION. However, a key bottleneck in DeltaNet arises in its full attention formulation, where the attention is computed as

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 M F & M B are Semiseperable with rank-1

For the lower triangular part of the selective mask M F , the upper triangular part can be filled such that it creates a full matrix with rank 1, which aligns with the definition of a semi-separable matrix with rank 1, as below:

M

F

=

1

λ

1

1

λ

1

λ

2

λ

2

1

.

.

.

.

.

.

.

.

.

.

.

.

λ

L

-

1

· · ·

λ

1

λ

L

-

1

· · ·

λ

2

λ

L

-

1

· · ·

λ

3

· · ·

1

= TRIL





























1

λ

-

1

1

λ

-

1

1

λ

-

1

2

· · ·

λ

-

1

1

· · ·

λ

-

1

L

-

1

λ

1

1

λ

-

1

2

· · ·

λ

-

1

2

· · ·

λ

-

1

L

-

1

λ

1

λ

2

λ

2

1

· · ·

λ

-

1

3

· · ·

λ

-

1

L

-

1

.

.

.

.

.

.

.

.

.

.

.

.

.

.

.

λ

L

-

1

· · ·

λ

1

λ

L

-

1

· · ·

λ

2

λ

L

-

1

· · ·

λ

3

· · ·

1



























= TRIL









































1

λ

1

λ

1

λ

2

.

.

.

λ

L

-

1

· · ·

λ

1

︸

︷︷

︸

1

λ

-

1

1

λ

-

1

1

λ

-

1

2

· · ·

λ

-

1

1

· · ·

λ

-

1

L

-

1

︸

︷︷

︸

u

⊤

= Inv(

l

)

⊤







































=

TRIL

(

l u

⊤

)

l

where TRIL is the function which masks the upper part of the matrix and set it to zero. Same is applied for the upper triangular part of the matrix M B as well. Also since decay mask is the specific case of selective mask by setting λ i = λ 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:

by apply chunking to queries/keys/values and defining the:

qIk1

qk

qIk3

qIk

qiks

qik6

qIk

qkg

qkg

1

入

x4

26

1

zs

qk1

qkz

qk6

qk

qkg

入

1

a3

25

A6

1

g . A6

=Ag

qk1

qk

gk3

qkA

qk6

qk

qke

qkg

x²

入

1

入

26

1

14

A4hs

4.ay

4-Ag

4 ..kg

qk1

q7kz

qiks

qk6

qk

qkg

Ikg

a2

1

x²

A3

x4

1

asg

as ..ag

qk1

qkz

gk6

9kg

23

a2

入

1

入

24

1

6

Agang

qk1

qkz

qEks

qk6

qk

qkg

qkg

x4

a3

x²

1

入

x²

2-as

A4s

1

qk1

q7k2

q5k4

q7k6

97k7

qk8

q7kg

26

a5

a4

a2

1

1 .6

A . A6

A4as26

As26

16

1

Ag

Agng

qgk1

qkz

qEks

qk6

qk

qake

qkg

x

26

x²

入

1

4 ..ay

Aga

1

2g

qk1

qk2

gks

qk6

qk

qkg

qkg

28

x”

a6

25

a4

a3

1

1 . g

入 ..Ag

4 .

sg

g2

Ag

1

we can simply rewrite the Eq. (128) in chunkwise form as:

where N is the number of total chunks and N = L C 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 as stored in C ij and the final output for chunk i is computed by using the last chunkwise hidden state S [ i ] divided by the scaling for that chunk C [ i ] which are equal to S [ iN ] , C [ iN ] .

To construct the chunkwise mask M ij , we define the chunk-level selective parameters as:

Since the full mask is composed of lower and upper triangular components:

we determine the appropriate chunkwise form based on relative chunk positions:

If i > j , the chunk falls entirely within the lower triangular part, requiring only M F , which can be efficiently computed as L F [ i ] 1 L F [ j ] .

If i < j , the chunk is fully in the upper triangular region, needing only M B , which follows from L B [ j ] 1 L B [ i ] .

If i = j , the chunk lies along the diagonal and requires both the lower triangular part of M F and the upper triangular part of M B , expressed as:

The full matrix is:

For the fixed decay mask, a simpler case of the general selective mask, L F and L B are identical and simplify to L i = λ i . Since the full mask follows M ij = λ | i -j | , the chunkwise mask for i, j can be written as:

Similarly, for the upper triangular part:

For diagonal chunks, the mask remains a fixed matrix Γ ∈ R C × C , where Γ ij = λ | i -j | , representing a smaller version of the full fixed decay mask M ∈ R L × L with M ij = λ | i -j | . 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 simultaneously, leading to the following matrix multiplications in the parallel chunkwise form:

and the mask M [ j ] is created based on:

Consider a real matrix X ∈ R n × n . The operator Tril( X , d ) returns the lower-triangular region of X , including all elements on and below the diagonal shifted by d . In other words, tril( X , d ) ij = X ij whenever j ≤ i + d and is 0 otherwise. Similarly, Triu( X , d ) returns the upper-triangular region, keeping elements on and above the diagonal shifted by d . Formally, Triu( X , d ) ij = X ij if j ≥ i + d and 0 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:

which can be rearranged to isolate h ( t ) :

By multiplying the equation by e -At , we get

Since ∂ ∂t e At = Ae At = e At A , Eq. (136) can be written as:

After integrating both sides and simplifications, we get

By multiplying both sides by e At to isolate h ( t ) and performing further simplifications, at the end we get

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 k in terms of h ( kT ) such that

Using the new notation, Eq. (139) becomes

Now we want to express the system in the form:

To start, let's write out the equation for x k +1 as

After multiplying by e A T and rearranging we get

Plugging this expression for x k +1 in Eq. (143) yields to

which can be further simplified to

Now, assuming that x ( t ) is constant on the interval [ kT, ( k +1) T ) , which allows us to take Bx ( t ) outside the integral. Moreover, by bringing the e A ( k +1) T term inside the integral we have

Using a change of variables v = ( k +1) T -τ , with dτ = -dv , and reversing the integration bounds results in

Finally, if we evaluate the integral by noting that d dt e A t = A e A t and assuming A is invertible, we get

Thus, we find the discrete-time state and input matrices:

And the final desecrate state space representation is:

As in case of LION-S (similar to choice of Mamba2 [6]) the matrix A is identity while the time step T is selective and equal to a i . And simply for LION-S scenario the term Bx ( t ) will change into k i v ⊤ i therefor considering Linear Transformer as continuous system like:

(155)

By applying the ZOH discritization the final descreate LION-S will be equal to:

DISCRETE

And it applies to both directions forward and backward.

C Additional experimental validation

C.1 LRA full dataset results for LION

C.2 LRA Configurations for LION

For the LRA task, we utilized the same model dimensions as specified in the S5 [43] paper, following the guidelines from the S5 GitHub repository 5 . Our state matrix was represented as a vector Λ i = λ i , where each element contains a scalar non-input dependent value e a . The value a was initialized based on HIPPO theory, alongside the input-dependent a i , as described in main body.

We employed the ADAMW optimizer with an initial learning rate of 5 × 10 -4 and a cosine learning rate scheduler [30]. The weights for the queries and keys, as well as the selective component of Λ , were initialized using a Gaussian distribution with a standard deviation of 0.1. For the values v , we initialized W v using zero-order hold discretization, represented as W init v = ( Λ -1 · (Λ -I ) ) . The non-selective parts of Λ 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 in LION, which scales the attention before applying the mask expressed as Y = SCALE ( QK ⊤ ) ⊙ M has proven ineffective in addressing the challenging PathX task, as shown in Table 6 [5] .

C.4 LRA full Results

We have evaluated LION variants against benchmarks for long-range sequence modeling across different categories, including softmax-based Transformers, RNNs, SSMs, and Linear Transformers.

C.5 Experimental details for the MLM/GLUE 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/LION-D/LION-S blocks while keeping all hyperparameters the

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.e., https://github.com/HazyResearch/m2 .

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 a 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 -4 , β 1 = 0 . 9 , β 2 = 0 . 98 , ϵ = 10 -6 and weight decay 10 -5 . As a scheduler, we perform a linear warm-up for 6% of the training steps and a 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 6 and BERT 24 finetuning 7 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 et al. [10], Section C.1 and the file this link.

M2-LARGE : Available in Fu et al. [10], Section C.1 and the file this link.

Hydra : Available in Hwang et al. [21], Section D.2 and the file this link.

Modified : Same as 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.

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-LARGE 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 -5 (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 in 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 v.s. 82 . 25 for the BASE model size and 81 . 58 v.s. 82 . 95 for the LARGE model size.

C.7 Additional experimental results for the MLM/GLUE 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), we

0

200

400

600

800

1000

Sequence Length

20

30

40

50

60

70

Training Sequence Length

Out of Positional Encodings

+36.38 accuracy

LION-S

BERT

LION-LIT

0

5000

10000

15000

Context Length

0

20

40

60

80

OOM

Better

- 84.10 % memory

BERT

LION

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

C.8 ImageNet classification results for Tiny scale

In Table 12, we present the image classification results of LION models on 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. At 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 on plain Python implementations, leading to increasing overhead as resolution grows.

C.10 Ablation studies with image classification

Choice of λ i values. In this section, we study the properties of the selectivity parameter a i on CIFAR-100 dataset. We tested, three cases: ( i ) fixed mask with scalars a i = a i , ( ii ) vector, inputdependent a i ∈ R d ( cf. , Appendix B.5) and iii) input dependent scalar a i ∈ R . The results, presented

224

512

738

1024

1248

0

200

400

600

800

1000

Inference Time (s)

ViT

LION-D

VIM

Resolution

in Table 13, show that while the input dependency is beneficial, the expansion of a i 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.

Understanding the power of non-linearity, softmax, and positional embeddings. In Table 14, we present additional ablations on certain design elements of a 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 a nice accuracy. Though positional embedding boosts the accuracy, a mask can easily replace it.

C.11 Hyperparameters for Training Image Classifiers

All experiments were conducted on a single machine for CIFAR-100 and multiple machines for ImageNet, using NVIDIA A100 SXM4 80GB GPUs. For LION models and Hydra, the ViT architecture 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 a 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 ( L 2 + LD 2 ) LION-S has O ( LD 2 ) . Note that in this calculation, the exponentials and other nonlinearities are considered as 1 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 a vision transformer:

Calculating Q , K , V : 6 LD 2 ,

Attention A = QK T : 2 L 2 D

Softmax (assuming 1 FLOP for exp): 2 L 2

Calculating Y : 2 L 2 D

TOTAL: L (6 D 2 +4 LD +2 L )

The number of FLOPs in the attention module for LION:

Calculating Q , K , V , λ : 6 LD 2 +2 LD , · For each token in one forward/backward recurrence: -Updating S F/B i : 3 D 2 -Updating z F/B i : 2 D -Calculating c F/B i : 4 D +2 -Calculating y F/B i : 2 D 2 +4 D +1 -Total: 5 D 2 +10 D +3 · L forward + backward recurrences: 2 L (5 D 2 +10 D +3) · Calculating Y : 2 L ( D +1) · TOTAL: L (16 D 2 +24 D +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 are as critical as the pixel values themselves. Positional embeddings provide a 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 to bottom as illustrated in the left panel of Figure 11.

Linear Transformers and SSMs such as 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 are usually applied in both vertical and horizontal directions [1, 29]. These scans act as separate

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 improvement 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 [1].

C.15 Ablation studies on importance of bi-directionality on image classification

To highlight the importance of bi-directionality 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 patches in the forward direction while even-numbered blocks process them in the backward direction (Forward-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 a significant boost of up to 10% .

Learned decay parameters in LION-D-LARGE

Learned decay parameters in LION-D-BASE

Learned decay parameters in LION-D-BASE Vision

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

0.0

0

2

4

6

8

10

0

5

10

15

20

0

2

4

9

8

10

Layer

Layer

Layer

C.16 LION-D Decay factor distribution

We visualize the distribution of 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.

D Future Direction & Limitation

While LION maps many Linear Transformers to 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 to larger models, potentially suppressing softmax Transformers in multiple domains and tasks.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nMajor 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.\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n# 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?\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n## 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?\n\nQUESTIONS:\nSee Weaknesses, I have incorporated questions in points brought up to avoid repetition.\n\nLIMITATIONS:\nThe 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])", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-22"}