机器学习课程学习周报八

摘要

本周的学习重点是自注意力机制的计算优化。我探讨了如何通过Local Attention、Stride Attention、Global Attention等方法减少计算量。此外，还介绍了自动选择注意力计算和Attention Matrix的线性组合方法。最后，补充了Batch Normalization的知识，为模型训练提供了更好的稳定性。

Abstract

This week’s focus is on optimizing the computation of the self-attention mechanism. I explored methods like Local Attention, Stride Attention, and Global Attention to reduce computational load. Additionally, we discussed automatic selection of attention computation and linear combinations in the Attention Matrix. Lastly, we supplemented our understanding with Batch Normalization, enhancing model training stability.

一、机器学习部分

1.1 self-attention的计算量

如果现在自注意力模型输入的序列长度为$N$，则对应的Query为$N$个，对应的Key也为$N$个。它们之间相互计算关联性（即注意力分数），可以得到上图中的Attention Matrix，这个矩阵的复杂度是$N^2$，当$N$的数值很大时，该矩阵的计算量就会变得很大。因此，这一节介绍多种方法以加速计算Attention Matrix的计算。

Notice：当$N$很大时，self-attention的计算才会主导整个模型中计算量。例如：在Transformer模型中，除了self-attention还有其他模块的计算量，self-attention模块的计算量占模型整体计算量是与$N$有关的，当$N$过小时，对self-attention的改进计算并不会明显提高Transformer模型的运算速度。

1.2 人类理解代替自注意力计算

根据人类对问题的理解，对Attention Matrix某些位置的值直接赋值，跳过计算步骤，从而减少计算量。

1.2.1 Local Attention/Truncated Attention

计算self-attention时，并非计算整个序列间的self-attention分数，而是只看自己和左右的邻居，其他的关联性都设定为0。下图在Attention Matrix中，表示为灰色的部分都人工设定为0，只计算蓝色部分的self-attention分数。这种方法叫做Local Attention或Truncated Attention。

Local Attention与CNN较为相似，主要体现在它们的局部关注机制上。这种机制使得模型在处理输入数据时，只关注输入数据的局部区域，而不是整体。卷积神经网络（CNN）中，卷积层通过滑动窗口的方式在输入数据上提取特征。这种操作也可以看作是一种局部关注机制，通过卷积核仅关注输入数据的局部区域来提取特征。Local attention相比于之前介绍的包含全序列的注意力，更加注重输入数据的局部关系，与卷积核的滑动也很类似。

1.2.2 Stride Attention

根据自己对问题的理解，计算局部的self-attention并不一定是左右邻居，如下图，可以是分别计算序列中两步前或两步后的关联性，也可以是分别计算序列中一步前或一步后的关联性，灰色的地方设定为0。这种方法叫做Stride Attention。

1.2.3 Global Attention

前面介绍的方法都是以某一个位置为中心，分别计算左右的关联性。Global Attention注重于整个序列，其会添加特殊的token到原始的序列中，特殊的token分别与整个序列计算self-attention，具体做法有两种：

• 从原来的token序列中，选择一部分作为特殊的token。
• 外加一部分额外的token。

从上图的Attention Matrix观察得到，在原始的序列中，第一和第二个位置被选择为特殊的token。从横轴的角度看，第一和第二个位置的Query与整个序列的Key分别做了self-attention。从纵轴的角度看，序列每一个位置的Query都与第一和第二位置的Key做了self-attention。灰色的位置设定为0。

在Big Bird中提出了Random attention并且将其与前面的Local Attention和Global Attention一并融合。

1.2.4 聚类Query和Key

第一步，根据相似度聚类Query和Key，上图中根据不同颜色聚类为了4类。

第二步，相同类之间的Query和Key才做self-attention。

1.3 自动选择自注意力计算

通过神经网络学习出一个0-1矩阵，深色位置代表1，浅色位置代表0。只有深色位置计算self-attention，浅色位置不计算。

输入序列中的每一个位置都通过一个神经网络产生一个长度为$N$的向量，然后将这些向量拼起来得到大小为$N\times N$的矩阵。然而现在这个由向量拼成得到的矩阵中的值，是连续值，要转换为0-1矩阵，这一部分是可以微分的，所以可以通过学习得到，具体需要看Sinkhorn Sorting Network的论文。

1.4 Attention Matrix中的线性组合

计算Attention Matrix的Rank（秩），得到Low Rank，说明该矩阵的很多列是其它列的线性组合。由此可得，实际上并不需要$N\times N$的矩阵，目前$N\times N$的矩阵中包含很多重复的信息，也许可以通过减少Attention Matrix的大小（主要是列数量）实现减少运算量。

