LION: A bidirectional framework that trains like a Transformer and infers like an RNN

We are glad that you enjoyed our paper and found our contribution valuable to the community! We are happy that our illustrations were helpful and our experiments thorough. Below, we have addressed the remaining questions raised by the reviewer:

---

1) Causal Language Modeling and Vision experiments are rather small scale

We agree with the reviewer appreciate their comment that larger scale models would strenghten our papers. Since pre-training is a compute intensive task, our experiments focus on smaller scales. After the initial submission, we expanded our experiments from tiny models (5.5M) to small scale (21.7M) for vision and from base (110M) to large(340M) models for causal language task. In the below tables and in Table 3-4 of our paper, you can find the results with these scaled models. In both cases, we observed great benefits with larger models. Our results show that as the model size increases, the gap between LION-S and the transformer architecture gets less and less insignificant.

We believe that these new results further proves the capabilities of the LION architecture.

---

2) Continuous domain of LION-S is rather loose

Thank you for the comment. As also requested by Reviewer sped, we have added details in Appendix Section B.6 explaining how discretization from the continuous domain using zero-order hold filtering leads to LION-S.

---

3) Terminology: "attention" and "backwards path"

Thanks for the suggestion, we have modified the paper to use a unified terminology for attention and to refer backward/forward path as backward/forward recurrence.

We appreciate your curiosity regarding the name "LION." It was chosen to reflect LInear AttentiON, and it aligns with the naming conventions used in the community, such as HIPPO, Mamba, Griffin, and others.

---

4) Default non-linearity for LION-S? Only Mask and no non-linearity or mask and non-linearity?

In Table 8, the applied non-linearity is elu()+1 for LION-LIT and silu normalized for LION-S on top of the mask.

---

5) Mask $M$ for Linear Transformer Equations 109 and 112 in new pdf version

Thank you for the detailed notice. The mask in Equation 101 of the original manuscript was missing, and we have now revised this equation for xLSTM, including the extension with the maximum operator as well.

Since the Equations 109 and 112 refer to the Linear Transformer with and without scaling, this model does not use the selective parameter $\lambda_i$; in other words, $\lambda_i = 1$, which results in a mask of all ones. Since the Hadamard product of the mask with the attention matrix does not alter the attention, the scaling does not affect the overall computation.

---

6) xLSTM max stabilizer

We thank you for your careful notice. We have updated our theorem for the xLSTM extension into the bidirectional setting by incorporating the maximum operator in the denominator for scaling in the Appendix C5 line 126 of our new pdf.

---

7) Using just forward recurrence but reversing sequence order

Thanks for the suggestion, we have not tried input sequence reversing but will include it in the final version of the paper.

Regarding Vision-LSTM [1], we considered a similar mechanism to traversal paths to improve locality in masking which we called LION-S (v2) which improved results. For further details of our patch reordering approach, please take a look at Appendix C.8.

We also considered larger sizes of the model which we call SMALL using the VIT-S dimensions. For reference, ViT-S is 21.7M while ViT-T is 5.5M. Please refer to the table in our reply to your first concern.

References:

[1] Alkin, B., Beck, M., Pöppel, K., Hochreiter, S., & Brandstetter, J. (2024). Vision-LSTM: xLSTM as Generic Vision Backbone. arXiv preprint arXiv:2406.04303

---

Once again, we would like to express our gratitude to the reviewer for their thoughtful comments and positive feedback. We are pleased to hear that you enjoyed reading our paper and hope that we have addressed all of your concerns. In case of any concerns remain, we would be more than happy to respond.

Concern 3: Claims about extrapolation, more efficient than Transformers and solving LRA are not surprising

1. Efficiency over Transformers: We would like to clarify that the novelty of our work is not the efficiency compared to Transformers. Rather, we demonstrate that LION enables the efficiency of Linear Transformers in a bidirectional setting, unlocking the efficient potential of the RNN format in this context. Since the efficiency of SSMs and RNNs during inference has already been established, we are proved that LION, being equivalent to an RNN during inference, is inherently efficient.

2. Extrapolation Beyond Context Length Without Additional Positional Embedding: Similar to the previous point, we do not claim LION is the only model able to extrapolate. Since LION-S does not use positional embeddings and can extrapolate beyond the context length, one might ask to see this ability in experiments, and we provide experimental evidence supporting that LION-S can indeed do this (Figure 3).

3. Solving LRA: Unlike the previous point, LION-S is indeed the first Linear Transformer capable of solving the LRA task, which is a significant result (Table 2). It is important to note that even Mamba, despite its strengths, is not able to solve the LRA task, as we have detailed in our response to the Concern 5 of this response.

Benefits of LION: The LION framework and the model LION-S enable fast training and parallelization of transformers while maintaining the advantages of RNN/SSMs during inference, which are both faster and more efficient. It combines the best of both worlds for bidirectional sequence modeling. As evidence, the training times for different models for CIFAR100 on a single NVIDIA-A100 GPU with batch size 1024 are presented below:

This shows that training using attention is significantly faster than SSMs like Hydra, which is considered a state-of-the-art bidirectional SSM.

---

Concern 4: Comparison against vision SSMs

Although LION is not an SSM, for completeness, we compare it with the HYDRA $8$ model, which is one of the best-performing state-of-the-art vision SSMs. We used HYDRA-T to match the parameter count. Results show that LION-S (v2) outperforms HYDRA on CIFAR-100 and performs very similarly on ImageNet. Note that the hyperparameters were taken from the original paper. LION_S (v2) is our extended approach that reorders patches to include more spatial information into calculations (for further details, please refer to Appendix C.8). As these results were obtained during the rebuttal period, we will continue tuning until our final submission. For the results LION-S is using the same parameters as ViT for training.

---

Concern 5: Mamba and S5 results in experiments:

