FLASH：一种高效的Transformer设计

最新推荐文章于 2025-07-26 13:45:23 发布

传道解惑也

最新推荐文章于 2025-07-26 13:45:23 发布

阅读量1.6k

点赞数 6

CC 4.0 BY-SA版权

分类专栏： NLP论文阅读文章标签： transformer 深度学习人工智能

本文链接：https://blog.csdn.net/shawroad88/article/details/126617904

NLP论文阅读专栏收录该内容

7 篇文章

订阅专栏

FLASH是一种针对长序列处理的高效Transformer模型，通过GAU单元融合注意力机制与前馈网络，实现更快速且节省显存的目标。文章介绍了FLASH的工作原理，包括局部与全局注意力计算方法，以及在不同任务上的实证表现。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

背景

近年来，Transformer凭借其优秀的设计，在文本、图像、语音等方向大杀四方。但是由于其attention的二次复杂度限制了其在长序列上的应用。本文提出了一种快(速度快)、省(省显存)的模型FLASH(Fast Linear Attention with a Single Head)，在长序列的表现远远高于标准的Transformer。

模型介绍

GAU(Gated Attention Unit)

在标准的Transformer结构中，多头注意力和FFN是交替连接的。GLU那篇论文中，将FFN替换成基于门控的线性单元，发现效果会变好。因此，我们先简单了解一下门控单元GLU的计算，如下左图:

具体计算:

也就是将输入X分别经过放射变换(线性映射＋激活函数)得到 $U, V$ 。然后再将 $U, V$ 进行点积，最后再进行线性映射，得到门控线性单元的输出。

上述的GLU中没有对token两两进行注意力计算，如果在上面的 $U, V$ 中引入注意力，那岂不是就省了前面的多头注意力计算了。如下式：

在 $U, V$ 进行点积计算的时候，如果给 $V$ 乘一个注意力矩阵 $A$ (维度为nxn，其中n为序列的长度)，那岂不是就引入了注意力信息。

基于此，本文就提出了一种新的结构GAU。主要是给出了一种注意力矩阵A的计算方法。具体计算如下:

对输入X进行放射变换(线性映射+激活函数)得到Z，然后对Z分别进行 $Q,K\mathcal{Q},\mathcal{K}$ 变换，就是对Z中的每一个标量进行平移等运算。这里的 $Q,K\mathcal{Q},\mathcal{K}$ 变化类似于LayerNorm中的 $α,β\alpha, \beta$ ，是可训练的。然后将两种变换的结果进行矩阵乘。最后再经过一个 $relu^2$ 的激活函数( $relu^2$ 是将relu的计算结果平方)，得到最终的注意力矩阵。然后和门控注意力单元中的V进行相乘，这样就在门控注意力单元中引入了注意力信息。具体结果如下图所示:

注意:可能是GLU对attention的依赖没有那么强，因此，作者在实验中只用了一个注意力头。

在GAU的实验中，作者固定e=2d，那"n层Attention+n层FFN"的标准Transformer模型，对应的就是"2n层GAU"的新模型，即该模型为FLASH-Quad。其中Quad表明复杂度依然是二次的。即：FLASH的二次复杂度版本。

Fast Linear Attention with GAU

可能有读者发现，在上述的GAU中，你只是将attention和FFN合并起来，替代了标准的attention+FNN，并没有解决attention的二次复杂度呀？对，因此，作者提出了一种快速计算注意力的方法。

过去，在解决注意力的二次复杂度问题上，有两种主流方法: (1)将注意力计算稀疏化、(2)将注意力计算线性化。稀疏化即人为根据先验知识规定哪些token可以进行注意力计算(典型代表: Longformer、BigBird等)。线性化则是提出另外的方法，去逼近标准注意力的效果（典型代表: Linformer、Performer等），如下公式所示:

正常的注意力是将 $Q, K$ 进行矩阵乘，接着经过softmax，最后乘V。如果将K和V先进行乘，则可以大大减少计算量。假设 $Q, K, V$ 的维度为: $(m, d)$ ,则标准注意力的计算量为:
$m * d * m + m * d * m$ ,即: $2dm^2$ ,是跟序列长度m成平方正比。如果先算K乘V，则计算量为: $d * m * d + d * m * d$ ,即: $2md^2$ ,是跟序列长度m成一次正比。所以，第二种方法随着序列的边长，效率会远高于第一种方法。

本文则是根据上述两种方式，结合"稀疏化"和"线性化"的优点，提出了一种"局部+全局"的分块混合的注意力计算方法。

首先是分块注意力的计算，假设序列长度为n，每个块的维度为c，则可分成n/c个块(默认可整除)。 $Ug,Vg∈Rc×e,Zg∈Rc×s\boldsymbol{U}_g,\boldsymbol{V}_g\in\mathbb{R}^{c\times e},\boldsymbol{Z}_g\in\mathbb{R}^{c\times s}$ ，其中g指的是第g个块。将 $Zg\boldsymbol{Z}_g$ 通过四个放射变换(线性映射+激活)分别得到 $Qgquad,Kgquad,Qglin,Kglin\boldsymbol{Q}_g^{\text{quad}},\boldsymbol{K}_g^{\text{quad}},\boldsymbol{Q}_g^{\text{lin}},\boldsymbol{K}_g^{\text{lin}}$ 。则块内注意力计算如下:

可以看出上述公式很好理解，不做过多描述。接下来算算其复杂度。每个块内注意力计算复杂度为 $c^2$ ，有n/c个块，则块注意力计算整体的复杂度为: $n/c) * c^2$ ，即nc,也就是正比于n。

接着用 $Qglin,Kglin\boldsymbol{Q}_g^{\text{lin}},\boldsymbol{K}_g^{\text{lin}}$ 进行全局的attention计算。这里采用的是上述介绍的注意力线性化方法。计算如下:

上述(7)式就是全局的两两计算。(8)式主要是在生成任务中，对标的是带有mask的多头注意力。

最后将两种attention结果整合到GAU中，得到线性版的GAU网络，计算如下:

作者在论文中还贴出来注意力计算的代码，如下所示:

至此，论文就介绍完了，下面简单看一下实验结果。

实验

首先，作者对比了GAU、多头注意力+FFN、以及多头注意力+GLU三种结构，在自回归任务和MLM任务上的表现，如下图:

横轴为速度，纵轴为效果，越靠右上，效果越好。上述实验是在长度为512上的效果对比。可以发现，GAU的在相同效果的前提下，速度更快；在相同速度的前提下，效果更好。

从上表中可以看出，虽然FLASH-Quad也是二次复杂度，但是也比标准的Transformer效果好，速度也更快。另外，随着序列的逐渐变长,FLASH的速度远远快于标准的Transformer。

最后看看消融实验，如下图所示:

上面MF-TFM++模型采用的是多头注意力+FFN的结构，只是多头注意力采用的线性注意力+块注意力。也就是本文提出的注意力计算。从消融实验中，可以看出一个很有用的信息，即localOnly attention比GlobalOnly attention更重要。

本文参考:
[1] 论文: https://arxiv.org/abs/2202.10447
[2] GLU: https://arxiv.org/abs/2002.05202
[3] 苏剑林. (Feb. 25, 2022). 《FLASH：可能是近来最有意思的高效Transformer设计》[Blog post]. Retrieved from https://kexue.fm/archives/8934