We thank the reviewer for their feedback and for recognizing our work as having significant impact and being among the first to present bi-directional sequence modeling for linear transformers in such detail.

The reviewer's feedback highlights the following key concerns:

• 1-3) Clarity of equations and explaining more on motivation behind bidirectional modeling.
• 4-6) Experimental results of LION compared to baselines.
• 7-10) Questions regarding Figures and some more ablations.

We address this in detail in the responses to each point below.

---

1) Presenting the idea in more clear way

We thank the reviewer for suggesting the inclusion of $Y = ((1 + I) \odot QK^\top)V$ as the parallel form of a naive bidirectional linear transformer. As suggested, we will incorporate this clarification to better highlight the limitations of naive bidirectional linear transformers and to further clarify LION’s formulation. This formulation suffers from imbalanced attention along the diagonal due to double counting introduced by the identity matrix $I$. In contrast, LION uses $Y = (QK^\top)V$ with correction terms in its RNN formulation to avoid this issue in the forward and backward passes.

---

2) Motivation behind bi-directional modeling

We thank the reviewer for recommending additional references on bidirectional modeling. As many tasks beyond next-token prediction (e.g., biomedical, image, and signal domains) benefit from bidirectional modeling we will highlight this with the suggested references.

---

3) Why not using Transformers instead of LION if SSMs are too slow

We agree with the reviewer that Transformers train faster than SSMs on bidirectional tasks while maintaining strong performance. However, their inference cost and memory requirements even by considering technical efficiency improvements remain theoretically quadratic with sequence length ($O(L^2)$). In contrast, SSMs offer linear inference complexity ($O(L)$), making them valuable for constrained resources. LION addresses this gap by building the first framework for bidirectional modeling with linear transformers combining the fast training of softmax transformers with the inference efficiency of SSMs, as shown in Figure 2 and Table 2.

---

4) LION's speed comparison with baselines and using custom kernels

Our main goal was to keep LION’s implementation simple and easy to integrate, which is why we relied on a pure PyTorch implementation. Notably, even for Transformers, custom CUDA implementations for bidirectional tasks show more modest speedups compared to causal tasks. As also shown in Table 1 of the FlashAttention paper [1], the speedup gain for BERT is only 1.17× compared to 3.5× for GPT, since bidirectional attention lacks the sparsity introduced by causal masks.

Importantly, even without custom kernels, LION achieves significantly faster training than SSMs like Vim and Hydra, which do rely on custom kernels for speedups. While we fully agree with the reviewer that custom implementations could further improve speed, our results already show that the pure PyTorch version outperforms existing SSMs.

---

5) Hydra's performance with different token sequence orders

We trained Hydra with different sequence orders at the Tiny and base scale (times are reported relative to a Transformer of the same size):

Hydra benefits from using different scans at the Tiny scale but still fails to surpass its original version, even with extensive tuning. This approach does not scale well to the base model, likely due to Hydra’s sensitivity to its decay factor. In contrast, LION-D$^\natural$ with multiple scans outperforms both Hydra variants at the Tiny scale, highlighting LION’s stability during training. Plus, unlike LION, which integrates multiple scans directly into its attention mask $M$, Hydra requires an extra pass in the rotated direction, doubling its training time.

---

6) LION is faster than Hydra, but delivers worse performance.

In most scenarios (except for image classification at the base scale) LION delivers competitive performance. For instance, LION-D$^\natural$ outperforms Hydra at the Tiny scale (Table above) by 3% and Small scale (Table 1 of our paper) by 1% in image classification while being ~2× faster to train. Similarly, LION-S is 3× faster to train on the MLM task, with only a small 0.1% gap in GLUE score compared to Hydra. We believe that even with these marginal accuracy differences, LION’s training speed gains remain highly valuable.

---

7) Figure 2 shows linear scaling for the RNN formulation. However, the curve does not look like a straight line.

