PaperlayerView exact revision on arXiv

Source-linked AI summary

Gated Linear Attention Transformers with Hardware-Efficient Training

Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, Yoon Kim

arXiv:2312.06635v6cs.LGcs.CL

TL;DR

Linear attention offers efficient training and linear-time inference but often underperforms softmax attention, while existing implementations are not sufficiently I/O-aware. The paper introduces hardware-efficient FLASHLINEARATTENTION and a gated linear attention Transformer, which performs favorably against strong Transformer and linear-time baselines and generalizes effectively to longer contexts.

  • Problem

    Linear attention combines parallel training with linear-time inference but generally underperforms softmax attention, and existing implementations lack I/O-awareness.

  • Method

    The paper develops I/O-aware chunkwise linear-attention algorithms and extends them with data-dependent gates for a gated linear attention Transformer.

  • Results

    GLA Transformer performs favorably against LLaMA-style Transformers, RetNet, and Mamba, with particularly strong length generalization.

  • Takeaways & Limitations

    GLA provides a competitive linear-time-inference alternative with efficient hardware utilization and strong extrapolation from shorter training contexts.

  • Takeaways & Limitations

    The experiments do not establish how GLA scales to larger models or datasets because larger-scale studies were limited by compute resources.

Abstract

from arXiv · show

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.

1 Introduction

Linear attention offers parallel training and linear-time inference through matrix-valued recurrent states, but often trails softmax attention and has lacked hardware-efficient implementations. The paper introduces I/O-aware FLASHLINEARATTENTION and gated linear attention to address these gaps.

  • Linear attention combines efficient parallel training with linear-time inference through a recurrent formulation using 2D hidden states.
  • Linear attention often underperforms softmax attention, while existing variants use global, data-independent decay factors and remain behind strong Transformer architectures.
  • The paper develops I/O-aware algorithms for hardware-efficient linear attention and applies them to a gated variant competitive with softmax attention.
  • FLASHLINEARATTENTION is faster than FLASHATTENTION-2 even on short 1K sequences, while GLA Transformers perform favorably against LLaMA-style Transformers, RetNet, and Mamba.

2 Background: Linear Attention

Linear attention replaces softmax similarity with a kernelized dot product, yielding a recurrent matrix-state formulation and a chunkwise form that interpolates between parallel and recurrent computation. Chunking reduces training complexity below quadratic when sequence length exceeds hidden dimension.

  • Linear attention replaces exponential query-key similarity with a kernel and feature map, simplifying attention computation.
  • The resulting layer can be written as a linear RNN whose hidden state is a matrix updated through outer products.
  • Causal parallel linear attention remains quadratic because the causal mask prevents reassociation from reducing its complexity.
  • Chunkwise parallel attention splits the sequence into length-C chunks and combines inter-chunk recurrence with intra-chunk parallel computation.
  • O(LCd+Ld2) is less than O(L2d) when L>d, with C=L recovering the parallel form and C=1 recovering the recurrent form.

3 Hardware-Efficient Linear Attention

Hardware-efficient linear attention must balance GPU occupancy, tensor-core use, and memory movement. FLASHLINEARATTENTION implements chunkwise computation with materialization choices that trade memory footprint against sequence-level parallelism and achieves higher speed than key baselines.

  • Hardware considerations: Efficient GPU implementations must account for compute models, memory hierarchies, and specialized units such as tensor cores.
  • Hardware considerations: Recurrent linear attention has low FLOPs but poor wall-time efficiency because elementwise updates cannot use tensor cores, while materializing hidden states increases memory I/O.
  • Chunkwise form: Chunkwise attention exposes a tunable trade-off between parallelism and computation, with matrix multiplications enabling tensor-core use when C is a multiple of 16.
  • FLASHLINEARATTENTION: FLASHLINEARATTENTION uses tiled on-chip reuse to reduce repeated HBM transfers during chunkwise computation.
  • Materialization trade-off: The non-materialization version is more memory-efficient but lacks sequence-level parallelism, whereas materialization enables chunk parallelism at an approximately 10-20% larger memory footprint.
  • Results: Both FLASHLINEARATTENTION versions are substantially faster than FLASHATTENTION-2 and pure PyTorch chunkwise linear attention.