The poor performance of Mamba in the LRA task is well-recognized within the community and has been acknowledged by Albert Gu, the author of Mamba [6], Mamba-2 [4] after community finds that it is not as performant of other SSMs (as stated in this GitHub issue). Additionally, other papers [12] have reported similar performance for Mamba as we observed in our experiments. Regarding S5, we have reported their results from version 1 of their submission, which were also presented in [9]. We have now added the updated version of the results, both below and in Table 2 of our paper.

---

Concern 6: "Transformers can be represented as SSMs"

Thanks to reviewer we have emphasized in the footnote of our background section (on page 3) that softmax-based attention cannot be linearized. However, we would like to highlight that the fact that softmax-based attention does not have an exact finite-dimension equivalent in RNN form has been established in the "Transformers are RNNs" paper [1]. Variants of linear recurrent models, such as Performer [5], approximate softmax-based attention by applying Random Fourier Features.

Additionally, we acknowledge that the sentence "Transformers can be presented as SSM" could be seen as imprecise, especially considering Mamba-2’s footnote. However, the title of Mamba-2, "Transformers are SSMs," can be considered even more imprecise than our statement following the same reasoning.

Concern 7: LRA Text Dataset

We thank the reviewer for their careful attention to the details of the paper. This was a typo, and it has been corrected in the new version of the paper. The Text dataset has a sequence length of 4096.

---

Concern 8: $exp(a_i)-1$ is not clear in equations 27 and 26 (LION-S section)

We would like to emphasize that the terms $exp(a_i)-1$ and $exp(a_i)$ are not arbitrary choices but are derived from the zero-order-hold (ZOH) discretization of a continuous-time system. This discretization method is well-established in signal processing, as noted in the HIPPO [14] paper and commonly referenced in standard signal processing textbooks such as [13]. This ZOH discretization also appears in the Mamba paper, where the discretized parameters of ZOH filtering are presented as follows: $\hat{B} = (\Delta A)^{-1}(\exp{\Delta A} - I)\Delta B$ $\hat{A} = \exp{\Delta A}$ with $\Delta$ representing the time step size. To provide further clarity, we have included a detailed proof of ZOH discretization in Appendix Section B.6. Additionally, we find the reviewers' comment—"we have the sum of the two coefficients summing to, ensuring an attention mask"—unclear. Specifically, the term $(\exp{a_i} - 1)$ does not contribute to the creation of an attention mask, and the phrase "ensuring the mask" lacks precision and clarity. Only the parameter $\lambda_i$ contributes to the attention mask [3,4], and a good example of this is xLSTM [11] where the state parameter $\lambda_i$ and $\gamma_i$ (in equation 4 of our paper) are different and do not sum to one.

---

Concern 9: Adding LION-LIT on LRA

The reason for not initially including LION-LIT as a starting point is that it is the bidirectional version of LinearTransformer, which already appears in the table. However, to address the reviewer's concern, we have now added LION-LIT to Table 2 as a result.

This addition further supports that the choice of LION-S is indeed significantly effective.

---

Concern 10: LSTM and GRU categorization

Thank you for the comment, which helped us clarify the definition and distinction of bidirectionality. We define bidirectional architectures as those that were initially introduced as bidirectional models. As a result, we have revised Table 1 of our paper to reflect this distinction, placing GRU and LSTM in the unidirectional category, as they were originally designed for unidirectional sequence modeling. We have separated these models into a different row of the table and added models like ELMO [10], which were designed as bidirectional RNNs, in a separate row.

---

Concern 11: SSMs dense information processing

Thank you for the comment. We would like to clarify that all linear attention-based models can be represented as RNNs [15], as the linear attention matrix is of rank $d$, where $d$ is the model hidden dimension, while softmax-based attention has rank $L$ due to the non-linear operation (exp) applied to the attention matrix. However, the reason for representing these models as RNNs is not just because of the rank of the matrix. We also want to highlight that studies like Performer [5] approximate softmax attention and can still be represented as an RNN. For clarity of our paper and claims and in response to the reviewer's suggestion, we have removed the sentence in our final manuscript.

---

Since some of the weaknesses are not itemized below, we have mentioned which concerns correspond to the relative questions/weaknesses in our response.

W1: Answered in Concern 5

W2: Answered in Concern 6

W3: Answered in Concern 9

W4: Answered in Concern 10

W5: Answered in Concern 11

---

Q1: Answered in Concern 1 and 2

Q2: Answered in Concern 7

Q3: Answered in Concern 8

Q4: Answered in Concern 2

Q5: Answered in Concern 5

Q6: Answered in Concern 9

Q7: Answered in Concern 10

---

References

$1$ : Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, François Fleuret. arXiv:2006.16236, 2020.

$2$ : Retentive Network: A Successor to Transformer for Large Language Models. Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, Furu Wei. arXiv:2307.08621, 2023.

$3$ : Random Feature Attention. Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A. Smith, Lingpeng Kong. arXiv:2103.02143, 2021.

$4$ : Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. Tri Dao, Albert Gu. arXiv:2405.21060, 2024.

$5$ : Rethinking Attention with Performers. Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, David Belanger, Lucy Colwell, Adrian Weller. arXiv:2009.14794, 2022.

$6$ : Mamba: Linear-Time Sequence Modeling with Selective State Spaces. Albert Gu, Tri Dao. arXiv:2312.00752, 2024.

$7$ : Efficiently Modeling Long Sequences with Structured State Spaces. Albert Gu, Karan Goel, Christopher Ré. arXiv:2111.00396, 2022.

$8$ : Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers. Sukjun Hwang, Aakash Lahoti, Tri Dao, Albert Gu. arXiv:2407.09941, 2024.

$9$ : Liquid Structural State-Space Models. Ramin Hasani, Mathias Lechner, Tsun-Hsuan Wang, Makram Chahine, Alexander Amini, Daniela Rus. arXiv:2209.12951, 2022.

