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 (with respect to output length) inference complexity. Recent works such as RetNet (Sun et al., 2023) and TransNormerLLM (Qin et al., 2023a) observe that adding a global decay term to the additive RNN update rule greatly improves performance, sometimes outperforming standard Transformers with softmax attention when trained at scale. In this work we show that adding a data-dependent gating mechanism further improves performance. We derive a parallel form of this gated linear attention layer that enables efficient training. However, a straightforward, numerically stable implementation of this parallel form requires generalized matrix multiplications in log-space for numerical stability, and thus cannot take advantage of tensor cores on modern GPUs which are optimized for standard matrix multiplications. We develop a hardware-efficient version of the parallel form that can still make use of tensor cores through block-parallel computations over sequence chunks. Experiments on moderate-scale language modeling (340M-parameter models trained on 15B tokens, 1.3B-parameter models trained on 100B tokens) show that gated linear attention (GLA) Transformers perform competitively against a strong LLaMA-architecture Transformer baseline (Touvron et al., 2023) as well as Mamba (Gu & Dao, 2023), a recently introduced state-space model with a data-dependent state transition mechanism. For training speed, our Triton-based implementation performs comparably to CUDA-optimized FlashAttention-2 (Dao, 2023) under the regular 2048 training length setting, while outperforming FlashAttention-2 when training on longer sequences beyond 4096.

翻译：具有线性注意力的Transformer模型既支持高效的并行训练，又可表述为具有二维（矩阵值）隐藏状态的循环神经网络，因此实现了与输出长度成线性关系的推理复杂度。近期研究如RetNet（Sun等人，2023）和TransNormerLLM（Qin等人，2023a）发现，在加性循环神经网络更新规则中加入全局衰减项可显著提升性能，在大规模训练时甚至能超越采用softmax注意力的标准Transformer模型。本文证明引入数据依赖的门控机制可进一步改善性能。我们推导了该门控线性注意力层的并行形式以实现高效训练。然而，直接实现该并行形式时为保证数值稳定性需要在对数空间进行广义矩阵乘法，因而无法利用现代GPU中针对标准矩阵乘法优化的张量核心。我们开发了该并行形式的硬件高效版本，通过序列分块的块并行计算仍可调用张量核心。中等规模语言建模实验（3.4亿参数模型使用150亿tokens训练，13亿参数模型使用1000亿tokens训练）表明，门控线性注意力（GLA）Transformer在性能上与强大的LLaMA架构Transformer基线（Touvron等人，2023）及最新引入数据依赖状态转移机制的状态空间模型Mamba（Gu & Dao，2023）具有竞争力。在训练速度方面，我们基于Triton的实现常规2048训练长度设置下与CUDA优化的FlashAttention-2（Dao，2023）相当，而在超过4096的长序列训练中则优于FlashAttention-2。