In Figure 2, we report memory usage with respect to image resolution, which refers to the number of pixels (or tokens) per row. Since sequence length $L \approx \text{res}^2$, LION-RNN scaling quadratically with resolution means LION scales linearly with sequence length $L$. In contrast, the Transformer in Figure 2 scales quartically ($O(res^4)$) with resolution (i.e., quadratically with sequence length). We use resolution on the x-axis as it is more intuitive for images, same visualization also is used in Vim [2] (Figure 1 c). Linearity is also clear in our MLM results in Figure 8 (right) in Appendix, where LION shows linear scaling and the Transformer quadratic scaling.

---

8) The RNN formulation only requires storage of the state which is independent of the sequence length.

RNNs have constant memory in causal tasks (particularly in autoregressive settings) as each token depends only on the forward hidden state $S^F_i$. In contrast, for bidirectional tasks, outputs depend on both $S^F_i$ and $S^B_i$, which are not available simultaneously. This requires storing all hidden states from both directions, leading to linear memory in sequence length $L$. While Section A.2 of our paper describes a strategy to reduce this, memory remains linear. This is also consistent with Vim [2] in Figure 1-c where its memory scales linearly.

---

9) How well do non-bidirectional and alternating models perform?

We investigate the importance of full bidirectional modeling in Appendix Section C.15, where we compare LION variants using only forward, only backward, alternating forward-backward layers, and full bidirectional processing on CIFAR classification.

As shown in Table 16, LION with full bidirectionality outperforms both unidirectional variants by up to 7%, and also surpasses the alternating forward-backward approach by 3.6%. This highlights the clear benefits of bidirectionality in each layer, rather than alternating directions or using only a single directions. We copy the results from paper below:

---

10) What chunk-size(s) are/were used to produce Figure 1 and how was this chosen?

We demonstrate how to choose the chunk size to balance memory and speed in Section C.16 of our Appendix, "Effect of Chunk Size on Speed and Memory." Figure 11 illustrates this trade-off, with the optimal range falling between 64 and 128, as used in the LION variants (LIT,D,S) for image classification tasks in Figure 1. We will bring this Figure to the main body thanks to the reviewers notice.

---

References

[1] FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré

[2] Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model

We sincerely thank the reviewer for their feedback and questions which helped us improve our manscript. Bellow we will answer the remaining questions:

---

1) Imbalanced Attention

Thank you for the pointer. Imbalanced attention indeed rise an issue for linear transformers. For imbalanced attention the diagonal elements of the attention mask become $M_{ii} = 2$, while off-diagonal entries satisfy $M_{ij} \leq 1$ for case of presence of decay (D,S) or $M_{ij}=1$ for case of no decay (Lit), creating a discrete gap between a token’s attention to itself versus others as opposed to smooth transition from diagonal to off diagonal elements in LION. In contrast, balanced attention ensures that all entries satisfy $M_{ij} \leq 1$ enabling a smooth transition from main diagonal to other elements.

Thanks to your suggestion we have tested effect of imbalanced attention in small-scale CIFAR image classification for LION, and the results are shown below:

As observed, Imbalanced Attention (double counting in the forward and backward paths) leads to worse performance compared to LION with balanced attention. We will add this as part of the motivation for balanced attention, along with the formulation suggested by the reviewer.

---

2) Adding the Motivation behind Bidirectionally in Introduction

Exactly, we will indeed include the following paragraph in the introduction to highlight the importance of bidirectionally:

"Several real-world tasks are inherently bidirectional and benefit from bidirectional sequence modeling. Examples include DNA and protein modeling [52], computer vision [51], and biological and chemical sequence modeling [53]. In these domains, bidirectional models often outperform causal ones, for instance, Bi-Mamba outperforms Mamba in DNA modeling [52]. This motivates the development of architectures specifically designed for bidirectional sequence modeling."

If you have any suggestions for making the motivation more clear, we are happy to adjust.

---

3.1) Memory Efficiency Gap being Addressed by Hydra to some Extent