$10$ : Deep Contextualized Word Representations. Matthew E. Peters, Mark Neumann, Mohit Iyyer, Matt Gardner, Christopher Clark, Kenton Lee, Luke Zettlemoyer. arXiv:1802.05365, 2018.

$11$ : xLSTM: Extended Long Short-Term Memory. Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, Sepp Hochreiter

[12]: State Space Models as Foundation Models: A Control Theoretic Overview. Carmen Amo Alonso, Jerome Sieber, and Melanie N. Zeilinge

[13]: Raymond A DeCarlo. Linear systems: A state variable approach with numerical implementation. Prentice-Hall, Inc., 1989.

[14]: HiPPO: Recurrent Memory with Optimal Polynomial Projections: Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Re

[15]: Parallelizing Linear Transformers with the Delta Rule over Sequence Length: Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, Yoon Kim

---

We hope that our clarifications have provided a clearer understanding of our contributions and that we have successfully addressed the reviewer's concerns. In case the reviewer has any remaining concerns, we are more than willing to address them

Thank you for your review. We believe there has been a misunderstanding regarding the contribution of LION and its training paradigm, which have led to a misevaluation of our work. We have clarified LION's contribution and training approach, and below, we address the reviewer’s comments:

Misunderstanding of Our Main Contribution

Thank you for the review and the concerns raised. We would like to clarify that the key contribution of our paper and the LION framework has been misunderstood. LION and its variants (such as LION-S) train like Transformers and do not use parallel scan. This contribution is clearly reflected in our title, Figure 1, Table 1, Section 3, and elsewhere in the paper. Below, we address this point along with other concerns raised by the reviewer:

---

Concern 1: "LION can be viewed simply as SSMs, LION-S is two SSMs summed together":

LION is not simply an SSM; it is a framework where training happens with attention (Theorem 3.3 Eq 25, i.e., $Y = SCALE(QK^\top ⊙ M)V$) and inference is performed using the equivalent RNN (Theorem 3.3 Eq 20-24) in a bidirectional setting. Additionally, LION-S is not simply two SSMs summed together, as this summation does not equate to full attention (as shown in Figure 2, parts b1 and b2, and Equation 7 of our paper). Instead, LION-S combines our framework with a specific selectivity (inspired by RFA[3] and Mamba-2 [4]), resulting in a model equivalent to full masked attention.

Based on the connection outlined in [1], which allows causal Linear Transformers to be trained using attention and inferred as equivalent RNNs (as presented in Equations 4a to 4c of our paper), these models were introduced prior to the development of SSMs [7], and specifically before S5, which introduced parallel scan for SSMs. Causal Linear Transformers and all their variants use Transformer-style training and RNN inference, as demonstrated in Proposition 3.2 of our paper.

Moreover, we would like to highlight that the equation provided by the reviewer:

$y_i^F + y_i^B = q_i^\top (S_i^F + S_i^B - k_i v_i)$

poses a problem because the outputs for the forward and backward recurrences are not available for token $i$ at the same time. Therefore, this equation cannot be simply calculated unless all the $q_i$, $k_i$, and $v_i$ values are stored in memory, which contradicts the fundamental motivation behind using SSMs and RNNs (please look at our response to reviewer tj5X 1st question). Additionally, the term $k_i v_i$ is mathematically incorrect, as these two vectors cannot be multiplied using matrix multiplication. It is also not a dot product, as it should be a matrix representation to correctly express linear attention [1].

---

Concern 2: Difference between LION and SSMs

LION shares similar properties with mentioned Linear Transformer variants (reffered as Transformers as Linear Recurrent Model in Table 2) and has key differences compared to SSMs, as shown below:

1. Linear Tranformers, including LION, use attention matrix parallelization during training, whereas SSMs use the parallel scan algorithm.
2. Linear Tranformers, including LION, are presented as RNNs with matrix-valued hidden states $S^{F/B}_i \in \mathbb{R}^{d \times d}$ during inference, while SSMs have vector-valued hidden states $h^{F/B}_i \in \mathbb{R}^{d}$.
3. Linear Tranformers, including LION, use multi-head attention, while SSMs employ state expansion.

Based on these differences, models like Linear Transformer $1$ , RFA $3$ , and RetNet $2$ can be considered distinct from SSMs, especially since these models were introduced before the formalization of SSMs. Since LION uses attention parallelization during training, it has matrix-valued hidden state during inference, and it uses multi-head attention instead of state expansion, LION is categorized as "Transformer as Linear Recurrent Model".

Moreover, LION extends these models into a bidirectional setting, where training uses attention parallelization (without the parallel scan algorithm), as proven in Theorem 3.3 and shown in Figure 1, and inference is done by a theoretically equivalent bidirectional RNN, as formalized in the same theorem. Additionally, our code is adapted from the original Transformer implementations of ViT and BERT, with modifications made only to the attention block to include masking during training (we have included a torch implementation of the mask in Appendix section C6), removing the positional embeddings, and adding the recurrent inference. The rest of the model architecture has not been altered. We will share the code upon acceptance.

---

• P4: xLSTM, Gateloop, Mamba's S6, GLA-Transformer, and LRU are not included in the paper

  Thank you for your comments. Since our experiments are extensive across multiple domains, we believe the reviewer was referring to the absence of certain models in the LRA task results. We would like to address each baseline in detail:

  • Mamba (S6): We have already included Mamba in the LRA results in Table 2 (line 389 of our paper), and we have now included the new version mentioned by the reviewer.
  • xLSTM: Thanks to the reviewer, we have added the xLSTM results to Table 2 of our paper for the LRA task.
  • GateLoop: We appreciate your comment. However, we have not found results for GateLoop on the LRA benchmark. Additionally, the GateLoop paper does not include experimental results for the LRA task. We are willing to include it if the reviewer can provide a reference of the results.
  • LRU: Thank you for pointing this out. We have added the LRU baseline to Table 2 as requested.

