Gated Linear Attention Transformers with Hardware-Efficient Training

• URL: http://arxiv.org/abs/2312.06635v6
• Date: Tue, 27 Aug 2024 01:27:29 GMT
• Title: Gated Linear Attention Transformers with Hardware-Efficient Training
• Authors: Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, Yoon Kim,
• Abstract summary: This work describes a hardware-efficient algorithm for linear attention that trades off memory movement against parallelizability. 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 Transformer is found to perform competitively.
• Score: 60.670102007737476
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Parallelizing Linear Transformers with the Delta Rule over Sequence Length [49.88826673324244]
  This work describes a hardware-efficient algorithm for training linear transformers with the delta rule. We train a 1.3B model for 100B tokens and find that it outperforms recent linear-time baselines.
  arXiv  Detail & Related papers  (2024-06-10T17:24:42Z)
• The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry [24.198536617002667]
  Linear attentions have shown potential for improving Transformer efficiency, reducing attention's quadratic complexity to linear in sequence length. We propose Hedgehog, a learnable linear attention that retains the spiky and monotonic properties of softmax attention while maintaining linear complexity.
  arXiv  Detail & Related papers  (2024-02-06T19:31:26Z)
• Ultra-Long Sequence Distributed Transformer [10.263668150008316]
  Transformer models trained on long sequences often achieve higher accuracy than short sequences. Existing methods for long sequence training offer limited speedup and memory reduction. This paper presents a novel and efficient distributed training method, the Long Short-Sequence Transformer.
  arXiv  Detail & Related papers  (2023-11-04T11:38:53Z)
• Linear attention is (maybe) all you need (to understand transformer optimization) [55.81555204646486]
  We make progress towards understanding the subtleties of training Transformers by studying a simple yet canonicalized shallow Transformer model. Most importantly, we observe that our proposed linearized models can reproduce several prominent aspects of Transformer training dynamics.
  arXiv  Detail & Related papers  (2023-10-02T10:48:42Z)
• Hyena Hierarchy: Towards Larger Convolutional Language Models [115.82857881546089]
  Hyena is a subquadratic drop-in replacement for attention constructed by interleaving implicitly parametrized long convolutions and data-controlled gating. In recall and reasoning tasks on sequences of thousands to hundreds of thousands of tokens, Hyena improves accuracy by more than 50 points over operators relying on state-spaces and other implicit and explicit methods.
  arXiv  Detail & Related papers  (2023-02-21T18:29:25Z)
• The Devil in Linear Transformer [42.232886799710215]
  Linear transformers aim to reduce the quadratic space-time complexity of vanilla transformers. They usually suffer from degraded performances on various tasks and corpus. In this paper, we identify two key issues that lead to such performance gaps.
  arXiv  Detail & Related papers  (2022-10-19T07:15:35Z)
• Linearizing Transformer with Key-Value Memory Bank [54.83663647680612]
  We propose MemSizer, an approach to project the source sequence into lower dimension representation. MemSizer not only achieves the same linear time complexity but also enjoys efficient recurrent-style autoregressive generation. We demonstrate that MemSizer provides an improved tradeoff between efficiency and accuracy over the vanilla transformer.
  arXiv  Detail & Related papers  (2022-03-23T18:10:18Z)
• Stable, Fast and Accurate: Kernelized Attention with Relative Positional Encoding [63.539333383965726]
  We propose a novel way to accelerate attention calculation for Transformers with relative positional encoding (RPE) Based upon the observation that relative positional encoding forms a Toeplitz matrix, we mathematically show that kernelized attention with RPE can be calculated efficiently using Fast Fourier Transform (FFT)
  arXiv  Detail & Related papers  (2021-06-23T17:51:26Z)
• Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention [22.228028613802174]
  Transformers achieve remarkable performance in several tasks but due to their quadratic complexity, they are prohibitively slow for very long sequences. We make use of the associativity property of matrix products to reduce the complexity from $mathcalOleft(N2right)$ to $mathcalOleft(Nright)$, where $N$ is the sequence length. Our linear transformers achieve similar performance to vanilla transformers and they are up to 4000x faster on autoregressive prediction of very long sequences.
  arXiv  Detail & Related papers  (2020-06-29T17:55:38Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.