We agree with the reviewer that bidirectional SSMs like Hydra and Vim address the gap in memory efficiency during inference to some extent. However, LION aims to bridge the gap in slow training speed of SSMs compared to transformers while retaining their inference efficiency. LION uses full attention, enabling fully parallel and faster training. In contrast, Vim (based on Mamba) relies on parallel scan, which involves recursion and leads to significantly slower training. Hydra (based on Mamba-2) adopts a chunkwise parallel training strategy (aka SSD) because the full attention form for Mamba2 is noted to be unstable (as also mentioned in the Mamba-2 blog post). While Hydra trains faster than Vim, since chunkwise parallel is theoretically faster than scan-based approaches, it remains slower than the full attention used in LION and softmax Transformer.

---

3.2) $O(L^2)$ Figure Relevance Despite FlashAttention and KV Caching

This is a great question. KV caching is primarily effective in causal language modeling, where past tokens remain fixed. In contrast, for bidirectional models, for every token the attention is computed over the entire sequence, so caching does not offer efficiency gains since the whole sequence attends to each token. Regarding FlashAttention, we included in our experiments as presented in ViT-Chunk (green line) in Figure 1, which uses chunked queries to reduce memory usage in inference. While it improves over naive softmax attention, it is still less memory efficient compared to SSMs and LION in RNN form, especially on longer sequences. Similar finding is also shown in the Vim [2] (Figure 1c), where FlashAttention remains less memory efficient than Mamba in inference.

---

4) Representation of the Results

The reviewer's representation of the results is correct. LION achieves significantly faster training than Hydra while maintaining comparable performance, marginally lower at larger scales and higher at smaller scales. We will clearly state this in the Experiments section. Additionally, an important part of LION’s contribution and novelty lies in its theoretical formulation, which enables mapping several models to their bidirectional counterparts, as shown in Table 1 as also noted by the reviewer which has not been studied to this detail.

---

Once again, we sincerely thank the reviewer for their valuable feedback and insightful questions. If there are any further questions, we would be happy to address them.

We highly thank the reviewer for considering our paper timely and valuable, the first to study linear transformers for bi-driectional setting, and for describing it as well-written. Below, we address the only major weakness and the minor weaknesses raised by the reviewer.

---

Major Weakness on NLP Experiment

We thank the reviewer for their feedback which helped us improve our results. Our setup follows M2-BERT (released in 2024 and relatively recent) for two key reasons:

1. It also serves as the benchmark and training recipe used in Hydra and previous works on bi-directional sequence modeling [1,2,3].
2. It offers a more GPU-friendly training setup. Due to resource constraints, we could not adopt heavier training configurations such as NeoBERT (2.8TB dataset and 4k sequence length) or RoBERTa (500K steps, 8K batch size), and thus followed the lighter M2-BERT recipe.

We chose the GLUE benchmark as a practical and commonly used evaluation for models at this scale [1,2]. While larger benchmarks may offer additional insights, GLUE remains appropriate given our scope and resources.

We thank the reviewer for the suggestion and have incorporated a longer sequence length into our experiments. We increased the training length to 512 for base scale, as suggested. Results are shown below:

The above results were obtained on sequences of length 512 in a single run, without any hyperparameter tuning, due to time constraints during the rebuttal period and limited GPU resources. All LION variants showed improved performance in GLUE score compared to the shorter 128-length setting reported in Table 3 of the paper. This demonstrates that LION has strong potential for handling longer context lengths, even without additional tuning.

---

Minor Weaknesses

• "LION is up to 10× faster than SSMs" is too strong: We will revise this statement to clarify that LION is 10× faster than Vision Mamba and 2× faster than Hydra, ensuring the claim is precise and accurately reflects the comparisons.
• New models relying on RoPE generalize better in context extrapolation: Context extrapolation is a general advantage of linear transformers, and LION inherits this property as well. We include Figure 8 in the Appendix to illustrate this. While we fully agree with the reviewer that modern positional embeddings like RoPE, YaRN, and ALiBi also support extrapolation, the focus of our study is not to compare positional encoding methods. Therefore, in Figure 8 we simply aimed to show that LION can extrapolate beyond trained sequence length, due to its use of decays instead of fixed positional embeddings.