For vision tasks, we would like to emphasize that models such as LRU and GateLoop have not been extended into the bi-directional setting. We include Hydra in our comparison, despite its release as pre-print on arxiv just four months ago. According to the reviewer's guidelines last FAQ point, comparing against Hydra may not be ideal, but we included it to extend our study in the rebuttal. Hydra is one of the few models that claims bi-directionality, which aligns with the scope of our contribution, making it relevant for comparison despite the timing of its release.

Moreover, since Hydra is already outperforming other vision models like Vision Mamba, S4-VIT, and Hyena-VIT, comparing against Hydra provides a comprehensive evaluation, especially given the limited time during the rebuttal.

---

• P5: ""Many of these input-dependent gating methods would outperform most of the Transformer models listed in Table 2 ... Additionally, many of these methods can generate a similar mask when applied bidirectionally (and normalized), which could lead to a result similar to a "Transformer" architecture""

As also mentioned by the reviewer, "generating a similar mask" is the key contribution of our paper. Specifically, we introduce a novel method for generating the mask and extending existing models into their bi-directional setting, as detailed in Appendix C.5. The models we compare against do not have such a mask in their architecture for the bi-directional case, nor do they have the capability to train in a bi-directional setting, which remains the core contribution of our work.

Asking to extend these models into a bi-directional setting and then comparing them against LION-S is not entirely fair, as we would essentially be comparing our own contribution to itself. In any case, both approaches are based on the LION framework, and the central focus of our work is extending existing models into the bi-directional setting, which represents a significant and novel contribution to the field.

LION-S is a representative example of using this framework with a selective mask, demonstrating the potential of gating mechanisms in LION framework.

---

• P6: "Memory comparison with Hydra"

We have added the memory consumption of Hydra in Figure 3. For this comparison, we used the memory efficient implementation in the Official Hydra repo. In below table, we present the inference GPU usages of models in terms of GB. It can be observed that Hydra requires a memory nearly 2.5$\times$ that of LION-S/LIT for high resolutions.

---

Final remarks

We appreciate the reviewer's efforts and feedback, which helped us improve our work by incorporating additional baselines such as Mamba results from the Appendix of xLSTM and LRU. These additions have certainly enhanced the quality of our paper. However, we believe it is unfair to assign a score of 3 and recommend rejection solely based on these points. If the concerns raised have been addressed, we kindly ask the reviewer to consider revising the rating accordingly.

---

Dear reviewer sped,

Thanks for your response. We answer to your points as follows:

• P1: Regarding Mamba's performance on LRA

Thank you for your comment. We would like to highlight that our results for Mamba are taken from [1,6], we have updated Table 2 of our paper based on the results requested by the reviewer, specifically from Table 6 in the xLSTM paper [2]. However, we would like to point out two key differences:

1. Mamba's performance is still lower than that of LION-S.
2. Neither us nor other authors managed to make Mamba solve the PathX problem, while LION-S is able to solve this task. Please note that the Mamba LRA results from xLSTM [2] are incomplete as PathX and Text are missing.

---

• P2: Comparison against vision SSMs

1. Which mask to apply at LION-S(v2): As we have mentioned in line 1498 of our Appendix C.8, we do not choose between these masks and we average two masks given in Figure 7 b and c.
2. Does Hydra have the same mask: We would like to emphasize that this mask, which is our contribution, is specific to our architecture. Since Hydra uses a matrix-form state $A$, we do not see a straightforward way of applying this technique to Hydra. As mentioned in the Mamba-2 paper [7], only the scale-valued $A$ can be represented as masked attention.
3. Hydra results compared to VIT: We appreciate the acknowledgement of our extra efforts towards including the results of this concurrent work (released less than 4 months before submission deadline, please note last FAQ in ICLR 2025 reviewer guidelines). However, we would like to highlight that in the Hydra paper [3], the ViT-Base family is considered. In our experiments, we consider the ViT Tiny and Small families. We reproduced Hydra based on their repository for Tiny scale and reported the findings. Due to our limited resources, we are only able to run Hydra in the Tiny scale.

---

• P3: LION-S uses cumsum (parallel scan)

We would like to clarify that LION-S employs cumsum, that uses a parallel scan, for efficiently creating the mask. On the other hand, SSMs employ the parallel scan as the key ingredient to parallelize their recurrence formula. The main parallelization mechanism in LION-S is equivalent to the one in the transformer, i.e., through computing the attention matrix.

It is worth noting that the cumsum operation is not unique to LION-S. Other prominent Transformer models, such as the SLiM Performer official implementation [5], Llava Hugging Face transformers implementation [4], and the x-transformers library, also use cumsum in their training and they use attention parallelization.

Moreover, we have demonstrated that our training time is significantly faster compared to parallel scan and is comparable to the full VIT training time, as shown in the table provided in our initial response to your concern 3.

Therefore, we believe that the use of cumsum (in our case for mask creation) which is indeed common in many Transformer implementations under the hood and should not be considered neither as applying parallel scan training nor as a reason for rejecting our approach since it does not change the fact that our general training structure follows a transformer training.

I appreciate the author's efforts in providing a comprehensive rebuttal.

The poor performance of Mamba in the LRA task is well-recognized within the community and has been acknowledged by Albert Gu

Mamba's author stated, "we believe it should perform well on data such as text and not as well on data such as images (e.g., the Image/Pathfinder tasks)." However, the recent xLSTM paper (accepted at NeurIPS) reports achieving 99.2% on Pathfinder with Mamba, a significant improvement over the 69.26% reported here. This raises some ambiguity about Mamba's performance expectations.

