Linear Attention for Efficient Bidirectional Sequence Modeling

TL;DR

LION is a novel framework that extends Linear Transformers to bidirectional sequence modeling, enabling faster training and inference while matching or surpassing softmax Transformer performance.

Contribution

The paper introduces LION, the first systematic approach to adapt Linear Transformers for bidirectional tasks, unifying multiple representations and demonstrating broad applicability.

Findings

LION achieves comparable or better performance than softmax Transformers on bidirectional tasks.

LION offers significantly faster training and more efficient inference than existing State Space Models.

The framework generalizes core Linear Transformer representations to the bidirectional setting.

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 arXiv:2006.16236; LION-D, based on arXiv:2307.08621; and LION-S, a variant using selective decay arXiv:2103.02143, arXiv:2312.0075. 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.