about
Gated Linear Attention Transformers with Hardware-Efficient Training (arxiv.org)
2 points by dataminer on Dec 12, 2023 | hide | past | pdf | discuss on HN

In plain words: Linear attention can run as a memory loop at test time but loses accuracy and speed; this version moves memory efficiently and adds gates that decide what to keep. It beats the fastest standard attention code at 1,000-token inputs and matches a usual transformer's language quality.

Abstract

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.

Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, Yoon Kim
arXiv:2312.06635 · cs.LG, cs.CL · submitted Dec 11, 2023 · updated Aug 27, 2024
abstract · pdf · html · minor update

add comment on HN