Additionally, Mamba's author claims that performance on Retrieval is comparable to S4. However, this does not align with the results presented in the current work—Mamba achieves 72.14%, while S4 reaches 89.46%. Furthermore, the xLSTM paper reports a 90.2% performance for Mamba on Retrieval. Given these discrepancies, the Mamba results presented here seem quite unusual.

Comparison against vision SSMs

LION-S (v2) introduces a new masking mechanism designed to encourage local information.

• Figure 7 suggests the use of two different masking mechanisms. How is the decision made on which one to apply?
• Are the Hydra experiments using the same masking mechanism as LION-S (v2)? If not, this could give LION-S (v2) an unfair advantage in comparison.
• The results in the Hydra paper indicate that Hydra significantly outperforms ViT on ImageNet. However, the numbers presented here show Hydra underperforming ViT. Could you provide some insight into why this discrepancy occurs?

Linear Tranformers, including LION, use attention matrix parallelization during training, whereas SSMs use the parallel scan algorithm. torch implementation of the mask in Appendix section C6

The PyTorch implementation of the mask in C6 uses cumsum (cumulative sum). While the paper argues that attention matrix parallelization is employed, there is still a parallel scan actually running under the hood. This scan computes the prefix sums (cumulative sums) in parallel, as an iterative approach would be too slow.

---

n terms of similarities, LION-S utilizes an input-dependent gating mechanism, which aligns it with methods that also employ such gating, including xLSTM, Gateloop, Mamba's S6, GLA-Transformer, and LRU. However, a direct comparison with these methods is notably absent.

While I agree with the authors' assertion that LION-S is not merely a sum of two SSMs, I apologize for framing it that way. However, I find it somewhat misleading that LION-S is described as a Transformer, with the discussion focusing exclusively on comparisons to Transformer-based models. Many of the input-dependent gating methods mentioned above are quite similar to the recurrence mechanisms described in Section 4. Furthermore, many of these input-dependent gating methods would outperform most of the Transformer models listed in Table 2 (e.g., LRU’s performance on Long Range Arena), and vision models based on these methods would likely be more efficient, as indicated by the results in Table 3.

Additionally, many of these methods can generate a similar mask when applied bidirectionally (and normalized), which could lead to a result similar to a "Transformer" architecture. Appendix C.5 demonstrates this in the context of the proposed LION framework. It is also worth noting that the memory usage of these input-dependent gating methods likely leads to GPU usage comparable to what is shown in Figure 3.

Since the primary focus of the paper is LION-S (not LION-LIT, which is based on the Linear Transformer), and since the core component of LION-S—the input-dependent recurrence—closely resembles these existing input-dependent gating methods, it would be natural for the paper to include comparisons with methods like xLSTM, Gateloop, Mamba's S6, GLA-Transformer, and LRU. I appreciate the inclusion of Hydra in the experiments, as I believe it provides a more fair comparison, given that it too is based on an input-dependent gating RNN.

In my view, presenting direct experimental comparisons with these input-dependent gating methods would make the comparisons more comprehensive and fair. For example, including a memory comparison with Hydra (and/or Vision Mamba) in Figure 3 would help provide a clearer and more balanced evaluation.

Dear reviewer sped, thanks for engaging actively in the discussion which we highly appreciate and helped us to add more analysis and details to our work. Since many points have been raised during this rebuttal period, let us start by summarizing the concerns you raised and have been addressed, to then answer to your remaining concerns.

We would also like to emphasize that some of the concerns raised have already been addressed in the appendix of both the original and revised versions of the paper, and we have simply pointed them out for clarity. By addressing these concerns, we believe the quality of the paper has increased. We kindly ask you to acknowledge that the previous concerns have been addressed.

The remaining concerns raised by the reviewer regarding the Vision Mamba and Hydra baselines, variants of LION (such as LION-RETNET), training time comparisons, and the focus of our paper on the LION framework are addressed below:

---

(1) The method referenced in the review was Vision Mamba

Thanks for the clarification. Since the reviewer believes that Vision Mamba is an important baseline to include in our study, we have added the results for Vision Mamba Tiny and Small in our new experimental results in Table 4 of our paper, based on their work [1]. Note that there are salient differences in their training setup, including data augmentation and other novelties, which are tailored towards improving the accuracy (which is not our focus nor our contribution here). However, we will still continue working on reproducing the Vision Mamba results and aim to include the reproduction by the end of the rebuttal period.

Moreover, the key advantage of LION lies in its efficient training with reduced time and memory resources. Our framework enables linear recurrent models like Linear Transformer, RETNET, and selective variants to also benefit from fast/parallelized training of Transformers. Our experiments illustrate that these models perform as well as vision models and outperform VIT in several tasks.

---

(2) Comparing against Hydra provides a comprehensive evaluation. Hydra paper asserts that Hydra significantly outperforms ViT on ImageNet, yet the results presented here do not support this comparison

As we mentioned in our previous response and as reflected in our paper (Table 4), the scale of Hydra that we reproduced is "Tiny," while the only scale represented in the original Hydra manuscript is "Base." Therefore, our results are not in contradiction with the Hydra paper, as they never claimed that "Tiny" size for Hydra outperforms VIT, and it does not exist in their publication. Moreover, after a careful reading, we failed to find a direct comparison of Hydra and Vision Mamba [1] in the Hydra paper.

Furthermore, since Hydra was suggested by reviewer tj5X, and considering the short rebuttal period and our limited resources, we decided to include Hydra. Hydra is more aligned with our study since it is applied to both vision and bi-directional language modeling tasks. For comparison, we have also added the Vision Mamba results from their paper in Table 4 of our manuscript.

---

(3) The paper presents two distinct methodological contributions:

- A framework (LION) that enables the modification of existing models into bidirectional, Transformer-like architectures. - A new linear recurrent model (LION-S) that leverages this framework.

