Linear Attention for Efficient Bidirectional Sequence Modeling (2502.16249v1)

Abstract: Transformers with linear attention enable fast and parallel training. Moreover, they can be formulated as Recurrent Neural Networks (RNNs), for efficient linear-time inference. While extensively evaluated in causal sequence modeling, they have yet to be extended to the bidirectional setting. This work introduces the LION framework, establishing new theoretical foundations for linear transformers in bidirectional sequence modeling. LION constructs a bidirectional RNN equivalent to full Linear Attention. This extends the benefits of linear transformers: parallel training, and efficient inference, into the bidirectional setting. Using LION, we cast three linear transformers to their bidirectional form: LION-LIT, the bidirectional variant corresponding to (Katharopoulos et al., 2020); LION-D, extending RetNet (Sun et al., 2023); and LION-S, a linear transformer with a stable selective mask inspired by selectivity of SSMs (Dao & Gu, 2024). Replacing the attention block with LION (-LIT, -D, -S) achieves performance on bidirectional tasks that approaches that of Transformers and State-Space Models (SSMs), while delivering significant improvements in training speed. Our implementation is available in http://github.com/LIONS-EPFL/LION.

Linear Attention for Efficient Bidirectional Sequence Modeling

The paper under review presents a comprehensive paper on the development and application of linear transformers for bidirectional sequence modeling. This development is significant because linear attention paradigms have traditionally been limited to causal sequence modeling, and therefore this paper extends the utility of linear transformers by formulating a bidirectional framework. The authors examine the inherent challenges of extending linear attention models to bidirectional scenarios, shedding light on mathematical and practical solutions.

Key Contributions

1. Bidirectional RNN Framework for Linear Transformers: The authors introduce a framework that enables linear transformers to operate in a bidirectional manner, equivalent to full linear attention. This involves constructing a bidirectional recurrent neural network (RNN) that leverages linear attention mechanisms without the quadratic scaling issue of conventional transformers. This transformation addresses the issue of high computational cost associated with standard transformer architectures.
2. Implementation of Efficient Chunking Strategies: The paper proposes a chunking strategy for training and inference, allowing for an efficient balance between memory use and computational efficiency. This strategy is critical in scaling the application of linear transformers across large sequences while maintaining their computational benefits.
3. Applications to Image Classification and Masked LLMing: The framework is applied to various neural architectures, showcasing its effectiveness in image classification tasks on ImageNet-1K and in masked LLMing using the C4 dataset. The results suggest that the proposed framework delivers significant speed improvements in training, up to 9 times faster than some traditional models, with competitive accuracy rates.
4. Scalability and Generalization: The paper further highlights the scalability of the proposed linear transformer approach, illustrating how models can efficiently generalize to inputs of varying resolutions and lengths without deteriorating performance.

Implications and Future Directions

The paper's findings have several implications for future research and application in AI:

• Practical Impact: On a practical level, this work facilitates the deployment of efficient and scalable models in real-time applications, such as natural language processing and computer vision tasks, where bidirectional context is advantageous.
• Theoretical Insights: The formulation of bidirectional RNNs equivalent to linear transformers provides a deeper understanding of how linear attention can bypass quadratic attention costs, offering a blueprint for further theoretical exploration in transformer models.
• Potential Extensions: Future research could extend this framework to other domains and tasks needing full attention. The framework's flexibility suggests applications in fields requiring efficient sequence modeling, such as genomics and audio processing.
• Integration with State-Space Models (SSMs): Integration with state-space models could enhance the stability and performance of linear transformers in various modeling contexts. The paper briefly hints at how such integrations can contribute to performance improvements, which warrants detailed exploration.

Overall, the proposed framework represents a step forward in making linear transformers more versatile and less computationally intensive. It addresses the common bottleneck of high computational demands when using traditional transformers in bidirectional contexts, thereby broadening the applicability of linear attention models in practical, real-world tasks. This research contributes to optimizing AI model efficiency, a significant priority as data scale and complexity continue to rise.