---

Minor Suggestions

• FlashAttention usage: Thank you for the pointer. We will clarify in the final version that FlashAttention is used for the softmax Transformer baseline in our experiments.
• Bringing Figure 5 to the main content: We appreciate the suggestion and will move Figure 5 to the main paper in the final submission.
• LION also refers to an optimizer: Thanks for the pointer. We are aware of the potential confusion. However, naming models after animals is common in the linear transformer community, and there is no prior model named LION in this context. Additionally, since the domains differ (sequence models vs. optimizers), we believe confusion is unlikely.

---

Questions

• Including MTEB benchmark: Due to time constraints and our limited compute resources, we were unable to support this benchmark. However, we agree it is a valuable direction for future work.

• Increasing sequence length in NLP: Thank you for the suggestion. We have ablated the 512-token sequence length for LION models in GLUE task and included the findings in the table above.

---

References

[1] Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers, Hwang, S., Lahoti, A., Dao, T., & Gu, A. (2024).

[2] Benchmarking and building long-context retrieval models with loco and m2-bert, J Saad-Falcon, DY Fu, S Arora, N Guha, C Ré

[3] Monarch Mixer: A Simple Sub-Quadratic GEMM-Based Architecture Daniel Y. Fu, Simran Arora, Jessica Grogan, Isys Johnson, Sabri Eyuboglu, Armin W. Thomas, Benjamin Spector, Michael Poli, Atri Rudra, Christopher Ré

We thank the reviewer for finding LION framework novel, general and effective in practical scenrios bellow we will address the questions and concerns mentioned in the review.

---

1) Scalability of LION to longer sequences and models sizes

Due to computational constraints and alignment with prior work [1,2,3], our experiments focus on models with ~100–300M parameters. We expect LION to scale well on non-causal tasks with larger models. Similarly, while we couldn’t train on longer sequences (>16K), LION’s strong LRA performance suggests effective generalization to longer inputs.

---

2) Speed comparison fairness

Despite using only PyTorch, LION still outperforms Hydra and Mamba in speed, highlighting the efficiency of its formulation. We chose PyTorch to keep LION's implementation simple and easy to integrate. We agree with the reviewer that custom kernels could offer additional speedups and represent a promising direction for future work. LION already achieves faster training even with a pure PyTorch implementation.

---

3) Experimental details

Thank you for the pointer. If there is any specific part of the appendix you would like to see included in the main body, please let us know and we would be happy to incorporate it. Due to space constraints, we focused on key content in the main body. For the final submission, we plan to bring in additional details from the appendix, as the NeurIPS 2025 Call for Papers allows one extra page.

---

4) Performance-efficiency tradeoff

LION achieves performance close to Transformers on most tasks and even outperforms them at smaller scales, i.e., in image classification with Tiny models (Table 12, Appendix). It also solves LRA tasks from scratch, where Transformers are not able to.

In general SSMs and linear transformers may slightly lag behind Transformers in accuracy, but their inference efficiency and flexibility to switch between attention, RNN, and chunkwise forms make them particularly valuable in practice, this is the core motivation behind LION. We believe that more fine-tuned architectural modifications, such as convolutions over $Q, K$ and additional gates, as used in DeltaNet (ablation of Table 1) [6], could further enhance LION's performance. However, we intentionally used a simple backbone to isolate and highlight the effect of LION’s attention and sequence mixing block to reflects more its generality across different decays.

---

5) Minor Issues

We thank the reviewer for pointing out the minor issues related to typos and sentence clarity. We will address these to improve the overall readability in the final version.

---

6) Questions

• 1) Performance on longer sequences As mentioned above, we believe LION can scale to longer sequences and larger model sizes; however, due to resource constraints, we were unable to explore this in our current work.