We have revised our main claims in the paper to better align with the study. Now, we introduce LION-S as a representative of our framework with a selective mask, ensuring the focus remains on our framework rather than solely on LION-S. Additionally, we clarified that LION-S serves as a running example illustrating how our framework works for both fixed and selective masks. As represented in our introduction now our main claims are clarified as:

• We propose a theoretical framework LION (Theorem 1), which expresses bidirectional Transformers as bidirectional RNNs, enabling efficient inference for long sequences while benefiting from well-established Transformer training (cf., Table 1).
• Our theoretical framework offers the foundations to transform a wide class of autoregressive recurrent models (cf., Section 2) into their bidirectional counterparts.
• We propose three main running examples of our framework, inspired by prior work, namely:
  1. LION-LIT: Scaled attention without masking, a bidirectional extension of Linear Transformer.
  2. LION-RETNET: Fixed masked scaled attention with scalar and learnable state parameter $\gamma$, an extension of RETNET into the bidirectional setting.
  3. LION-S: Selective masked scaled attention with input-dependent mask $\lambda_i$, inspired by the selectivity of Mamba-2.
• Through extensive experiments in the Long Range Arena, Vision Tasks, and Masked Language Modeling, we have demonstrated the capabilities of the LION framework and the models built upon it, as outlined above.

---

We thank the reviewer for their comments. Below, we have addressed the points raised in their feedback:

---

1. Some equations/notations are not clear enough.

We thank the reviewer for the helpful pointer. We have clarified in the text that in Eq. 6, $a_1/a_2$ are without the mask, and in Eq. 7, $b_1/b_2$ correspond to the masked attention. Additionally, $\lambda_i$ is the scalar version of $\Lambda_i$, as described in Eq. 5a. We have added the exact definition of $\lambda_i$ to the paper and further clarified Equations 6 and 7.

---

LION-Lit performs worse than BERT or ViT, and the gap is not negligible.

Thank you for the important notice. As outlined in Equations 20-25 and lines 336-339, LION-LIT is the bidirectional sequence modeling equivalent of the Linear Transformer. LION-LIT is formulated as Scale($QK^\top$), whereas models like ViT and BERT are formulated with Softmax($QK^\top$). Thus, LION-LIT is equivalent to linear attention but does not apply softmax, as it prevents the the attention matrix from being linearized, as discussed in Theorem 3.3 of our paper (lines 145-154) and in more detail in Appendix B.3. It was observed in autoregressive and unidirectional models [1] that Linear Attention can be less performant than softmax-based attention. However, Linear Attention offers greater efficiency, as it can be expressed in an RNN format, enabling faster inference and lower memory usage [1,2]. To address the reviewer's concern, we have carried out additional experiments on the GLUE benchmark with BERT-large models, which indicate that the gap between LION-LIT and BERT is becoming smaller by increasing the model dimension to LARGE from Base as bellow:

---

3. The advantage of adding selectivity is not well explained. Table 2 only shows LION-S results, but it does not compare it with LION.

Thank you for the suggestion. We have now included the LION-LIT results (which is a bi-directional version of Linear Transformer), which indeed support our claim of improvement by using the selectivity in LION-S. As observed, the Linear Attention variants of linear recurrent models do not perform as well as selective ones, such as Mamba-2 $2$ or Gated-RFA $3$ . By incorporating selectivity, we introduce LION-S. The results of LION-LIT are as follows:

The LION-LIT results are presented in revised Table 2, clearly demonstrating the effectiveness of the selective parameters and their choices.

---

4. The definitions of the architectural elements of figure 1 should be added to its caption.

Thank you for the helpful pointer. The definitions have been added to the caption in our revised paper.

---

5. Is this method limited to bidirectional models? Most generative models are not bidirectional. How does it apply to single-directional models?

We have already evaluated our selective parameter choice (LION-S 1D) in Appendix A.1 for the single-directional scenario of generative language modeling (in our first and latest submission), showing significant improvement over Linear Attention, with performance close to GPT-2 [4] in terms of perplexity (PPL). Additionally, since many variants of Transformers and SSMs have already been explored for autoregressive tasks (as shown in Table 5 of the Appendix), we believe that extending these models to their bidirectional alternatives is the key missing piece, which constitutes the main contribution of our paper. We will repost it here:

As seen, LION-S significantly outperforms Linear Attention and achieves performance comparable to GPT-2.

6. Other than the theoretical improvement in complexity, how much is the improvement in reduction in latency, computation FLOPs, etc? Figure 3 only shows GPU memory for image classification. How does GPU memory change on language tasks?

Thank you for the question. In Appendix A.1, Figure 5 we have shown the latency and GPU memory of different gerneration modes on causal language modelling task. We also added the memory usage graph for non-causal language task in Figure 3. For number of FLOPs theorethical calculations are included in Appendix D.9. The calculations show that while transformer requires $O(L^2+Ld^2)$ FLOPS, LION-S needs only $O(Ld^2)$ with $d$ being the model dimension and $L$ the sequence length. The same calculations apply to the non-causal language task where $L$ corresponds to context length.

---

References

$1$ : Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, François Fleuret. arXiv:2006.16236, 2020.

$2$ : Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. Tri Dao, Albert Gu. arXiv:2405.21060, 2024.

$3$ : Random Feature Attention. Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A. Smith, Lingpeng Kong. arXiv:2103.02143, 2021.

$4$ : Language Models are Unsupervised Multitask Learners. Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever. 2019.

---

We again extend our appreciation for the reviewer's comments and concerns and hope that our efforts to address these concerns have been effective. In case the reviewer has any remaining concerns, we are happy to address them.

---

Dear Reviewer z393,

We hope you are doing well. As today marks the last day of the extended discussion period for ICLR 2025, we would like to kindly request your final feedback on our submission. We sincerely appreciate your thoughtful comments, and we have made every effort to address your concerns in our rebuttal.

