计算机体系结构技术杂谈（下）
2.8 加速Transformer：稀疏注意力加速器分析
1. 稀疏注意力加速器简介
近年来， Transformer模型在深度学习的各个领域，包括自然语言处理、图像分类、图像及语音生成等方面，都取得了远超于传统神经网络模型的表现。最近的ChatGPT和各类基于Transformer的AIGC应用也由于其惊人的表现被大家广泛关注。
Transformer模型之所以能取得如此惊人表现的主要原因之一是其中的注意力机制模块（Attention Mechanism），然而注意力机制模块由于需要进行大量的矩阵乘计算，导致其成为影响Transformer模型性能的主要瓶颈。如何设计硬件加速器来高效地运行大型注意力机制模块便成为了近期研究关注的重点。在中，将与大家一起分享探讨近期各大顶会中发表的稀疏注意力机制加速器的实现方法，希望能给各位读者一些启发。
2. TRANSFORMER和自注意力机制的工作原理
先来看一看Transformer和自注意力机制的工作原理。Transformer模型通常可以由encoder模块堆叠而成。在encoder模块中，可以通过自注意力机制，计算出一个sequence中每个token和其他token之间的两两对应的关联性，从而得到需要被“注意”的token-pair。n个token间的两两对应的关联性可以用一个n x n的矩阵来表示，如图2-49所示。
图2-49 n个token之间的注意力关系
Transformer模型架构，如图2-50所示。
图2-50 Transformer模型架构
那么Transformer模型是如何一步一步计算出这些token之间的关联呢？如图2所示，假设Transformer的输入有n个token，首先，它们会先被一个embedding模块转换为n x d大小的X矩阵；随后，X会在Linear Transformation模块中分别与三个权重矩阵
相乘，并线性转换为三个n x n大小的矩阵，分别称为Query（Q），Key（K）和Value（V）：
（1）
将
和
（既K矩阵的转置）相乘得到一个n x n的注意力得分（Attention Score）S矩阵，并对其每一行计算SoftMax，便可以得到一个表达token和token之间的关联性注意力权重（Attention Weight）A矩阵：
（2）
最后，将A矩阵与V矩阵相乘，即可得到最终的n x d大小的自注意力输出Z。
（3）
如图2所示，在Transformer模型中，注意力的输出还会通过一个残差连接与其输入相加，并通过计算LayerNorm函数加上一个前向网络（Feed-Forward Network，FFN）得到Transformer的encoder模块的最终输出。通常一个模型中会叠加多个encoder模块，并在最后使用一个分类器（Classifier）来完成一些分类预测的任务。
3. 注意力机制的计算复杂度
由于注意力机制的复杂度是随着token的数量增加（既sequence的长度）呈平方数增长的（如图1所示），从而增加了传统Transformer模型在硬件部署的难度。进一步分析一下一个Transformer中每一步所带来的计算开销：在一个传统的Transformer模型中，encoder模块有三个主要的矩阵乘计算，分别为：a）QKV的线性转换（公式1）；b）自注意力计算（公式2和3）；和c）FFN前向网络。如图2-51所示，通过计算在不同输入sequence长度下各个矩阵乘的计算量（FLOPs）对比可以看到，随着sequence长度的增加，自注意力计算的占比（蓝色部分）逐渐增加，最终远超于线性转换和前向网络计算量的总和（橙色部分）。
图-51 自注意力计算的开销占比
4. 加速自注意力机制：稀疏化注意力权重矩阵
通过上面的分析可以看到，在sequence长度大于一定量级时，注意力机制的计算便会成为Transformer模型计算的主要瓶颈，学术界和产业界也都提出了各种方法来加速注意力机制，其中最受关注的方法之一便是注意力权重剪枝。与普通的权重剪枝（Weight Pruning）有所不同的是，注意力矩阵只能在运行时由公式2计算得出，而无法像权重剪枝一样在运行之前就可以完成。所以，有一部分的研究关注在如何高效准确地在运行时对注意力矩阵进行剪枝，如[1]和[2]。当然，运行时进行剪枝会带来额外的开销以及硬件设计的难度，所以也有很多研究关注在如何使用固定的pattern进行剪枝而不损失Transformer模型的预测精确度。
2.9 DOTA：检测和忽略可扩展Transformer加速的薄弱环节
2.9.1 主要思想
提出了一种运行时剪枝注意力矩阵的方法。作者在训练Transformer模型时，同步训练了一个简单的弱注意力检测器（Weak Attention Detector）。在进行推理时，这个检测器会在计算注意力权重矩阵之前，预测出token间的弱关联性（即弱注意力），并在计算注意力时忽略这些token的计算。具体检测方法如下：首先，输入X会先被线性转换为低维度的
和
：
（4）
其中，P为d x k大小的矩阵，其作用是将X的尺寸从n x d减小为n x k（其中k远小于d）。P矩阵的数据会在
内做随机采样来决定。此外，
和
也都是k x k大小的权重矩阵。可以用
和
相乘得到一个n x n大小的预估注意力得分
：
（5）
这个检测器中的
和
，可以通过一个MSE（Mean Squared Error）损失函数在训练Transformer时同步进行优化：
（6）
其中B为mini-batch的大小。由于 和公式2中A矩阵的尺寸相同，可以用 矩阵来预测A矩阵中的弱注意力位置，并在计算A矩阵时，跳过相应的计算。弱注意力的阈值可以通过一个top-K算法搜索得到。
2.9.2 加速器设计
DOTA整体加速器的主体设计，如图2-52所示。
图2-52 DOTA注意力机制加速器和弱注意力检测器
加速器整体分成多个lane实现token-level数据并行计算，每个lane中有一个可重构的矩阵乘单元（RMMU，Reconfigurable Matrix Multiplication Unit）用来计算浮点和整数精度的矩阵乘，以及一个弱注意力检测器。检测器会将预估的注意力得分与阈值进行比较，并由此生成一个n x n大小的bit mask用0和1来表示弱注意力和强注意力在注意力权重矩阵的位置，以便在之后的调度器（scheduler）中跳过弱注意力的计算。此外，检测器中还使用了乱序执行等方法做进一步的加速。