• 2) Long training times of Vim and Hydra Thanks for the sharp observation. The slow training time of Vim, even with custom kernels, is due to the use of parallel scan in Mamba’s [4] training algorithm, which is less efficient on GPUs compared to attention. This issue has been addressed in Mamba2 [5] by using a scalar decay for $A$ and adopting a chunkwise parallel training algorithm (aka SSD), significantly improving training speed (which is also shown in details in Mamba2 blogpost). Hydra, built on Mamba2, benefits from this and trains faster. However, both still rely on recurrent processing over the state (also stated in related work of LION and DeltaNet [6]) and are not as fast as attention, which is a key motivation behind LION, which aims to match the training throughput of softmax attention.

• 3) Custom kernel implementation of LION As we addressed this concern in details in concern 4 above, LION is implemented solely in PyTorch. Our goal is to keep LION simple and fast to use within standard PyTorch workflows.

• 4) Figure 8 (left) does not include the Hydra Thanks for the pointer we will add Hydra to Figure 8 (left) and it also can extrapolate alike LION.

---

References

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

[2] MosaicBERT: A Bidirectional Encoder Optimized for Fast Pretraining Jacob Portes, Alex Trott, Sam Havens, Daniel King, Abhinav Venigalla, Moin Nadeem, Nikhil Sardana, Daya Khudia, Jonathan Frankle

[3] Monarch Mixer: A Simple Sub-Quadratic GEMM-Based Architecture Daniel Y. Fu, Simran Arora, Jessica Grogan, Isys Johnson, Sabri Eyuboglu, Armin W. Thomas, Benjamin Spector, Michael Poli, Atri Rudra, Christopher Ré

[4] Mamba: Linear-Time Sequence Modeling with Selective State Spaces Albert Gu, Tri Dao

[5] Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality, Tri Dao, Albert Gu

[6] Parallelizing Linear Transformers with the Delta Rule over Sequence Length", Yang et al. NeurIPS 2024

We thank the reviewer for the positive feedback on our work and for recognizing its novelty, practical efficiency, and comprehensive experiments. We also appreciate the thoughtful suggestions, and below we address the remaining points:

---

1. LRA Experiments: LION solves long-context tasks in LRA, while Transformers fail to improve beyond random guessing, including the Path-X with at 16K sequence length. The $✗$ symbol indicates failure to improve random guessing, not an OOM error.

2. Speed & ViT Implementation: Thanks for the pointer. ViT and DeiT baselines follow the official DeiT [1] repo and recipe, using the original Python implementation. We will move these details from the Appendix to the main paper for clarity.

3. Experimental Details: We appreciate the suggestion and fully agree. We will use the extra camera-ready page to bring more details regarding implementations and experimental results.

Finally, we sincerely thank the reviewer for their constructive and positive feedback, which has greatly helped us improve our paper.

---

References

[1] Training data-efficient image transformers & distillation through attention: Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, Hervé Jégou

We thank the reviewer for their detailed feedback and for appreciating our work, particularly noting its rigor, technical quality, and clarity. Below, we address the main concerns raised.

---

1) Design of decay terms

We like to clarify that, both the selective and fixed decay factors are learned during training: the fixed decay uses non-input dependent $\lambda = \sigma(a)$ with a learned scalar $a$, while the selective decay is input-dependent, defined as $\lambda_i = \sigma(Wx_i)$ with learned weights $W$.

Regarding performance, we find the choice of decay factor to be task-dependent. The selective variant (LION-S) generally outperforms the fixed variant (LION-D), by 0.25% on masked language modeling (Table 3) and 0.44% on LRA (Table 4),likely due to its greater representational power and inductive bias. In contrast, LION-D performs slightly better in image classification, where fixed decay appears sufficient to capture relative token distances in signal and image modalities [1].