If there are any remaining questions or clarifications needed, we would be more than happy to respond. We would greatly appreciate it if you could review the rebuttal and provide your final thoughts by the end of the discussion period.

Additionally, if you feel that we have adequately addressed all concerns, we would greatly appreciate it if you could consider increasing the score accordingly.

Best regards,

Authors

---

Similar to Mamba, we have chosen the selectivity of the proposed bidirectional model based on the connection to ZOH discretization, and our results demonstrate that it works for bidirectional sequence modeling as well. Our claim is not that we invented this approach, but rather that it is a traditional method for sampling continuous signals, which has been successfully used in models like Mamba and other SSMs for causal sequence modeling. We are therefore showing that using LION with this approach also works for well-known bidirectional sequence modeling tasks, such as vision and masked language modeling.

---

3: No Position Encodings required is a consequence of using SSMs

We would like to emphasize again that the primary contribution of LION is extending causal Linear Transformers and their selective variants into a bidirectional setting. In this setting, the model trains like a Transformer using attention and infers like an RNN/SSM. As a result, LION leverages the advantages of SSMs/RNNs during inference, such as the elimination of the need for positional embeddings. While we do not claim this as the novelty of our work, we demonstrate that because LION is equivalent to an RNN during inference, it inherently possesses these advantages.

LION train using attention parallelization (which is faster than scan algorithm, please see our response to question 5) while benefiting from the memory and computational efficiency of RNNs/SSMs during inference. Our experimental results provide evidence that, if the model is equivalent to an RNN during inference, it should inherit the advantages of RNNs/SSMs, such as not requiring positional embeddings and the ability to extrapolate.

We demonstrate that LION and its selective variant LION-S (inspired by Mamba) indeed exhibit these features. Aside from positional embeddings, details about efficiency of LION are also provided at our response to concern 3 of reviewer sped.

---

4: Practicality of Recurrent Inference

Thanks for the comment. We would like to clarify that techniques like chunking can also be applied to LION. Since LION uses a scalar-valued state parameter, it is $d$ times more memory-efficient than Mamba and comparable to Mamba-2 with a scalar value $a_i$. This efficiency allows LION to achieve $O(Ld)$ memory requirements without needing chunking. However, chunking can still be applied for parallelization and further acceleration during inference when memory demands are high.

As you can see, LION enables Linear Transformers to be extended into a bidirectional setting, while retaining all the advantages of linear attention during inference, such as chunk parallelization, and allowing for parallelized training using attention.

The key advantage of LION over other SSMs that use parallel scan is that LION leverages attention parallelism during training, which is significantly faster than parallel scan, as shown in the table in our response to Question 5 of your review. This results in LION being much faster to train compared to parallel scan-based approaches.

---

5: Experimental Validity Concerns

• GLUE Scores: In our initial results we had employed the original pretraining and finetuning recipes from the M2-BERT repository. We have updated our results with the pretraining and finetuning recipes employed in the M2 models for the parameter-matching BASE and LARGE sizes. We would like to highlight that this repository differs from Hydra's recipe, which is based on MosaicBERT.
  • It is worth noting that these hyperparameters are highly tuned for the M2 model family. For the LARGE models, we even had to reduce the learning rate due to divergence during training, see Appendices D.4-D.5.
  • Similarly, Hydra [3] tunes the hyperparameters specifically for their architecture by exploring different learning rates, weight decays and number of epochs for each GLUE task.
  • Our computational resources do not allow for such expensive tuning. In our new results, we can observe that both the BERT and LION models obtain respectively a GLUE score of $82.25$ and $80.31$ for the BASE (110M) size and $82.95$ and $81.58$ for the LARGE (340M) size. For additional results, please check Table 3 and Appendices D.4-D.5. The BERT large results are: |Model|MLM Acc.|MNLI|RTE|QQP|QNLI|SST2|STSB|MRPC|COLA|Avg.| |-|-|-|-|-|-|-|-|-|-|-| |BERT|69.88|85.68|67.44|89.9|91.89|93.04|88.63|90.89|56.14|82.95| |LION-LIT|67.11|83.73|57.18|89.85|89.93|91.86|88.02|90.18|55.36|80.76| |LION-S|69.16|84.38|57.69|89.57|90.30|92.93|87.68|90.57|59.54|81.58|

This demonstrates how closely LION-S aligns with the original BERT, with the gap smaller in the larger model hidden size ($d$), while LION-S offers significant inference speed benefits.

• Vision Task Comparison:
  • We agree with the reviewer that larger scale models would strengthen our paper. Since pre-training is a compute intensive task, our experiments focus on smaller scales. After the initial submission, we expanded our experiments from 5.5M (Vit-T) to 21.7M (ViT-S). In the below table and in the updated Table 4 of our paper, you can find the results with larger models. Our results show that as the model size increases, the gap between LION-S and ViT gets less and less significant from 2.28% to 1.33%. |Model|ImageNet| |-|-| |ViT-S|72.19| |LION-LIT (SMALL) |69.62| |LION-S (SMALL)|70.86|
  • To further improve the performance of LION architecture, we introduced LION-S (v2) that changes the order of patches to capture spatial information. For details of the implementation, please see the Appendix C.8.
  • Additionally, we have tried distillation to tiny LION-S model using DeiT framework [4] and we have observed that distilled LION-S can outperforms ViT-T. We believe that this result further proves the capabilities of the LION architecture. |Model|ImageNet| |-|-| |ViT-T|70.23| |LION-S (TINY)|67.95| |LION-S (v2) (TINY)|69.22| |LION-S-Distil (TINY) |70.44|
  • Moreover, we have added the training times for each model using the CIFAR100 data, highlighting that training with attention (like ViT & LION) is significantly faster than using parallel scan. This underscores the motivation for training like a Transformer while inferring like an RNN. |Training Strategy (Model)|Time (s)/Epoch| |-|-| |Attention (VIT) |24.6| |Attention (LION-S) |35.8| |Attention (LION-LIT)|26.6| |Parallel Scan (Hydra) |43.4|

