[论文解读] Gated Linear Attention Transformers with Hardware-Efficient Training

Yang Song-lin, Bailin Wang+3 more|arXiv (Cornell University)|Dec 11, 2023

Neural Networks and Applications被引用 5

一句话总结

本文提出门控线性注意力（GLA），一种硬件高效的Transformer变体，用数据相关的门控线性注意力机制替代标准的softmax注意力机制。通过优化内存访问模式并引入一种新颖的I/O感知算法——FLASHLINEARATTENTION，GLA在短序列（例如1K）上训练速度甚至快于FlashAttention-2，同时保持线性时间推理。GLA-Transformer在语言建模任务中表现与强基线模型（如LLaMA、RetNet和Mamba）相当或更优，尤其在长序列泛化和召回密集型任务中表现突出。

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.

研究动机与目标

• 解决线性注意力与标准softmax注意力在Transformer中性能差距的问题。
• 通过设计一种I/O感知、硬件优化的算法，减少内存移动，提升线性注意力的训练效率。
• 通过数据相关的门控机制增强线性注意力的表达能力，提升其在长上下文和召回密集型任务中的性能。
• 在保持线性时间推理复杂度的同时，实现与标准Transformer相当的竞争力表现。
• 展示模型在未见过的更长序列长度上具备强大的泛化能力，且性能无明显下降。

提出的方法

• 提出FLASHLINEARATTENTION，一种针对现代GPU优化内存访问与并行性的硬件高效线性注意力算法。
• 引入一种分块并行训练形式，结合块间递归与块内并行计算，在I/O效率与训练速度之间实现平衡。
• 设计一种门控线性注意力机制，其中隐藏状态更新受数据相关的门控调制，提升建模能力。
• 利用前缀和技巧推导出可学习参数α和β的闭式梯度，避免在高带宽内存中存储中间状态。
• 将FLASHLINEARATTENTION算法推广至门控变体，实现GLA-Transformer的高效训练。
• 采用一种改进的注意力计算方式，在保持类似RNN的递推结构以实现线性时间推理的同时，通过分块实现并行训练。

实验结果

研究问题

• RQ1 在短序列上，硬件优化的线性注意力实现是否能超越高度优化的softmax注意力实现（如FlashAttention-2）？
• RQ2 在线性注意力中引入数据相关的门控机制，是否能显著提升性能，相比固定衰减因子或原始线性注意力？
• RQ3 GLA-Transformer能否在保持线性时间推理的同时，实现与标准Transformer（如LLaMA）相当的性能？
• RQ4 GLA-Transformer在长于其训练长度的序列上泛化能力如何，特别是在召回密集型任务中？
• RQ5 与SOTA线性时间模型（如Mamba和RetNet）相比，GLA-Transformer的训练吞吐量如何？

主要发现

• FLASHLINEARATTENTION在短序列（如1K token）上也比FlashAttention-2更快，证明其具有更优的I/O效率。
• GLA-Transformer在语言建模基准测试中表现优异，性能与LLaMA架构模型相当或更优，并超越近期的线性时间模型（如RetNet和Mamba）。
• 在340M和1.3B参数模型上，GLA-Transformer在11个零样本任务上的平均准确率分别为48.0%和55.5%，在多个基准测试中优于Mamba和RetNet。
• GLA-Transformer在超过20K token的长序列上表现出有效的泛化能力，即使在2K序列上进行训练，困惑度下降也极小。
• GLA-Transformer的训练吞吐量高于同等规模的Mamba模型，表明其在大规模预训练中具备更好的可扩展性。
• 模型在BoolQ和ARC等召回密集型任务中表现强劲，表明其具备更强的记忆保持能力和长距离依赖建模能力。