External Attention：外部注意力机制

Posted 2021-05-12 巴蜀秀才

External Attention：外部注意力机制

作者： elfin

---

目录

• 1、External Attention
  • 1.1 自注意力机制
  • 1.2 外部注意力机制
• 2、思考

---

​ 最近Transformer在CV领域的研究非常热，如ViT、BoTNet、External Attention等。使用Transformer的传统印象就是慢，这种慢还往往是我们不能接受的推理速度。在最近的实验中，基于单张2080TI，使用ResNet-34实现了batch_size为\\(288\\)，图片大小为\\(32\\times200\\)，每个epoch的训练用时大约为：5个小时；同样的实验环境下，我实验了Swin Transformer的结构，在stage1实验卷积进行下采样(实现embedding的效果)，stage2~stage4使用Swin block，block的数量分别为[2, 8, 2]，stage5使用两层卷积进行简单的任务适配。这个结构在训练时，batch_size设置到\\(42\\)，图片大小修改为\\(32\\times240\\)，每个epoch的训练用时大约为17个小时。Swin Transformer原作者声称其比卷积网络性能好，速度快，经过实验，实际上CNN的推理速度要比Transformer快，性价比方面Transformer确实是更高一些，但是其庞大的参数量也导致模型不太好训练。经过实验对比，在我的数据集上，Swin Transformer提升了将近7个百分点(epoch较少)。

​ 虽然自注意力有那么多优点，但是这推理速度慢，耗资源的特点让我们这些缺硬件的渣渣很难受，然后就在前几天，突然看见清华、谷歌等大佬团队在做\\(Transformer\\rightarrow MLP\\)的操作，当然这里不是指简单返祖现象，以清华的工作为例，它将二次复杂度降到线性复杂度，这个特性有成功吸引到我，所以，下面是清华的External Attention的相关情况。

---

Top  ---  Bottom

1、External Attention

1.1 自注意力机制

首先回顾 self-attention机制，常见的自注意力如图所示：

对于给定的输入特征图\\(F\\in \\mathbb{R}^{N\\times d}\\)，其中\\(N\\)是像素个数，\\(d\\)是特征维度；自注意力产生了查询矩阵\\(Q\\in \\mathbb{R}^{N\\times {d}\'}\\)，键矩阵\\(K\\in \\mathbb{R}^{N\\times {d}\'}\\)，和值矩阵\\(V\\in \\mathbb{R}^{N\\times d}\\)。则自注意力机制的计算公式为：

\\[\\begin{align} A = \\left ( \\alpha _{i,j} \\right ) &= softmax\\left ( QK^{T} \\right )\\\\ F_{out} &= AV \\end{align} \\]

简化版的注意力机制如下：

它的计算公式为：

\\[\\begin{align} A &= softmax\\left (FF^{T} \\right )\\\\ F_{out} &= AF \\end{align} \\]

基于上面的公式不难发现，注意力特征图是计算像素级的相似度，输出是输入的精确特征表示。

即使公式(3)(4)简化了，但是其计算复杂度仍然是\\(O\\left ( dN^{2} \\right )\\)。我们注意到随着特征图的增大，计算量的增大往往是我们无法接受的。以我上面的实验为例：

• \\(32\\times200\\)的计算量为：\\(192\\times\\left ( 32\\times200 \\right )^{2}=7864320000\\)
• \\(32\\times240\\)的计算量为：\\(192\\times\\left ( 32\\times240 \\right )^{2}=11324620800\\)

大约增加了\\(40\\%\\)的计算量，这还只是一层，当然上面的计算方式有待商榷，但这也说明了计算量的增长是非常迅猛的，而且在高级语义特征图中，通道数会更大，一定程度上几乎会导致在所有层级上计算量都有很大的提升。当然这个增长是呈现指数倍率增长，所以我们需要提升精度的同时，提升速度、优化资源占用。

---

Top  ---  Bottom

1.2 外部注意力机制

经过实验得知：自注意力机制是一个\\(N-to-N\\)的注意力矩阵，可视化像素之间的关系，可以发现这种相关性是比较稀疏的，即很多是冗余信息。因此清华团队提出了一个外部注意力模块。

它的注意力计算是在输入像素与一个外部记忆单元\\(M\\in \\mathbb{R}^{S\\times d}\\)之间：

\\[\\begin{align} A &= Norm\\left (FM^{T} \\right )\\\\ F_{out} &= AM \\end{align} \\]

注意与自注意力机制不同，上式(5)是第\\(i\\)个像素点与\\(M\\)第\\(j\\)行的相似度。这里\\(M\\)是一个输入Input的可学习相关性参数，作为训练数据集的全局记忆。\\(A\\)是从先验信息得来的注意力特征图，Norm操作和自注意力一样。最终，通过\\(A\\)来更新\\(M\\)。

另外，我们用两种不同的记忆单元：\\(M_{k}\\)和\\(M_{v}\\)来增加网络的建模能力。

最终外部注意力机制的计算公式为：

\\[\\begin{align} A &= Norm\\left (FM^{T}_{k} \\right )\\\\ F_{out} &= AM_{v} \\end{align} \\]

经过上面的公式，外部注意力机制的复杂度是\\(O(dSN)\\)。这里的\\(d\\)、\\(S\\)是一个超参数，经过实验，作者发现\\(S\\)设置为64效果挺好。因此外部注意力机制比自注意力机制更高效，并且它可以直接应用于大尺寸的输入。

---

Top  ---  Bottom

2、思考

1. 计算量的降低和特征图的大小关系很大，这里的外部注意力机制的\\(S\\)设置为64，那么特征图的像素个数低于64时，实际它的计算量就更大。
2. Swin Transformer是基于窗口进行自注意力计算的，window size作者使用的是7，像素个数\\(N=49\\)，这也小于\\(64\\)，所以同等情况下，外部注意力会有更大的计算量。

---

Top  ---  Bottom

完！