---

References

[1]: Retentive Network: A Successor to Transformer for Large Language Models. Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, Furu Wei. We

[2]: Rudolph Emil Kalman. A new approach to linear filtering and prediction problems. 1960.

[3]: Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers. Sukjun Hwang, Aakash Lahoti, Tri Dao, Albert Gu

[4]: Training Data-Efficient Image Transformers & Distillation Through Attention. Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, Hervé Jégou. arXiv:2012.12877, 2021.

[5]: Raymond A DeCarlo. Linear systems: A state variable approach with numerical implementation. Prentice-Hall, Inc., 1989.

---

We hope that through our evaluations, we have been able to address the concerns raised by the reviewer. In case the reviewer has any remaining concerns, we are happy to address them.

We thank the reviewer for their feedback. Below, we have addressed the concerns raised and provided clarifications that we believe may help resolve any misunderstandings and better highlight the contribution of our work.

---

1: Recurrence Form of Linear Attention is a straightforward extension

Thanks for the comment. We would like to stress that the reviewer might have missed key details, causing a misevaluation. For this purpose, we would like to clarify:

• While our approach may seem straightforward in retrospect, it is not an obvious extension, as it has not been formally addressed before. Theorem 3.3 and the proof in Section 3 provide key insights into expanding matrices to achieve bidirectional equivalence to the RNN format.
• We consider the straightforward nature of our approach as an advantage, making it easier to understand. However, below we show that the naive extension of Linear Attention (which we believe what the reviewer has in mind) has significantly higher memory requirements than LION, rendering LION a non-obvious extension.
• In short, compared to the naive extension (cf., below), LION scales linearly with the model dimension $d$, whereas the naive extension and other alternatives for bidirectionality (as shown in Table 1 of the revised PDF) exhibit memory demands that scale with $d^2$.

Naive extension vs LION's Linear Memory Requirement Relative to Model Dimension:

Given this review, we see that an important benefit of LION, regarding its memory requirement, particularly in comparison to the naive extension of Linear Transformers into a bidirectional setting, has not been sufficiently highlighted. The following formulation for Linear Attention, presented in the "Transformers are RNNs" section, serves as the basis for comparison:

$S_i = S_{i-1} + k_i v_i^\top$

$z_i = z_{i-1} + k_i$

$y_i = \frac{q_i^\top S_i}{q_i^\top z_i}$

A naive extension of this approach, accounting for both forward and backward states, would be:

$S_i^{F/B} = S_{i-1}^{F/B} + k_i v_i^\top$

$z_i^{F/B} = z_{i-1}^{F/B} + k_i$

OUTPUT: $y_i = \frac{q_i^\top (S^F_i + S^B_i - k_i v_i^\top)}{q_i^\top (z^F_i + z^B_i - k_i)}$

We demonstrate below that this formulation is highly inefficient in terms of memory.

Since the forward and backward states for token $i$ are not available simultaneously (except for the middle token in even-length sequences), both the forward and backward hidden states $S^F_i$ and $S^B_i$ ($O(2d^2L)$), and both scaling states $z^F_i$ and $z^B_i$ ($O(2Ld)$). Additionally, the memory required to store the output of each token ($O(Ld)$) results in a total memory demand of $O(2d^2L + 3Ld)$.

Even with the memory allocation strategy described in Appendix B.5 of our paper, the memory of above naive extention still scales quadratically with $d$, which contradicts the core motivation of using bidirectional RNNs. The goal of bidirectional RNNs is to store only one output of dimension $d$ for each token, updating the state without requiring quadratic memory demands of $O(Ld^2)$.

To address this, we store the states $c^{F/B}_i$ as scalar values and $y^{F/B}_i$ as vectors of dimension $d$. This approach only requires $O(Ld + L)$ memory units, which is significantly more efficient than the naive extension and scales linearly with $d$ (Also added to Table 1 of our paper).

In retrospect, LION is easy to understand (i.e., perhaps straightforward but in a good sense) and yet is highly effective (but non-obvious). This is our main contribution (which is also extensible): extending existing autoregressive recurrent models, such as Linear Attention (LA) or RetNet $1$ , into a bidirectional setting, where the attention mechanism is theoretically equivalent to a bidirectional RNN. LION achieves low memory requirements during inference, scaling linearly with the model dimension $d$.

---

2: LION's Continuous-Domain Selectivity is derived in Mamba

Thank you for the comment. We would like to clarify that this discretization is a zero-order hold (ZOH) discretization of continuous-time dynamic systems $2$ , which is also used in Mamba. However, as mentioned in the HIPPO paper on page 21, this discretization was derived much earlier (before Mamba) and has been well-established in signal processing community and can also be found in relevant textbooks [5].

For better understanding of the reader, we have included an explicit mention of this fact along with the proof of this discretization for both continuous state-space dynamics and transformer dynamics in Equations 26 and 27 in the appendix B.6 of our paper for completeness.

We would like to gently remind you that the author-reviewer discussion period for our manuscript is approaching.

We appreciate your valuable contributions. We believe that our responses address the concerns raised. If you have any further suggestions or concerns, please let us know.

Best,

Authors

---

Dear Reviewer tj5X,

We hope you are doing well. As today marks the last day of the extended discussion period for ICLR 2025, we would like to kindly request your final feedback on our submission. We sincerely appreciate your thoughtful comments, and we have made every effort to address your concerns in our rebuttal.

If there are any remaining questions or clarifications needed, we would be more than happy to respond. We would greatly appreciate it if you could review the rebuttal and provide your final thoughts by the end of the discussion period.

Additionally, if you feel that we have adequately addressed all concerns, we would greatly appreciate it if you could consider increasing the score accordingly.

Best regards,

Authors

---