4 Gated Linear Attention

GLA adds data-dependent gating to linear attention while preserving a hardware-efficient chunkwise form for training. Its implementation uses secondary tiling and memory-efficient gradient computation to balance expressiveness, numerical stability, and GPU efficiency.

  • Recurrent and Parallel Form of GLA: GLA introduces a data-dependent 2D forget gate that generalizes linear attention while retaining a hardware-efficient chunkwise training form.The gate varies over time and the formulation encompasses several recent RNNs.
  • Recurrent and Parallel Form of GLA: The gate parameterization balances parameter efficiency, hidden-state size, and training efficiency, using a middle ground between scalar and fully low-rank gating.The paper uses an outer-product-based parameterization with low-rank prediction of α_t.
  • Recurrent and Parallel Form of GLA: GLA’s parallel form is numerically stabilized by computing the attention-like matrix P in log space because cumulative gate products can become extremely small.Without this stabilization, the corresponding K and B factors can explode; the formulation also cannot use standard half-precision matmuls directly.
  • Chunkwise Parallel Form of GLA: The chunkwise form separates intra-chunk computation from inter-chunk recurrence, propagating prior hidden states through cumulative decay factors.Λ encodes decay from a chunk’s start, while Γ encodes decay toward its end.
  • Chunkwise Parallel Form of GLA: Secondary-level chunking enables half-precision matmuls for interactions between sub-chunks while retaining full-precision log-space computation within sub-chunks.This two-level tiling reduces non-half-precision matmul FLOPs and improves wall-clock performance.
  • Hardware-Efficient Training: GLA provides a memory-efficient closed form for dlogα_t gradients, avoiding materialization of all matrix-valued hidden states in HBM.The paper derives dlogα_t directly and computes dqt and dkt as in the ordinary linear-attention algorithm.

5 Empirical Study

The empirical study evaluates GLA against Transformer and linear-time baselines across language modeling, recall-intensive tasks, length extrapolation, and hardware efficiency. GLA is competitive with strong baselines, extrapolates effectively, and achieves favorable throughput while maintaining similar memory usage.

  • Main results: GLA improves over RetNet on all evaluated tasks, while both GLA and Mamba perform comparably to Transformer++.All baselines use the same dataset and token budget for fair comparison.
  • Recall-intensive tasks: GLA outperforms RetNet on the synthetic MQAR recall task, while matrix-valued-state models outperform Hyena and RWKV.Standard quadratic attention achieves perfect scores and is omitted from the comparison.
  • Recall-intensive tasks: Subquadratic models significantly underperform Transformers on FDA and SWDE, although GLA outperforms the other subquadratic models on these recall-intensive tasks.The authors associate GLA’s advantage with its larger recurrent state and selection mechanism.
  • Long sequence training and length extrapolation: Models trained on 2K contexts show that GLA extrapolates better than Mamba and RetNet on most PG19 position buckets, while GLA and RetNet generalize to 18K on SlimPajama.Mamba struggles beyond 4K, and Transformers cannot extrapolate beyond their training length.
  • Long sequence training and length extrapolation: Long-sequence pretraining consistently improves perplexity, while GLA shows only marginal differences between direct long-context training and TBPTT settings.Mamba benefits significantly from 8K training and performs similarly to GLA in that setting.
  • Ablations: The ablation finds that both data-dependent scalar gates and finer-grained gating matter, while one head performs best but requires much more GPU memory.The experiments therefore use four heads as a trade-off between perplexity and memory.
  • Training efficiency: GLA has similar GPU memory usage to the compared models and higher training throughput than Mamba, with greater advantages beyond sequence length 4096.The comparison uses 1.3B models on a single H100 GPU and FLASHLINEARATTENTION with hidden-state recomputation.
  • Limitations and future work: Larger-scale validation remains unresolved because the experiments could not be extended beyond the studied scales due to limited compute resources.The authors state that scaling behavior to larger models and datasets is unclear.