As the reviewer noted, different linear transformers use different decay parameterizations, as shown in Table 1 of our paper and Table 2 of DeltaNet [2]. For LION, we selected three diverse decay forms to demonstrate the generality of our framework. We believe each decay choice introduces different inductive biases, making them suitable for different tasks, for example, DeltaNet performs well on retrieval. We believe that evaluating different decay choices is a valuable direction for future work.

---

2) Experiment takeaways

DeiT is second fastest while outperforming LION D and S

The main takeaway from Table 2 is that LION, and more broadly, linear transformers are competitive with softmax-based transformers like DeiT in both training speed and accuracy. However, they offer a significant advantage in inference efficiency. As shown in Theorem 3.1, LION can be expressed in RNN form, enabling linear-time inference ($O(L)$) compared to the quadratic complexity ($O(L^2)$) of softmax transformers. This is illustrated in Figure 2 and 8.

We agree that Table 2 alone may not fully capture LION's advantage over softmax transformers. Since we aim to highlight three key aspects, 1) accuracy, 2) memory usage, and 3) training time, presenting all of them in a single plot would reduce clarity. Therefore, we separated them: Table 2 focuses on performance and training speed, while Figure 2 emphasizes inference memory efficiency and scalability. Together, they demonstrate that LION combines the strengths of transformers in speed and accuracy with the inference efficiency of RNNs.

Chunking Tradeoff is Generic; Figure 1 Serves as Motivation Rather Than a Result

We agree with the reviewer that the tradeoff between RNNs, chunking, and attention is generic, which also applies to bidirectional models, as noted in the background section of our paper. Figure 1 is primarily intended to motivate the value of supporting multiple representations depending on resource constraints. Thanks to the reviewer’s suggestion, we will move this figure to Section 3, where it better fits as a motivation.

Final Conclusion from Experiment Section

Our main contribution, as noted by the reviewers (particularly in the second point), is the novel formulation of LION as the framework for extending linear transformers to the bidirectional setting while supporting and analyzing tradeoffs across different regimes.

Our experiments shows that LION works across multiple modalities. Key takeaways of experiments include:

• LION achieves training speed and accuracy close to softmax transformers (Table 1), while being significantly more efficient at inference.

• LION matches or exceeds the inference efficiency of SSMs (Figure 2), while training much faster.

---

3) Discussion on DeltaNet

We appreciate the reviewer's sharp observation on this point. Our theorem is specifically focused on linear transformers with diagonal decay, and models outside this class generally cannot be expressed in the full attention formulation. However, for the special case of DeltaNet, our theorem supports an equivalent bidirectional RNN formulation with LION-style correction terms to avoid double counting, as follows:

$S_i^{F/B} = S_{i-1}^{F/B}(I - \beta_i k_i k_i^\top) + \beta_i k_i v_i^\top, \quad y_i^{F/B} = q_i^\top S_i^{F/B} - \frac{1}{2} q_i^\top k_iv_i^\top$

The chunkwise parallel form of DeltaNet (Eqs. 8–11 of the DeltaNet paper) can also be extended to the bidirectional setting by removing causal masks and applying Theorem 3.2 for the chunkwise form of LION.

However, a key bottleneck in DeltaNet arises in its full attention formulation. As described in the original paper (page 6, section “Fully Parallel Form for DeltaNet”), the attention is computed as:

$A = QK^\top T$

Computing the matrix $T$ (defined in equation 10 of DeltaNet) is costly, which directly citing from paper: "$T$ requires a matrix inverse (Eq. 10), which scales cubically with sequence length without further algorithmic changes. Due to the above we avoid using the fully parallel form for training DeltaNet". The same limitation applies in the bidirectional case, where the inverse undermines the fast-training advantage of LION.

We thank the reviewer for raising this point and will include the discussion, along with a note on models outside $\text{TC}^0$ with non-diagonal decay, in the limitations section.

---

References

[1] Mamba: Linear-Time Sequence Modeling with Selective State Spaces Albert Gu, Tri Dao

[2] Parallelizing Linear Transformers with the Delta Rule over Sequence Length", Yang et al. NeurIPS 2024]