Gated Linear Attention Transformers with Hardware-Efficient Training

使用线性注意力的Transformer可以实现高效的并行训练，但同时也可以被表述为具有二维（矩阵值）隐藏状态的RNN，因此可以享受线性（相对于输出长度）推理复杂度。最近的研究（例如Sun等人的RetNet（2023年）和Qin等人的TransNormerLLM（2023a））观察到，将全局衰减项添加到加性RNN更新规则中可以极大地提高性能，有时在大规模训练时优于具有softmax注意力的标准Transformer。在这项工作中，我们展示了添加数据相关门控机制可以进一步提高性能。我们推导出了这个门控线性注意力层的并行形式，从而实现了高效的训练。然而，这个并行形式的直接、数值稳定的实现需要在对数空间中进行广义矩阵乘法以实现数值稳定性，因此无法利用现代GPU上针对标准矩阵乘法进行优化的张量核心。我们开发了一种硬件高效的并行形式，通过对序列块进行块并行计算仍然可以利用张量核心。在中等规模语言建模（340M参数模型在15B标记上训练，1.3B参数模型在100B标记上训练）的实验中，门控线性注意力（GLA）Transformer与强大的LLaMA架构Transformer基线（Touvron等人，2023）以及最近引入的具有数据相关状态转移机制的状态空间模型Mamba（Gu和Dao，2023）相比表现出竞争力。对于训练速度，我们基于Triton的实现在常规的2048训练长度设置下与CUDA优化的FlashAttention-2（Dao，2023）性能相当，而在训练超过4096的较长序列时优于FlashAttention-2。

论文旨在探索一种新的基于数据门控机制的线性注意力机制，以提高Transformer的性能。

通过在线性注意力机制中添加数据门控机制，可以进一步提高Transformer的性能。研究人员提出了一种并行化的门控线性注意力层，可以实现高效的训练。

论文通过实验表明，门控线性注意力（GLA）Transformer可以与当前最先进的Transformer模型相媲美，并且在一些任务上表现更好。此外，研究人员还提出了一种硬件有效的并行化实现方法，可以利用现代GPU的张量核心进行加速。

最近的相关研究包括RetNet、TransNormerLLM、Touvron等人的LLaMA-architecture Transformer以及Gu和Dao的Mamba模型。