6 Related Work

Related work frames linear RNNs, state-space models, and linear Transformers as scalable subquadratic alternatives enabled by parallel temporal training. It also emphasizes data-dependent gating and hidden-state dimensionality as important design axes, while contrasting GLA’s hardware alignment with Mamba’s limitations.

  • Linear RNNs and state-space models: Linear RNNs, state-space models, and Transformers remove nonlinear temporal dependencies, enabling parallel training as subquadratic alternatives to traditional Transformers.Traditional RNNs are difficult to scale because of nonlinear hidden-state dependencies and sequential matrix updates.
  • Data-dependent gating: Data-dependent decay or gating has long been important in RNNs, but some approaches restrict gates to current inputs to preserve parallel training.HGRN, RWKV-v6, and Mamba are cited as using data-dependent mechanisms.
  • Hardware-aware implementations: GLA’s multi-head structure is more amenable to tensor parallelism than Mamba’s non-multi-head design.The paper identifies this as a reason GLA may become more efficient than Mamba at larger model sizes.
  • Hidden-state parameterization: Linear Transformers expand recurrent memory through outer-product hidden states, whereas linear SSMs expand it through a single-input-single-output strategy.Hidden-state dimensionality acts as a proxy for memory capacity and influences expressive power.

7 Conclusion

The paper proposes an efficient training algorithm for gated linear-attention Transformers with data-dependent gates. Language-modeling experiments show these models perform respectably against strong baselines.

  • The proposed algorithm efficiently trains linear-attention Transformers with data-dependent gating mechanisms.It balances FLOPs against parallelism while retaining half-precision matrix multiplications on GPU tensor cores.

Impact Statement

The paper targets improved training efficiency for gated linear-attention models, while noting that effects on biased and harmful language-model outputs remain unexplored.

  • The work aims to improve training efficiency for gated linear-attention models.Such efficiency could help democratize access to language models.
  • Whether these architectures affect biased and harmful language-model outputs remains an unexplored research question.

A.1 Linear Attention

Linear attention offers efficient sequence modeling but faces performance, memory, attention-concentration, and hardware-efficiency challenges. Gated, I/O-aware, chunkwise methods address parts of these limitations while introducing trade-offs and open directions.

  • Feature maps: The paper uses the identity feature map and leaves integrating non-identity feature maps into GLA for future work.
  • Attention concentration: Linear attention can suffer attention dilution, producing overly uniform distributions that fail to concentrate on relevant tokens.
  • Memory capacity: Linear attention has bounded memory size, whereas softmax attention has unbounded memory; increasing memory size and utilization may help bridge their performance gap.
  • Gating and memory: Matrix-valued gates provide finer-grained memory control but make training harder because existing methods materialize hidden states and cannot use tensor cores.
  • Hardware-efficient training: Most chunkwise linear-attention methods are not I/O-aware, with LIGHTNINGATTENTION2 identified as a concurrent similar exception.
  • Hardware-efficient training: FLASHLINEARATTENTION adds a materialization variant that uses sequence-level parallelism for higher throughput at a slightly larger memory footprint.

A.2 Sequence parallelism

The paper develops chunkwise and sequence-parallel training methods that trade memory use against parallelism. Materialization improves sequence-level parallelism, while non-materialization saves memory; gated variants extend the same framework.

  • Gated linear attention: The main gated linear-attention parameterization fixes β to 1, while making β learnable showed no empirical performance gain.
  • Gated linear attention: The generalized chunkwise form supports vanilla linear attention, RetNet-style decay, and fine-grained data-dependent decay through different parameter settings.

D Additional Experimental Results

Table 5 presents complete results across all 11 tasks, including 5-shot results for the 1.3B models.

  • Table 5 reports results on all 11 tasks, including 5-shot results for the 1.3B models.
Loading 2312.06635v6…