Vision Transformer 必读系列之图像分类综述(二): Attention-based
时间：2022-12-30 09:00:01
文 @ 000007
号外号外：awesome-vit 新的，欢迎大家 Star Star Star ~
https://github.com/open-mmlab/awesome-vitgithub.com/open-mmlab/awesome-vit
0 前言
在Vision Transformer 必读系列图像分类综述(1):概述一文中对 Vision Transformer 在图像分类中的发展进行了概述性总结，本文则对其中涉及的 Attention-based 详细说明部分。下一篇文章将解释概述中涉及的其他部分。
ViT 如下图所示下图所示：
注：本文涉及的思维导图可通过https://github.com/open-mmlab/awesome-vit下载。
1 Transformer
论文题目： Attention is All You Need
论文地址：https://arxiv.org/abs/1706.03762
Transformer 结构是 Google 在 2017 年为解决机器翻译任务(例如英文翻译为中文)而提出，从题目中可以看出主要是靠 Attention 注意机制最大的特点是放弃了传统 CNN 和 RNN，整个网络结构完全由 Attention 机制组成。因此，在分析模型结构之前，需要解释什么是注意机制。
1.1 Attention 注意力机制
生来就有注意力机制，看看任何图片，我们都会自动关注特定位置的特定物体。 Attention 该机制也具有相同的含义。对于所需的任何模式，无论是图像、文本、点云还是其他模式，我们都希望网络能够通过训练自动关注有意义的位置，如图像分类和检测任务。
注意机制并不是一个新概念，它已经广泛应用于视觉算法中，如典型的SENet。
利用 Squeeze-and-Excitation 模块计算注意力权重概率分布，然后在特征图上重新加权每个通道。
可以举一个更简单的例子，假设有一个训练好的分类网络，输入一张图片，训练好的分类网络权重 W 和图片 X 注意力计算，从 X 提取有助于分类的特征，最终可作为分类的依据。W 和 X 都是矩阵，要想利用 W 实现重加权的矩阵 X 等价于计算的目的 W 和 X 相似度(点乘)，然后将相似度转换为权重概率分布，再次作用于 X 上就可以。
以简单的猫狗二分类为例。网络的最终输出是 2x1 向量，第一个数字表示猫类，否则是狗类，假设网络已经训练好了，它 W 为 shape 为 2x1 值为[[0.1, 0.5]]，X 表示输入图片 shape 也是 2x1，其值为[[0.1, 0.8]]，可见其类别为狗，采用以下计算步骤即可正确分类：
- W 和 X 转移相乘，即计算 W 中每个值和 X 获得每个值的相似度 2x2 矩阵，值为 [[0.01,0.08], [0.05,0.4]]。
- 第二个维度 Softmax，将其转化为概率权重图 [[0.4825, 0.5175], [0.4134, 0.5866]]。
- 乘以上述概率权重 X，得到 shape 为 2x1 输出，值为 [[0.4622, 0.5106]]。
- 这时，由于第二个值大，正确分类为狗。
X 它是一个包含狗的图像矩阵。正确分类的前提是训练良好 W 矩阵中的第二个数大于第一个数。上述过程可以简单地理解计算 W 和 X 如果两个向量相似(都是第二个大于第一个数)，那么就分为狗，否则就分为猫。
import torch W = torch.tensor([[0.1, 0.5]]).view([1, -1, 1]) X = torch.tensor([[0.1, 0.8]]).view([1, -1, 1]) # 1 计算两个向量相似度 attn_output_weights = torch.bmm(W, X.transpose(1, 2)) # 2 转换为概率分布 attn_output_weights = torch.softmax(attn_output_weights, dim=-1) # 3 注意力加权 cls = torch.bmm(attn_output_weights, X)[0] print(cls)
以下公式可以表示上述计算过程：
对应上述例子，Q 训练好 W 矩阵，K 图片输入，V 和 K 相等的一般解释是使用 Q 查询矩阵和 K 矩阵计算相似度，然后转换为概率分布。此时，概率值大的位置表示两者相似度大的部分，然后乘以概率分布 V 用注意力权重分布加权的值矩阵 V 矩阵，也就改变了 V 矩阵本身的分布。如果注意力机制训练的很好，那么提取的 V 这应该是我们想要的信息。若注意机制训练良好，则提取 V 应该是我们想要的信息。 d_k 当向量值很大时，平方根是为了避免梯度消失，Softmax 函数将几乎所有的概率分布分配到最大值对应的位置，即所谓的锐化，可以有效避免梯度消失和稳定训练过程。
上述公式是论文中最重要的公式Scaled Dot-Product Attention计算公式首先使用点乘计算 QK 除以分母 d_k 平方根进行 Scaled 操作，然后 Softmax 操作将其转换为概率乘 V 实现 Attention 功能。
需要注意： 为使上述公式不报错，其 Shape 关系必须为Q - (N, M)，K - (P, M),V - (P, G), 一般来说 K 和 V Shape 相同，但是 Q Shape 不一定和 K 相同。注意力层的计算复杂性可以通过灵活地改变这些维度来控制，大多数后续改进算法都使用这一点。
1.2 Transformer 结构分析
Transformer 它是为了解决机器翻译任务而提出的。机器翻译是一个历史悠久的问题，可以理解为序列转序列问题，我们常说 seq2seq 一般采用结构解决这类问题 encoder-decoder 结构，Transformer 这种结构也被使用。翻译任务的常规解决方案如下：
对应到 Transformer 更具体的结构之一是：
编码器主要包括自注意模块(QKV 同一输入)类似于前向网络，解码器类似于编码器，但内部有更多的交叉注意模块与解码器交互。
一般来说，标准 Transformer 包括 6 个编码器和 6 串行解码器。
- 编码器内部接收源翻译输入序列，通过自注意模块提取必要特征，通过前网络进一步抽象特征。
- 解码器端输入包括两部分，一个是自注意模块提取的目标翻译序列的特征，另一个是编码器提取的整体特征。这两个输入特征的向量将交叉计算，提取有利于目标序列分类的特征，然后通过前网络进一步抽象特征。
- 堆叠多个编码器和解码器，下一个编码器接收上一个编码器的输出，形成串行结构，最后使用解码器输出进行分类。
Transformer 完整结构如下：
编码器的基本组件包括：嵌入模块的源句子词 Input Embedding、位置编码模块 Positional Encoding、多头自注模块 Muti-Head Attention、前向网络模块 Feed Forward 以及必要的 Norm、Dropout 和残模块。
类似的解码器基本组件包括：嵌入模块的目标句子词 Output Embedding、位置编码模块 Positional Encoding、带 mask 的自注意力模块 Masked Muti-Head Attention、交叉互注意模块 Muti-Head Attention、前网络模块 Feed Forward 、分类头模块 Linear Softmax 以及必要的 Norm、Dropout 和残差模块。
因为本文的重点是分析视觉 Transformer，因此，没有必要深入分析机器翻译过程。读者只需要了解每个模块的功能，并进行视觉分类 Transformer 任务和 NLP 与解码器模块相比，机器翻译任务不同。 NLP 任务要简单得多。
1.2.1 编码器基本组件
(1) 嵌入模块的源句子词 Input Embedding
机器翻译是句子输入和句子输出。每个句子都由单词组成。将句子编码成程序可以理解的向量过程称为单词嵌入过程，通常称为 Word2Vec，在图像中对应称为 Token 化学过程是如何将图像转换为更语义的 Token，Token 概念会在 ViT 详细描述。
在 1.1 小节已经详细说明了注意力计算过程。左边是最简单的 Scaled Dot-Product Attention，单纯看上图你可以发现没有任何可学习参数，那么其存在的意义是啥？ 实际上可学习参数在 QKV 映射矩阵中，在自注意力模块中会对输入的向量分别乘上可学习映射矩阵 W_Q、W_K 和 W_V，得到真正的 Q、K 和 V 输入，然后再进行 Scaled Dot-Product Attention 计算。
为了增加注意力特征提取的丰富性，不会陷入某种局部特性中，一般会在注意力层基础上(单头注意力层)引入多个投影头，将 QKV 特征维度平均切分为多个部分(一般分成 8 部分)，每个部分单独进行自注意力计算，计算结果进行拼接 。在特征维度平均切分，然后单独投影、计算，最后拼接可以迫使提取的注意力特征更加丰富。也就是上面的多头注意力模块 Multi-Head Attention。
(3) 前向网络模块 Feed Forward
前向网络模块主要是目的是对特征进行变换，其单独作用于每个序列(只对最后一个特征维度进行变换)。由于没有结构图，故直接贴相关代码，包括两个 Position-wise FC 层、激活层、Dropout层和 LayerNorm 层。
class PositionwiseFeedForward(nn.Module):
''' A two-feed-forward-layer module '''
def __init__(self, d_in, d_hid, dropout=0.1):
super().__init__()
# 两个 fc 层，对最后的维度进行变换
self.w_1 = nn.Linear(d_in, d_hid) # position-wise
self.w_2 = nn.Linear(d_hid, d_in) # position-wise
self.layer_norm = nn.LayerNorm(d_in, eps=1e-6)
self.dropout = nn.Dropout(dropout)
def forward(self, x):
residual = x
x = self.w_2(F.relu(self.w_1(x)))
x = self.dropout(x)
x += residual
x = self.layer_norm(x)
return x
(4) Norm、Dropout 和残差模块
在每个注意力层后面和前向网络后都会接入 Dropout 、残差模块和 Layer Norm 模块。这些必要的措施对整个算法的性能提升非常关键。至于为啥用 Layer Norm 而不是 Batch Norm ，原因是机器翻译任务输入的句子不一定是同样长的，组成 Batch 训练时候会存在大量 Padding 操作，如果在 Batch 这个维度进行 Norm 会出现大量无效统计，导致 Norm 值不稳定，而 Layer Norm 是对每个序列单独计算，不考虑 Batch 影响，这样比较符合不定长序列任务学习任务。当然如果换成图像分类任务，则可以考虑使用 BN 层，后续有算法是直接采用 BN 的。
(5) 位置编码 Positional Encoding
考虑一个分类任务，输入一段句子判断是疑问句还是非疑问句？现在有两条语句分别是：
- 不准在地铁上吃东西
- 在地铁上吃东西准不
自注意力层的计算不会考虑字符间的顺序，因为每个字符都是单独和全局向量算相似度，也就是说上面两个句子输入进行注意力计算，输出的向量值是相同的，只不过相对位置有变化。如果我们对输出向量求和后值大于 0 还是小于 0 作为分类依据，那么上面两个句子输出相加值是完全相同的，那就始终无法区分到底是疑问句还是非疑问句，这就是我们常说的 Transformer 具有位置不变性。要解决这个问题，只需要让模型知道输入语句是有先后顺序的，位置编码可以解决这个问题。
加入位置信息的方式非常多，最简单的可以是直接将输入序列中的每个词按照绝对坐标 0,1,2 编码成相同长度的向量，然后和词向量相加即可。作者实际上提出了两种方式：
- 网络自动学习，直接全 0 初始化向量，然后和词向量相加，通过网络学习来学习位置信息。
- 自己定义规则，规则自己定，只要能够区分输入词顺序即可，常用的是 sincos 编码。
实际训练选择哪一种位置编码方式发现效果一致，但是不管哪一种位置编码方式都应该充分考虑在测试时候序列不定长问题，可能会出现测试时候非常长的训练没有见过的长度序列，后面会详细说明。
1.2.2 解码器基本组件
其大部分组件都和编码器相同，唯一不同的是自注意力模块带有 mask，还额外引入了一个交叉注意力模块以及分类头模块。
(1) 带 mask 的自注意力模块
注意这个模块的输入是目标序列转化为词向量后进行自注意力计算。机器翻译是一个 seq2seq 任务，其真正预测是：最开始输入开始解码 token 代表解码开始，解码出第一个词后，将前面已经解码出的词再次输入到解码器中，按照顺序一个词一个词解码，最后输出解码结束 token，表示翻译结束。
也就是当解码时，在解码当前词的时候实际上不知道下一个词是啥，但在训练时，是将整个目标序列一起输入，然而注意力计算是全局的，每个目标单词都会和整个目标句子计算自注意力，这种训练和测试阶段的不一致性无法直接用于预测。为此我们需要在训练过程中计算当前词自注意力时候手动屏蔽掉后面的词，让模型不知道后面词。具体实现就是输入一个 mask 来覆盖掉后面的词。
由于这种特性是只存在于 NLP 领域，图片中不存在，故不再进行更深入分析。
(2) 交叉注意力
交叉注意力模块和自注意力模块相同，只不过其 QKV 来源不同，Q 来自解码器，KV 来自编码器，交叉注意力模块会利用 Q 来提取编码器提取的特征 KV，然后进行分类。
(3) 分类头
分类头就是普通的线性映射，转换输出维度为分类个数，然后采用 CE Loss 进行训练即可。
1.3 总结
Transformer 结构内部存在多个组件，但是最核心的还是注意力模块，在原始论文中作者也引入了大量的可视化分析来证明注意力模块的作用，有兴趣的建议阅读原文。可能作者自己也没有想到这篇论文会在视觉领域引起另一个全新的风尚，开辟出一条新的看起来前途一片光明的道路。
图片来自 A Survey of Visual Transformers。
2 Vision Transformer
论文题目： An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
论文地址： https://arxiv.org/abs/2010.11929
ViT 是第一篇成功将 Transformer 引入到视觉领域且成功的尝试，开辟了视觉 Transformer 先河。其结构图如下所示：
其做法非常简单，简要概况为：
- 将图片分成无重叠的固定大小 Patch (例如 16x16)，然后将每个 Patch 拉成一维向量， n 个 Patch 相当于 NLP 中的输入序列长度(假设输入图片是 224x224，每个 patch 大小是 16x16，则 n 是 196)，而一维向量长度等价于词向量编码长度(假设图片通道是 3, 则每个序列的向量长度是 768) 。
- 考虑到一维向量维度较大，需要将拉伸后的 Patch 序列经过线性投影 (nn.Linear) 压缩维度，同时也起到特征变换功能，这两个步骤可以称为图片 Token 化过程 (Patch Embedding) 。
- 为了方便后续分类，作者还额外引入一个可学习的 Class Token，该 Token 插入到图片 token 化后所得序列的开始位置 。
- 将上述序列加上可学习的位置编码输入到 N 个串行的 Transformer 编码器中进行全局注意力计算和特征提取，其中内部的多头自注意模块用于进行 Patch 间或者序列间特征提取，而后面的 Feed Forward(Linear+ GELU+Dropout+ Linear+ Dropout) 模块对每个 Patch 或者序列进行特征变换。
- 将最后一个 Transformer 编码器输出序列的第 0 位置( Class Token 位置对应输出)提取出来，然后后面接 MLP 分类后，然后正常分类即可。
可以看出图片分类无需 Transformer 解码器，且编码器几乎没有做任何改动，针对图像分类任务，只需单独引入一个 Image to Token 操作和 Class Token 的概念即可。
如何理解 Token？ 个人觉得任何包括图片更加高级的语义向量都可以叫做 Token，这个概念在 NLP 中应用非常广泛，表征离散化后的高级单词语义，在图像中则可以认为是将图像转化为离散的含更高级语义的向量。
ViT 证明纯 Transformer 也可以取得非常好的效果，相比 CNN 在数据量越大的情况下优势更加明显，但是 ViT 也存在如下问题：
- 不采用超大的 JFT-300M 数据集进行预训练，则效果无法和 CNN 媲美，原因应该是 Transformer 天然的全局注意力计算，没有 CNN 这种 Inductive Bias 能力，需要大数据才能发挥其最大潜力。
- ViT 无法直接适用于不同尺寸图片输入，因为 Patch 大小是固定的，当图片大小改变，此时序列长度就会改变，位置编码就无法直接适用了，ViT 解决办法是通过插值，这种做法一般会造成性能损失，需要通过 Finetune 模型来解决，有点麻烦。
- 因为其直筒输出结构，无法直接应用于下游密集任务。
后面的文章对上述缺点采用了各种各样的改进，并提出了越来越先进的处理手段，推动了视觉 Transformer 的巨大进步。
3 全局概述
由于内容非常多，为了更容易理解，我在拆解模块的基础上对每个模块进行分析，而不是对某篇文章进行概括。综述部分的分析流程按照结构图顺序描述，我将近期图像分类视觉 Transformer 发展按照 ViT 中是否包括自注意层模块来划分，包括：
- Attention-based， 这类算法是目前主流研究改进方向，包括了 Transformer 中最核心的自注意力模块
- MLP-based， 这类算法不需要核心的自注意力模块，而是简单的通过 MLP 代替，也可以取得类似效果
- ConvMixer-based，这类算既不需要自注意力模块，也不是单纯依靠 MLP，而是内部混合了部分 Conv 算子来实现类似功能
- General architecture analysis，在这三类算法基础上也有很多学者在探讨整个 Transformer 架构，其站在一个更高的维度分析问题，不局限于是否包括自注意力模块，属于整体性分析
在上述三个方向中，Attention-based 是目前改进最多，最热门的，也是本综述的核心。本文按照 3 个分类顺序依次分析，最后进行通用架构分析。通过 General architecture analysis 部分可以深化Attention-based、MLP-based 和 ConvMixer-based 三者的联系和区别。本文仅仅涉及 Attention-based 部分。
3.1 Attention-based
Attention-based 表示这类算法必然包括注意力模块，我们将按照广度优先顺序进行一次分析。
继 ViT 后，我们将其发展分成两条线路：训练策略和模型改进。
其中训练策略表示目前主流对 ViT 模型的训练改进方式，而模型改进则是对各个部件进行改进。
- 训练策略包括两篇论文：DeiT 和 Token Labeling。两者提出的出发点一致，都是为了克服 ViT 需要 JFT-300M 大数据集进行预训练的缺点。DeiT 是通过引入蒸馏学习解决，而 Token Labeling 通过引入显著图然后施加密集监督解决。后续发展中大部分算法都是参考了 DeiT 的训练策略和超参设置，具有非常大的参考价值。
- 模型改进方面，我将其分成了 6 个组件以及其他方面的改进，6 个组件包括：
- Token 模块，即如何将 image 转 token 以及 token 如何传递给下一个模块
- 位置编码模块
- 注意力模块，这里一般都是自注意力模块
- Fead Forward (FFN) 模块
- Norm 模块位置
- 分类预测头模块
下面按照训练策略和模型改进顺序分析。
3.1.1 训练策略
训练策略解决 ViT 需要大数据先预训练问题以及超参有待优化问题。
3.1.1.1 DeiT
如果说 ViT 开创了 Transformer 在视觉任务上面的先河，那么 DeiT 的出现则解决了 ViT 中最重要的问题：如果不采用超大的 JFT-300M 数据集进行预训练，则效果无法和 CNN 媲美。在单个节点 8 张 V100 且无需额外数据的情况下，用不到 3 天的时间训练所提的 ViT（86M 参数），在 ImageNet 上单尺度测试达到了 83.1% 的 top-1 准确率。
DeiT 核心是引入蒸馏手段加上更强的 Aug 和更优异的超参设置。其蒸馏的核心做法如下所示：
ViT 的 Class Token 是加到图片输入序列的前面，那么蒸馏 Token 可以插到输入序列的后面，当然插入到哪个位置其实无所谓，你也可以插入到 Class Token 后面，经过 Transformer 编码器输出的序列相比 ViT 也会多一个，然后额外的一个输出 Token 经过线性层输出相同类别通道，最后进行蒸馏学习。
对于蒸馏学习来说，做法通常有两个：
- Soft 蒸馏，即学生模型和教师模型预测的 Softmax 概率分布值计算 KL Loss。
- Hard 蒸馏，即教师模型预测的 Softmax 概率分布值中，值最大对应的类别作为标签，然后和学生模型预测的 Softmax 概率分布值计算 CE Loss。
蒸馏学习中，通常教师模型会选择一个比学生模型性能更强的且已经提前训练好的模型，教师模型不需要训练，通过蒸馏 loss 将教师模型知识以一种归纳偏置的方式转移给学生模型，从而达到提升学生模型性能的目的。因为引入了额外的蒸馏 Token，而且该 Token 训练任务也是分类，所以实际上 DeiT 在推理时，是将 Class Token 和 Distillation Token 的预测向量求平均，再转换为概率分布。
为了证明 Distillation Token 的有效性，而不是只由于多了一个 Token 或者说多了一个可学习参数导致的，作者还做了对比试验，不加 Distillation Token，而是再加一个 Class Token，相当于有两个分类头，两个 Token 独立且随机初始化，实验发现他们最终收敛后两个分类 Token 的相似度达到 0.999，并且性能更弱，这样证明了加入 Distillation Token 的意义。
通过大量实验，作者总结了如下结论：
- 蒸馏做法确实有效，且 Hard 蒸馏方式效果会更好，泛化性能也不错。
- 使用 RegNet 作为教师网络可以取得更好的性能表现，也就是说相比 Transformer，采用卷积类型的教师网络效果会更好。
除了上述蒸馏策略，还需要特别注意 DeiT 引入了非常多的 Aug 并且提供了一套更加优异的超参，这套参数也是后续大部分分类模型直接使用的训练参数，非常值得学习，如下所示：
总而言之 DeiT 算法非常优异，实验也非常多(建议去阅读下)，最大贡献是通过蒸馏策略省掉了 ViT 需要 JFT-300M 数据集进行预训练这个步骤，并且提供了一套非常鲁棒且实用的超参配置，深深地影响了后续的大部分图像分类视觉 Transformer 模型。
3.1.1.2 Token Labeling
DeiT 不是唯一一个解决 ViT 需要大数据量问题的算法，典型的还有 Token Labeling，其在 ViT 的 Class Token 监督学习基础上，还对编码器输出的每个序列进行额外监督，相当于将图片分类任务转化为了多个输出 Token 识别问题，并为每个输入 Patch 的预测 Token 分配由算法自动生成的基于特定位置的监督信号，简要图如下所示：
从上图明显可以看出，相比 ViT 额外多了输出 Token 的监督过程，这些监督可以当做中间监督，监督信息是通过 EfficientNet 或者 NFNet ( F6 86.3% Top-1 accuracy) 这类高性能网络对训练图片提前生成的显著图，每个显著图维度是和类别一样长的 C 维，辅助 Loss 和分类一样也是 CE Loss。当然最终实验结果表明性能比 DeiT 更优异，而且由于这种密集监督任务，对于下游密集预测任务泛化性也更好。
在此基础上 DeiT 已经证明通过对 ViT 引入更多的强 Aug 可以提升性能，例如引入 CutMix，但是本文的做法无法直接简单增加 CutMix，为此作者还专门设计了一个 MixToken，大概做法是在 Pathc Embedding 后，对 Token 进行了相应的 CutMix 操作。性能表如下所示：
LV-ViT 即为本文所提模型。相比 DeiT，作者认为本文做法更加优异，体现在：
- 不需要额外的教师模型，是一个更加廉价的做法。
- 相比于单向量监督，以密集的形式监督可以帮助训练模型轻松发现目标物体，提高识别准确率，实验也证明了对下游密集预测任务（例如语义分割）更友好。
下表是对训练技术的详细分析：
简而言之，Token Labeling 的核心做法是通过引入额外的显著图来监督每个 patch 输出的预测 token，虽然不需要教师模型，但是依然需要利用更优异的模型对所有训练图片生成显著图。
3.1.2 模型改进
在 DeiT 提出后，后续基于它提出了大量的改进模型，涉及到 ViT 的方方面面。前面说过 ViT 模型主要涉及到的模块包括：
- Token 模块，即如何将 image 转 token 以及 token 如何传递给下一个模块
- 位置编码模块
- 注意力模块，这里一般都是自注意力模块
- Fead Forward (FFN) 模块
- Norm 模块位置
- 分类预测模块
3.1.2.1 Token 模块
Token 模块包括两个部分：
- Image to Token 模块即如何将图片转化为 Token，一般来说分成有重叠和无重叠的 Patch Embedding 模块
- Token to Token 模块即如何在多个 Transformer 编码器间传递 Token，通常也可以分成固定窗口 Token 化过程和动态窗口 Token 化过程两个
下面是完整结构图：
3.1.2.1.1 Image to Token 模块
首先需要明确：Patch Embedding 通常包括图片窗口切分和线性嵌入两个模块，本小结主要是说图片窗口切分方式，而具体实现不重要，常用的 2 种实现包括 nn.Conv 和 nn.Unfold，只要设置其 kernel 和 stride 值相同，则为非重叠 Patch Embedding，如果 stride 小于 kernel 则为重叠 Patch Embedding。
(1) 非重叠 Patch Embedding
ViT 和目前主流模型例如 PVT 和 Swin Transformer 等都是采用了非重叠 Patch Embedding，其简要代码为：
# 非重叠只需要设置Conv kernel_size 和 stride 相同即可
_conv_cfg = dict(
type='Conv2d', kernel_size=16, stride=16, padding=0, dilation=1)
_conv_cfg.update(conv_cfg)
self.projection = build_conv_layer(_conv_cfg, in_channels, embed_dims)
x = self.projection(x).flatten(2).transpose(1, 2)
通过设置 16x16 的 kernel 和 stride 可以将图片在空间维度进行切割，每个 patch 长度为 16x16x3，然后将每个 Patch 重排拉伸为一维向量后，经过线性层维度变换，输出 shape 为 (B, Num_Seq, Dim)。
在 TNT 中作者提出了一种更精细的非重叠 Patch Embedding 模块，如下图所示：
他的基本观点是自然图像的复杂度相较于自然文本更高，细节和颜色信息更丰富，而 ViT 的非重叠 Patch Embedding 做法过于粗糙，因为后续自注意力计算都是在不同 Patch 间，这会导致 Patch 内部的局部自注意力信息没有充分提取，而这些信息在图像中也是包含了不同的尺度和位置的物体特征，是非常关键的。故我们不仅要考虑 Patch 间自注意力，还要考虑 Patch 内自注意力，为此作者在 外层 Transformer 中又嵌入了一个内层 Transformer，相应的非重叠 Patch Embedding 也分成两步：整图的非重叠 Patch Embedding 过程和 Patch 内部更细粒度的非重叠 Patch Embedding 过程。
通过上图大概可以看出其具体做法，内部相当于有两个 Transformer，第一个 Transformer (Outer Transformer )和 ViT 完全一样，处理句子 Sentences 信息即图片 Patch 级别信息，第二个 Transformer (Inner Transformer，也需要额外加上 Inner Transformer 所需要的位置编码) 处理更细粒度的 Words 信息即图片 Patch 内再切分为 Patch，为了能够将两个分支信息融合，内部会将 Inner Transformer 信息和 Outer Transformer 相加。将上述 Transformer block 嵌入到 PVT 模型中验证了其对下游任务的适用性，通过进一步的可视化分析侧面验证了分类任务上 TNT 相比 DeiT 的优异性。
(2) 重叠 Patch Embedding
在常规的 CNN 网络中一般都是采用重叠计算方式，为此是否采用重叠 Patch Embedding 会得到更好的性能？
直接将非重叠 Patch Embedding 通过修改 Unfold 或者 Conv 参数来实现重叠 Patch Embedding 功能的典型算法包括 T2T-ViT 和 PVTv2，这两个算法的出发点都是非重叠 Patch Embedding 可以加强图片 Patch 之间的连续性，不至于出现信息断层，性能应该会比重叠 Patch Embedding 高。PVTv2 内部采用 Conv 实现，而 T2T ViT 是通过 Unfold 方式实现(论文中称为 soft split)。
前面说过 CNN 网络中一般都是采用重叠计算方式，那么是否可以用 ResNet Stem 替换非重叠 Patch Embedding过程，性能是否会更好？
在 Early Convolutions Help Transformers See Better 论文中作者进行了深度分析，虽然作者只是简单的将图片 Token 化的 Patch Embedding 替换为 ResNet Conv Stem，但是作者是从优化稳定性角度入手，通过大量的实验验证上述做法的有效性。作者指出 Patch Embedding 之所以不稳定，是因为该模块是用一个大型卷积核以及步长等于卷积核的卷积层来实现的，往往这个卷积核大小为 16*16，这样的卷积核参数量很大，而且随机性很高，从某种程度上造成了 Transformer 的不稳定，如果用多个小的卷积来代替则可以有效缓解。结构如下所示：
考虑了和 ViT 公平对比，新引入的 Conv Stem 计算量约等于一个 transformer block，故后续仅仅需要 L-1 个 transformer block。作者通过大量分析实验得到一些经验性看法：
(a) ViT 这类算法对 lr 和 wd 超参的选择非常敏感，而替换 Stem 后会鲁棒很多。
(b) ViT 这类算法收敛比较慢，而本算法会快很多，例如都在 100 epoch 处本文性能远优于 ViT。
ViT_p 即为 Patch Embedding 模式，ViT_c 即为 Conv Stem 模式，可以看出在不同 flops 下模型收敛速度都是 ViT_c 快于 ViT_p，虽然到 400 epoch 时候性能都非常接近。
(c) ViT 这类算法只能采用 AdamW 训练，而本文更加通用，采用 SGD 后性能没有显著下降。
众所周知，ViT 类模型都只能用 AdamW 训练，其占据显存是 SGD 的 3 倍，所以一般在 CNN 网络中都是用过 SGD 模型，性能通常不错，而通过替换 Patch Embedding 后也可以用 SGD 训练了
(d) 仅仅采用 ImageNet 训练 ViT 性能难以超越 CNN，而本文可以进一步提升 ViT 性能。
与上述论文持相同观点的也包括 ResT 、Token Learner、CSWin Transformer 等算法，他们都采用了完全相同的做法。 更进一步在 PS-ViT 中为了能够方便后续的渐进采样模块稳定提取更好的特征点，作者在 Image to Token 模块中不仅仅引入了 ResNet 的 Conv Stem 模块，还在后面再使用了 ResNet 第一个 stage 的前两个残差 block，在 Token to Token 模块中会详细说明 PS-ViT。
在 CeiT 中作者出发点是 CNN 中的诸多特性已经被证明是很成功的，纯粹的 Transformer 需要大量的数据、额外的监督才能达到和 CNN 相同的精度，出现这种问题的原因可能是 NLP 中的 Transformer 直接搬到图像任务中可能不是最合适的，应该考虑部分引入 CNN 来增强 Transformer。具体来说，在图片转 Token 方案中提出 Image-to-Tokens (I2T) 模块，不再是从图片中直接进行 Patch Emeding ，而是对 CNN 和 Pool 层所提取的底层特征进行 Patch Embedding，借助图像特征会比直接使用图片像素更好。
上图的上半部分是 ViT 的 Patch Embedding 过程，下图是 CeiT 所提出的做法，核心就是引入卷积操作提取底层特征，然后在底层特征上进行 Patch Embedding 操作。
既然采用 Conv Stem 可以解决很多问题，那么理论上经过精心设计的 Conv 结构也必然是有效的，例如 ViTAE 中就采用了空洞卷积做法，本质上是希望能够利用卷积提供多尺度上下文信息，这有助于后续模块信息提取，如下图所示：
对图片或者特征图应用多个不同空洞率的卷积提取信息后，进行拼接和 GeLU 激活函数后，直接拉伸为一维向量，从而转换为序列，并且由于空洞卷积可以实现下采样功能，故也可以有效地减少后续注意力模块计算量。
3.1.2.1.2 Token to Token 模块
大部分模型的 Token to Token 方案和 Image to Token 做法相同，但是也有些算法进行了相应改造。经过整理，我们将其分成两种做法：
- 固定窗口 Token 化
- 动态窗口 Token 化
固定窗口是指 Token 化过程是固定或者预定义的规则，典型的重叠和非重叠 Patch Embedding 就是固定窗口，因为其窗口划分都是提前订好的规则，不会随着输入图片的不同而不同，而动态窗口是指窗口划分和输入图片语义相关，不同图片不一样，是一个动态过程。
(1) 固定窗口 Token 化
这个做法通常和 Image to Token 模块完全一样，也可以分成非重叠 Patch Embedding 和重叠 Patch Embedding，大部分算法都属于这一类，例如 PVT、Swin Transformer 等。
(2) 动态窗口 Token 化
动态窗口 Token 化过程典型代表是 PS-ViT 和 TokenLearner。
前面说过，Vision Transformer with Progressive Sampling (PS-ViT) 中为了方便后续的渐进采样模块能够稳定提取更好的特征点，在 Image to Token 模块中不仅仅引入了 ResNet 的 Conv Stem 模块，还在后面再使用了 ResNet 第一个 stage 的前两个残差 block。在特征图 F 后，作者在 Token to Token 环节引入了一个渐进式采样模块，其出发点是 ViT 采用固定窗口划分机制，然后对每个窗口进行 Token 化，这种做法首先不够灵活，而且因为图片本身就是密集像素，冗余度非常高，采用固定划分方法对于分类来说可能就某几个窗口内的 Token 实际上才是有意义的，假设物体居中那么物体四周的 Token 可能是没有作用的，只会增加无效计算而已。基于此作者设计一个自适应采样的 Token 机制，不再是固定的窗口采样，而是先初始化固定采样点，如下图红色点所示，然后通过 refine 机制不断调整这些采样点位置，最终得到的采样点所对应的 Token 就是最有代表力的。其完整分类网络结构图如下所示：
得到特征图 F 后，经过渐进采样模块，不断 refine 采样点，最终输出和采样点个数个序列，将该序列作为 ViT 模型的输入即可。简单来看渐进采样模块起到了 Token to Token 作用。其中的渐进采样模块结构图如下所示：
详细计算过程如下：
- 首先图片经过 ResNet Conv Stem + ResNet 第一个 stage 的前两个残差块进行特征提取，得到 F
- 在特征图或者原图上先设置初始均匀固定间隔采样点 pt，上图是 9 个采样点，表示最终序列长度是 9
- 利用 pt 值对 F 进行采样，提取对应位置的特征向量，加上位置编码输入到编码器中，输出 T_t
- 将 T_t 经过一个 FC 层生成 offset，将该 offset 和初始位置 pt 相加就可以得到 refine 后的 p_t+1
- 将 3-4 步骤重复 N 次，下一个采样模块的输入包括 refine 后的 pt、特征图 F 和上一个采样模块的输出 T，三者相加
- 经过 N 次 refine 后，将该 token 序列拼接上 class token，然后再经过 M 个编码器模块
- 最后对 class token 对应位置输出 token 进行分类训练即可
可以发现，和 ViT 的主要差异就在于其采样点不是固定的均匀间隔，而是基于语义图自适应，从而能够在减少计算量的前提下进一步提升性能。PS-ViT 在 top-1 精度方面比普通 ViT 高 3.8%，参数减少约 4 倍，FLOP 减少约 10 倍，性能比较优异。
基于类似出发点，TokenLearner 提出可以基于空间注意力自适应地学习出更具有代表性的 token，从而可以将 ViT 的 1024 个 token 缩减到 8-16 个 token，计算量减少了一倍，性能依然可以保持一致，这也侧面说明了 ViT 所采样的固定窗口 token 有大量冗余，徒增计算量而已。其核心示意图如下所示：
假设想仅仅采样出 8 个 token，首先采用 Conv Stem 提取图片特征，然后分别输入到 8 个空间注意力分支中，空间注意力分支首先会应用一系列卷积生成空间 attention 图，然后逐点和输入特征相乘进行特征加权，最后通过空间全局 pool 层生成 1x1xC 的 Token，这样就将 HXWXC 的特征图转换为了 8 个通道为 C 的 Token。
为了进一步提高信息，作者还额外提出一个 TokenFuser 模块，加强 Token 和 Token 之间的联系以及恢复空间结构，整个分类网络的结构如下所示：(a) 为不包括 TokenFuser 的改进 ViT 结构，(b) 为包括 TokenFuser 的改进 ViT 结构。
从上述结构可以发现 TokenLearner 模块起到了自适应提取更具语义信息的 Token，并且还能够极大地减少计算量，而 TokenFuser 可以起到加强 Token 和 Token 之间的联系以及恢复空间结构的功能，TokenLearner+Transformer+ TokenFuser 构成 Bottleneck 结构。其中 TokenFuser 示意图如下所示：
其接收两个输入，一个是 TokenLearner 前的保持了空间信息的 1024 个 token 特征，一个是 TokenLearner 后经过自适应采样的 8 个 token 特征，然后以注意力模式两者进行乘加操作，融合特征以及恢复空间结构。
作者的分类实验依然采用了 JFT-300M 数据集进行预训练，然后在 ImageNet 1k上面微调，也就是说和最原始的 ViT 进行比较。
TokenFuser 也进行了相应的对比实验。
at 6 表示 TokenLearner 插入到第 6 个 Transformer Encoder 后。
3.1.2.2 位置编码模块
位置编码模块是为 Transformer 模块提供 Patch 和 Patch 之间的相对关系，非常关键。在通用任务的 Transformer 模型中认为一个好的位置编码应该要满足以下特性：
- 不同位置的位置编码向量应该是唯一的
- 不能因为不同位置位置编码的值大小导致网络学习有倾向性
- 必须是确定性的
- 最好能够泛化到任意长度的序列输入
ViT 位置编码模块满足前 3 条特性，但是最后一条不满足，当输入图片改变时候需要微调，比较麻烦。基于此也出现了不少的算法改进，结构图如下所示：
按照是否显式的设置位置编码向量，可以分成：
- 显式位置编码，其中可以分成绝对位置编码和相对位置编码。
- 隐式位置编码，即不再直接设置绝对和相对位置编码，而是基于图片语义利用模型自动生成能够区分位置信息的编码向量。
隐式位置编码对于图片长度改变场景更加有效，因为其是自适应图片语义而生成。
3.1.2.2.1 显式位置编码
显式位置编码，可以分成绝对位置编码和相对位置编码。
(1) 绝对位置编码
绝对位置编码表示在 Patch 的每个位置都加上一个不同的编码向量，其又可以分成固定位置编码即无需学习直接基于特定规则生成，常用的是 Attention is all you need 中采用的 sincos 编码，这种编码方式可以支持任意长度序列输入。还有一种是可学习绝对位置编码，即初始化设置为全 0 可学习参数，然后加到序列上一起通过网络训练学习，典型的例如 ViT、PVT 等等。
(2) 相对位置编码
相对位置编码考虑为相邻 Patch 位置编码，其实现一般是设置为可学习，例如 Swin Transformer 中采用的可学习相对位置编码，其做法是在 QK 矩阵计完相似度后，引入一个额外的可学习 bias 矩阵，其公式为：
Swin Transformer 这种做法依然无法解决图片尺寸改变时候对相对位置编码插值带来的性能下降的问题。在 Swin Transformer v2 中作者做了相关实验，在直接使用了在 256 * 256 分辨率大小，8 * 8 windows 大小下训练好的 Swin-Transformer 模型权重，载入到不同尺度的大模型下，在不同数据集上进行了测试，性能如下所示 (Parameterized position bias 这行)：
每个表格中的两列表示没有 fintune 和有 fintune，可以看出如果直接对相对位置编码插值而不进行 fintune，性能下降比较严重。故在 Swin Transformer v2 中引入了对数空间连续相对位置编码 log-spaced continuous position bias，其主要目的是可以更有效地从低分辨权重迁移到高分辨率下游任务。
相比于直接应用可学习的相对位置编码，v2 中先引入了 Continuous relative position bias (CPB)，
B 矩阵来自一个小型的网络，用来预测相对位置，该模块的输入依然是 Patch 间的相对位置，这个小型网络可以是一个 2 层 MLP，然后接中间采用激活函数连接。
其性能如上表的 Linear-Spaced CPB 所示，可以发现相比原先的相对位置编码性能有所提升，但是当模型尺度继续增加，图片尺寸继续扩大后性能依然会下降比较多，原因是预测目标空间是一个线性的空间。当 Windows 尺寸增大的时候，比如当载入的是 8*8 大小 windows 下训练好的模型权重，要在 16*16 大小的 windows 下进行 fine-tune，此时预测相对位置范围就会从 [-7,7] 增大到 [-15,15]，整个预测范围的扩大了不少，这可能会出现网络不适应性。为此作者将预测的相对位置坐标从 linear space 改进到 log space 下，这样扩大范围就缩小了不少, 可以提供更加平滑的预测范围，这会增加稳定性，提升泛化能力，性能表如上的 Log-Spaced CPB 所示。
在 Swin Transformer 中相对位置编码矩阵 shape 和 QK矩阵计算后的矩阵一样大，其计算复杂度是 O(HW)，当图片很大或者再引入 T 时间轴，那么计算量非常大， 故在 Improved MViT ，作者进行了分解设计，分成 H 轴相对位置编码，W 轴相对位置编码，然后相加，从而将复杂度降低为 O(H+W)。
关于绝对位置编码和相对位置编码到底哪个是最好的，目前还没有定论，在不同的论文实验中有不同的结论，暂时来看难分胜负。但是从上面可以分析来看，在 ViT 中不管是绝对位置编码和相对位置编码，当图片大小改变时候都需要对编码向量进行插值，性能都有不同程度的下降 ( Swin Transformer v2 在一定程度上解决了)。
3.1.2.2.2 隐式位置编码
当图片尺寸改变时候，隐式位置编码可以很好地避免显式位置编码需要对对编码向量进行插值的弊端。其核心做法都是基于图片语义自适应生成位置编码。
在论文 How much position information do convolutional neural networks encode? 中已经证明 CNN 不仅可以编码位置信息，而且越深的层所包含的位置信息越多，而位置信息是通过 zero-padding 透露的。既然 Conv 自带位置信息，那么可以利用这个特性来隐式的编码位置向量。大部分算法都直接借鉴了这一结论来增强位置编码，典型代表有 CPVT、PVTv2 和 CSwin Transformer 等。
CPVT 指出基于之前 CNN 分类经验，分类网络通常需要平移不变性，但是绝对位置编码会在一定程度打破这个特性，因为每个位置都会加上一个独一无二的位置编码。看起来似乎相对位置编码可以避免这个问题，其天然就可以适应不同长度输入，但是由于相对位置编码在图像分类任务中无法提供任何绝对位置信息(实际上相对位置编码也需要插值)，而绝对位置信息被证明非常重要。以 DeiT-Tiny 模型为例，作者通过简单的对比实验让用户直观的感受不同位置编码的效果：
2D PRE 是指 2D 相对位置编码，Top-1@224 表示测试时候采用 224 图片大小，这个尺度和训练保持一致，Top-1@384 表示测试时候采用 384 图片大小，由于图片大小不一致，故需要对位置编码进行插值。从上表可以得出：
- 位置编码还是很重要，不使用位置编码性能很差
- 2D 相对位置编码性能比其他两个差，可学习位置编码和 sin-cos 策略效果非常接近，相对来说可学习绝对位置编码效果更好一些(和其他论文结论不一致)
- 在需要对位置编码进行插值时候，性能都有下降
基于上述描述，作者认为在视觉任务中一个好的位置编码应满足如下条件：
- 模型应该具有 permutation-variant 和 translation-equivariance 特性，即对位置敏感但同时具有平移不变性
- 能够自然地处理变长的图片序列
- 能够一定程度上编码绝对位置信息
基于这三个原则，CPVT 引入了一个带有 zero-padding 的卷积 ( kernel size k ≥ 3) 来隐式地编码位置信息，并提出了 Positional Encoding Generator (PEG) 模块，如下所示：
将输入序列 reshape 成图像空间维度，然后通过一个 kernel size 为 k ≥ 3, (k−1)/2 zero paddings 的 2D 卷积操作，最后再 reshape 成 token 序列。这个 PEG 模块因为引入了卷积，在计算位置编码时候会考虑邻近的 token，当图片尺度改变时候，这个特性可以避免性能下降问题。算法的整体结构图如下所示：
基于 CPVT 的做法，PVTv2 将 zero-padding 卷积思想引入到 FFN 模块中。
通过在常规 FFN 模块中引入 zero-padding 的逐深度卷积来引入隐式的位置编码信息(称为 Convolutional Feed-Forward)。
同样的，在 CSWin Transformer 中作者也引入了 3x3 DW 卷积来增强位置信息，结构图如下所示：
APE 是 ViT 中的绝对位置编码，CPE 是 CPVT 中的条件位置编码，其做法是和输入序列 X 相加，而 RPE 是 Swin Transformer 中所采用的相对位置编码，其是加到 QK 矩阵计算后输出中，而本文所提的 Locally-Enhanced Positional Encoding (LePE)，是在自注意力计算完成后额外加上 DW 卷积值，计算量比 RPE 小。
LePE 做法对于下游密集预测任务中图片尺寸变化情况比较友好，性能下降比较少。
除了上述分析的诸多加法隐式位置编码改进， ResT 提出了另一个非常相似的，但是是乘法的改进策略，结构图如下所示：
对 Patch Embedding 后的序列应用先恢复空间结构，然后应用一个 3×3 depth-wise padding 1的卷积来提供位置注意力信息，然后通过 sigmoid 操作变成注意力权重和原始输入相乘。代码如下所示：
class PA(nn.Module):
def __init__(self, dim):
super().__init__()
self.pa_conv = nn.Conv2d(dim, dim, kernel_size=3, padding=1, groups=dim)
self.sigmoid = nn.Sigmoid()
def forward(self, x):
# x 是已经恢复了空间结构的 patch embedding
return x * self.sigmoid(self.pa_conv(x)) 作者还指出这个Pixel Attention (PA) 模块可以替换为任意空间注意力模块，性能优异，明显比位置编码更加灵活好用。
3.1.2.3 自注意力模块
Transformer 的最核心模块是自注意力模块，也就是我们常说的多头注意力模块，如下图所示：
注意力机制的最大优势是没有任何先验偏置，只要输入足够的数据就可以利用全局注意力学到泛化性能不错的特征。当数据量足够大的时候，注意力机制是 Transformer 模型的最大优势，但是一旦数据量不够就会变成逆势，后续很多算法改进方向都是希望能够引入部分先验偏置辅助模块，在减少对数据量的依赖情况下加快收敛，并进一步提升性能。同时注意力机制还有一个比较大的缺点：因为其全局注意力计算，当输入高分辨率图时候计算量非常巨大，这也是目前一大改进方向。
简单总结，可以将目前自注意力模块分成 2 个大方向：
- 仅仅包括全局注意力，例如 ViT、PVT 等
- 引入额外的局部注意力，例如 Swin Transformer
如果整个 Transformer 模型不含局部注意力模块，那么其主要改进方向就是如何减少空间全局注意力的计算量，而引入额外的局部注意力自然可以很好地解决空间全局注意力计算量过大的问题，但是如果仅仅包括局部注意力，则会导致性能下降严重，因为局部注意力没有考虑窗口间的信息交互，因此引入额外的局部注意力的意思是在引入局部注意力基础上，还需要存在窗口间交互模块，这个模块可以是全局注意力模块，也可以是任何可以实现这个功能的模块。其结构图如下所示：
3.1.2.3.1 仅包括全局注意力
标准的多头注意力就是典型的空间全局注意力模块，当输入图片比较大的时候，会导致序列个数非常多，此时注意力计算就会消耗大量计算量和显存。以常规的 COCO 目标检测下游任务为例，输入图片大小一般是 800x1333，此时 Transformer 中的自注意力模块计算量和内存占用会难以承受。其改进方向可以归纳为两类：减少全局注意力计算量以及采用广义线性注意力计算方式。
(1) 减少全局注意力计算量
全局注意力计算量主要是在 QK 矩阵和 Softmax 后和 V 相乘部分，想减少这部分计算量，那自然可以采用如下策略：
- 降低 KV 维度，QK 计算量和 Softmax 后和 V 相乘部分计算量自然会减少 。
- 减低 QKV 维度，主要如果 Q 长度下降了，那么代表序列输出长度改变了，在减少计算量的同时也实现了下采样功能。
(a) 降低 KV 维度
降低 KV 维度做法的典型代码是 PVT，其设计了空间 Reduction 注意力层 (SRA) ，如下所示：
其做法比较简单，核心就是通过 Spatial Reduction 模块缩减 KV 的输入序列长度，KV 是空间图片转化为 Token 后的序列，可以考虑先还原出空间结构，然后通过卷积缩减维度，再次转化为序列结构，最后再算注意力。假设 QKV shape 是完全相同，其详细计算过程如下：
- 在暂时不考虑 batch 的情况下，KV 的 shape 是 (H'W', C)，既然叫做空间维度缩减，那么肯定是作用在空间维度上，故首先利用 reshape 函数恢复空间维度变成 (H', W', C)。
- 然后在这个 shape 下应用 kernel_size 和 stride 为指定缩放率例如 8 的二维卷积，实现空间维度缩减，变成 (H/R, W/R, C), R 是缩放倍数。
- 然后再次反向 reshape 变成 (HW/(R平方), C)，此时第一维(序列长度)就缩减了 R 平方倍数。
- 然后采用标准的多头注意力层进行注意力加权计算，输出维度依然是 (H'W', C)。
而在 Twins 中提出了所谓的 GSA，其实就是 PVT 中的空间缩减模块。
同时基于最新进展，在 PVTV2 算法中甚至可以将 PVTv1 的 Spatial Attention 直接换成无任何学习参数的 Average Pooling 模块，也就是所谓的 Linear SRA，如下所示：
同样参考 PVT 设计，在 P2T 也提出一种改进版本的金字塔 Pool 结构，如下所示：
(b) 即为改进的 Spatial Attention 结构，对 KV 值应用不同大小的 kernel 进行池化操作，最后 cat 拼接到一起，输入到 MHSA 中进行计算，通过控制 pool 的 kernel 就可以改变 KV 的输出序列长度，从而减少计算量，同时金字塔池化结构可以进一步提升性能(不过由于其 FFN 中引入了 DW 卷积，也有一定性能提升)。
从降低 KV 空间维度角度出发，ResT 算法中也提出了一个内存高效的注意力模块 E-MSA，相比 PVT 做法更近一步，不仅仅缩减 KV 空间维度，还同时加强各个 head 之间的信息交互，如下所示：
其出发点有两个：
- 当序列比较长或者维度比较高的时候，全局注意力计算量过大。
- 当多头注意力计算时，各个头是按照 D 维度切分，然后独立计算最后拼接输出，各个头之间没有交互，当 X 维度较少时，性能可能不太行。
基于上述两点，作者引入 DWConv 缩放 KV 的空间维度来减少全局注意力计算量，然后在 QK 点乘后引入 1x1 Conv 模块进行多头信息交互。其详细做法如下：
- 假设输入序列 X Shape 是 nxd，n 表示序列长度，d 表示每个序列的嵌入向量维度。
- 假设想将特征图下采样 sxs 倍，可以将 X 输入到 kernel 为 (s+1,s+1)，stride 为 (s, s), padding 为 (s//2, s//2) 的 DW 卷积和 LN 层中，假设输出变成 (h'w', d)。
- 将其经过线性映射，然后在 d 维度切分成 k 个部分，分别用于 k 个头中。
- QK 计算点积和 Scale 后，Shape 变成 (k, n, n')，然后对该输出采用 1x1 卷积在头的 k 维度进行多个head 之间的信息聚合。
- 后续是标准的注意力计算方式。
其核心代码如下所示：
class Attention(nn.Module):
def __init__(self,
dim=32,
num_heads=8,
qkv_bias=False,
attn_drop=0.,
proj_drop=0.,
sr_ratio=2):
super().__init__()
self.num_heads = num_heads
head_dim = dim // num_heads
self.scale = head_dim ** -0.5
# （sr_ratio+1）x （sr_ratio+1） 的 DW 卷积
self.sr = nn.Conv2d(dim, dim, kernel_size=sr_ratio + 1, stride=sr_ratio, padding=sr_ratio // 2, groups=dim)
self.sr_norm = nn.LayerNorm(dim)
# 1x1 卷积
self.transform_conv = nn.Conv2d(self.num_heads, self.num_heads, kernel_size=1, stride=1)
self.transform_norm = nn.InstanceNorm2d(self.num_heads)
def forward(self, x, H, W):
B, N, C = x.shape
q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
# 1 空间下采样
x_ = x.permute(0, 2, 1).reshape(B, C, H, W)
x_ = self.sr(x_).reshape(B, C, -1).permute(0, 2, 1)
x_ = self.sr_norm(x_)
kv = self.kv(x_).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
k, v = kv[0], kv[1]
# 2 输出维度为 (B,num_head, N, N')
attn = (q @ k.transpose(-2, -1)) * self.scale
# 3 在 num_head 维度进行信息聚合，加强 head 之间的联系
attn = self.transform_conv(attn)
attn = attn.softmax(dim=-1)
attn = self.transform_norm(attn)
# 4 子注意力模块标准操作
attn = self.attn_drop(attn)
x = (attn @ v).transpose(1, 2).reshape(B, N, C)
x = self.proj(x)
x = self.proj_drop(x)
return x
(b) 降低 QKV 维度
Multiscale Vision Transformers (MViT) 也考虑引入 Pool 算子来减少全局注意力计算量。MViT 主要为了视频识别而设计，因为视频任务本身就会消耗太多的显存和内存，如果不进行针对性设计，则难以实际应用。正如其名字所言，其主要是想参考 CNN 中的多尺度特性(浅层空间分辨率大通道数少，深层空间分辨率小通道数多)设计出适合 Transformer 的多尺度 ViT。自注意力模块如下所示：
相比 Transformer 自注意力模块，其主要改变是多了 Pool 模块，该模块的主要作用是通过控制其 Stride 参数来缩放输入的序列个数，而序列个数对应的是图片空间尺度 THW。以图像分类任务为例，
- 任意维度的序列 X 输入，首先和 3 个独立的线性层得到 QKV，维度都是 (HW, D)。
- 将QKV 恢复空间维度，变成 (H, W, D)，然后经过独立的 3 个 Pool 模块，通过控制 stride 参数可以改变输出序列长度，变成 (H', W', D)，设置 3 个 Pool 模块不同的 Stride 值可以实现不同大小的输出。
- 将输入都拉伸为序列格式，然后采用自注意力计算方式输出 (H'W‘, D)。
- 为了保证输出序列长度改变而无法直接应用残差连接，需要在 X 侧同时引入一个 Pool 模块将序列长度和维度变成一致。
由于 MViT 出色的性能，作者将该思想推广到更多的下游任务中（例如目标检测），提出了改进版本的 Imporved MViT，其重新设计的结构图如下所示：
Imporved MViT 在不同的下游任务提升显著。
(2) 广义线性注意力计算方式
基于 NLP 中 Transformer 进展，我们可以考虑将其引入到 ViT 中，典型的包括 Performer，其可以通过分解获得一个线性时间注意力机制，并重新排列矩阵乘法，以对常规注意力机制的结果进行近似，而不需要显示构建大小呈平方增长的注意力矩阵。在 T2T-ViT 算法中则直接使用了高效的 Performer。
在 NLP 领域类似的近似计算方式也有很多，由于本文主要关注 ViT 方面的改进，故这部分不再展开分析
3.1.2.3.2 引入额外局部注意力
引入额外局部注意力的典型代表是 Swin Transformer，但是卷积模块工作方式也可以等价为局部注意力计算方式，所以从目前发展来看，主要可以分成 3 个大类：
- 局部窗口计算模式，例如 Swin Transformer 这种局部窗口内计算。
- 引入卷积局部归纳偏置增强，这种做法通常是引入或多或少的卷积来明确提供局部注意力功能。
- 稀疏注意力。
结构图如下所示：
需要特别注意的是：
