Gated Linear Attention Transformers with Hardware-Efficient Training

By Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, Yoon Kim. International Conference on Machine Learning, 2023

480 citations recorded by the source corpus.

Transformers with linear attention allow for efficient parallel training but can simultaneously be formulated as an RNN with 2D (matrix-valued) hidden states, thus enjoying linear-time inference complexity. However, linear attention generally underperforms ordinary softmax attention. Moreover, current implementations of linear attention lack I/O-awareness and are thus slower than highly optimized implementations of softmax attention. This work describes a hardware-efficient algorithm for linear attention that trades off memory movement against parallelizability. The resulting implementation, dubbed FLASHLINEARATTENTION, is faster than FLASHATTENTION-2 (Dao, 2023) as a standalone layer even on short sequence lengths (e.g., 1K). We then generalize this algorithm to a more expressive variant of linear attention with data-dependent gates. When used as a replacement for the standard attention layer in Transformers, the resulting gated linear attention (GLA) Transformer is found to perform competitively against the LLaMA-architecture Transformer (Touvron et al., 2023) as well recent linear-time-inference baselines such as RetNet (Sun et al., 2023a) and Mamba (Gu&Dao, 2023) on moderate-scale language modeling experiments. GLA Transformer is especially effective at length generalization, enabling a model trained on 2K to generalize to sequences longer than 20K without significant perplexity degradations. For training speed, the GLA Transformer has higher throughput than a similarly-sized Mamba model.

Subfields: Artificial Intelligence, Electrical and Electronic Engineering

View Paper: https://www.semanticscholar.org/paper/62b18cc55dcc7ffe52c28e1086aee893b7bc4334

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

1. Yunjie Tian
2. Qixiang Ye
3. Yunfan Liu
4. Yuzhong Zhao
5. Yue Liu
6. Hongtian Yu
7. Lingxi Xie
8. Yaowei Wang

In Electrical and Electronic Engineering

1. Qixiang Ye
2. Lingxi Xie

Citations in the corpus

• VMamba: Visual State Space Model (2024, background): “19 reduces to the form of Gated Linear Attention (GLA) [60], which indicates that GLA is also a special case of Mamba.”
• Gated Delta Networks: Improving Mamba2 with Delta Rule (2024, unknown): “…underperformed in language modeling tasks compared to standard Transformers, recent enhancements—such as incorporating data-dependent gating mechanisms akin to those in LSTMs, exemplified by models like GLA (Yang et al., 2024a) and Mamba2 (Dao & Gu, 2024a)—have shown promising improvements.”
• Parallelizing Linear Transformers with the Delta Rule over Sequence Length (2024, unknown): “We scale DeltaNets to moderate-scale language modeling benchmarks (1.3B models trained on 100B tokens), where DeltaNet is found to obtain better language modeling and zero-shot downstream task performance than strong linear recurrent models such as Mamba [30] and GLA [116].”
• Learning to (Learn at Test Time): RNNs with Expressive Hidden States (2024, unknown): “Some more recent examples are RWKV [58, 59], xLSTM [4], and Gated Linear Attention (GLA) [82].”
• Quasi Linear Kernel Attention with Infinite Capacity (2026, unknown): “Building on this, Arora et al. (2024b) combine linear attention with sliding windows, while gated linear attention (Yang et al., 2024a) and DeltaNet (Yang et al., 2024b) use data-dependent gating and delta-rule updates of the recurrent state, respectively, to better use a fixed-size state.”