选择具有代表性的Key，得到K个Key，即得到大小为$N\times K$的Attention Matrix。接下来考虑self-attention这一层的输出，同样地要从N个Value中挑出具有代表性的K个Value，一个Key对应一个Value向量。然后用Value矩阵乘上Attention Matrix可以得到self-attention层的输出。

为什么我们不能挑出K个代表的Query呢？

输出序列的长度与Query的数量是一致的，如果减少Query的数量，输出序列的长度就会变短。

挑选具有代表性的Key的方法为：

卷积降维和线性组合（K个向量是N个向量的K种线性组合，下图右）

1.5 通过矩阵乘法推导自注意力计算

简要复习一下自注意力机制的矩阵计算过程：第一步，输入序列分别做三种不同的变换，得到$d\times N$大小的Query和$d\times N$大小的Key，其中$d$是Query和Key的维度，$N$代表序列的长度。并得到$d^{\prime }\times N$大小的Value，其中特别用$d^{\prime }$表示Value的维度，是因为Value的维度可以与Query、Key不一样。第二步，$K^{\mathrm{T}}$乘上$Q$得到Attention Matrix，然后通过softmax做归一化。第三步，用$V$乘上归一化后的Attention Matrix($A^{\prime }$)得到自注意力层的输出$O$。

如果我们先忽略softmax的操作，self-attention的计算方法就是上图中第一行的计算过程，现在考虑第二行运算，先算$V$乘上$K^{\mathrm{T}}$的结果，再乘上$Q$，这样的计算顺序与第一行有何不同？得到的结果是一样的，运算量是不一样的。

尽管$A\left(CP\right)=\left(AC\right)P$，但是第一种计算方式的计算量是$10^6$，第二种计算方式的计算量的$10^3$，两者计算量之间的差异很大。因此我们这里先忽略softmax操作，考虑self-attention中矩阵计算的改进。

根据上图证明，$V(K^{\mathrm{T}}Q)$的计算量通常大于$(V{K}^{\mathrm{T}})Q$的计算量。

接下来加入softmax，写出计算self-attention的数学表达式：

下面通过数学证明的角度说明更换矩阵乘法顺序，计算self-attention的过程：

还有一个问题是，$\exp (q\cdot k)\approx \Phi (q)\cdot \Phi (k)$是如何实现的，具体需要参考下面的论文。

1.6 Batch Normalization

在Transformer的编码器中使用到了Layer Normalization，在上一周的周报中并将其与Batch Normalization做了比较，这里特别补充Batch Normalization的知识。

做标准化的原因是，希望能把不同维度的特征值规范到同样的数值范围，从而使得error surface比较平滑，更好训练。

Batch Normalization是对不同特征向量的同一维度，计算平均值和标准差，然后将特征值减去平均值再除以标准差，实现标准化。标准化后，同一维度上的数值的平均值是0，方差是1，接近高斯分布。

在神经网络中，输入特征${\tilde{x}}^1$、${\tilde{x}}^2$、${\tilde{x}}^3$已经做过了标准化，在经过$W^1$层后，且输入$W^2$层之前仍需要做标准化。至于是对激活函数前的$z^1$、$z^2$、$z^3$还是之后的$a^1$、$a^2$、$a^3$做标准化，差别不是很大。以$z^1$、$z^2$、$z^3$为例，$z^1$、$z^2$、$z^3$都是向量，做标准化的方法如下：

$\mu =\frac{1}{3}\sum _{i=1}^3{z}^i$是对向量$z^i$中对应元素进行相加，然后取平均。$\sigma =\sqrt{\frac{1}{3}\sum _{i=1}^3{\left(z^i-\mu \right)}^2}$是向量$z^i$与$\mu$相减，然后逐元素平方，求和平均后，再对向量的逐元素开根号。如果直接看公式会有一些歧义，因为$z^i$、$\mu$、$\sigma$都是向量，其中的求和，平方，开根号都是对向量中逐元素操作。最后标准化公式为：

$${\tilde{z}}^i=\frac{z^i-\mu }{\sigma }$$

实际上，GPU的内存不足以把整个dataset的数据一次性加载。因此，只考虑一个batch中的样本，对一个batch中的样本做Batch Normalization。在inference中，不可能等到整个batch数量的输入才做推理，具体方法为：在训练时计算$\mu$和$\sigma$的moving average，训练时的第一个batch为${\mu }^1$，第二个batch为${\mu }^1$，直到第t个batch为${\mu }^t$，且不断地计算moving average：

$$\bar{\mu }\leftarrow p\bar{\mu }+\left(1-p\right){\mu }^t$$

inference中标准化的公式变为：

$${\tilde{z}}^i=\frac{z^i-\bar{\mu }}{\bar{\sigma }}$$

总结

通过本周的学习，我对自注意力机制的优化策略有了更深入的了解，不同的注意力方法提供了多样化的计算选择，有助于提高模型的效率。下周还会围绕自注意力机制进行拓展学习。