Vision Transformer 必读系列之图像分类综述(二): Attention-based（下）

奇怪现象

作者以 Swin Tiny 版本为例，将其中的局部窗口注意力模块 LSA 替换为 DW Conv、decoupled dynamic filter (DDF)，从上图可以看出 DWConv 和 DDF 性能都比 LSA 强的，特别是 DW Conv，在参数量和 FLOPs 更小的情况下性能会比 Swin Transformer 高。

原因分析

作者试图从 2 个角度来统一分析 LSA、DWConv 和 DDF，分别是通道数设置和空间处理方式 spatial processing。通道数设置 channel setting 可以用于统一建模  LSA、DWConv 在通道数上的差异。DWConv 是逐通道(深度)上计算 Conv，也就是说在不同的通道上应用了不同的滤波器参数，而 DDF 和 LSA 是将通道数分组，然后在同一个组内采用相同的滤波器参数，LSA 的通道数分组实际上就是常说的多头机制即DwConv 可以视作一种特殊的多头策略。如果将 DWConv 的通道数设置为 LSA 的头个数，那么通道数设置原则上就是相同效果了。对比结果如下所示：

从上述曲线可以看出：

• 在相同通道配置下(例如1x、2x，这里是特指4 个 stage 中 head 个数，典型的 1x 是 3, 6, 12, 24)，DwConv 版本仍与 LSA 版本具有相似性能。
  
• 从上述 1x, 2x 结果可以看出，通常来说 DWConv 的通道数肯定比 LSA 的 Head 个数多，这可能是 DWConv 性能比 LSA 高的原因，但是当设置为相同的 Head 个数时，DWConv 性能甚至比 LSA 类型更好一些。，这说明通道配置并非导致前述现象的主要原因。
  
• 当 Head 个数配置大于 1x 时，LSA 性能反而下降，但是 DWConv 性能确实能够提升，这说明两个问题：(1) LSA 中直接提升头数并不能改善通道容量与性能，但是 DWConv 可以；(2) LSA 中需要一种新的策略以进一步提升通道容量和性能。
  

空间处理方式 spatial processing 是指如何得到滤波器或者注意力图并对空域进行信息聚合。DwConv采用静态的滤波器即一旦训练完成，不管输入啥图片都是采用固定 kernel，而其他两者则采用动态滤波器。为了方便统一建模分析，作者采用统一的表示式来说明，如下所示：

Conv 和 DwConv 计算公式为：

当仅仅使用 r^b_j-i 相对位置偏置，Norm 设置为恒等变换， Θ 表示滑动窗口，计算方式则可以表示 Conv 和 DwConv，j-i 表示相对偏移。
Dynamic filters 计算公式为：

W 是 1x1 卷积，当仅仅使用 q_ir^k_j-i ( r^k_j-i 表示 k 的相对位置嵌入向量 ),  Norm 设置为恒等变换则可以表示 Dynamic filters，注意 Dynamic filters 有非常多种做法，以上写的是其中一种。
LSA 计算公式为：

我们可以很容易地将公式 1 退化为 LSA 模式， Ω 表示非重叠局部窗口计算方式。
自从可以将 DwConv 、Dynamic filters 和 LSA 统一起来，并且将其分成三个核心不同部分：参数形式 parameterization, 规范化方式 normalization 和滤波器应用方式 filter application。

参数形式 parameterization

从上表可以看出：

• 动态滤波器的参数策略要比标准 LSA 策略具有更优的性能(Net2 vs Net1)
  
• 动态滤波器变种策略 (Net6) 具有与 SwinT 相当的性能
  
• LSA 参数策略与动态滤波器参数策略的组合 (Net7) 可以进一步提升模型性能
  

规范化方式 Normalization

• 当采用 Net7 的参数形式组合 Identity 时，模型训练崩溃
  
• 相比 FilterNorm，Softmax 规范化具有更优的性能
  
• 规范化方式并非 LSA 平庸的原因
  

滤波器应用方式 filter application

当将滤波器采用滑动窗口形式时，Net6 与 Net7 均得到了显著性能提升。这意味着：近邻处理方式是空域处理的关键。
基于上述实验，使 LSA 变平庸的因素可以分为两个因素：

• 相对位置嵌入是影响性能的一个关键因素。
  
• 另外一个关键因素是滤波器使用方式，即滑动窗口 vs 非重叠窗口。
  

DwConv 能够与 LSA 性能相媲美的原因在于：它采用了滑动窗口处理机制。当其采用非重叠窗口机制时，性能明显弱于 LSA（见 Table1 中的 Net4）。
动态滤波器性能优于 LSA 的原因在于相对位置嵌入与近邻滤波器使用方式。两者的集成 (Net7) 取得了最佳的性能。

对比非重叠局部窗口与滑动窗口，局部重叠的峰值性能要弱于滑动窗口。局部窗口的一个缺点在于：窗口间缺乏信息交互，限制了其性能；而滑动窗口的缺陷在于低吞吐量。那么，如何避免点乘同时保持高性能就成了新的挑战

基于上述分析，作者提出一种增强的局部自注意力模块 ELSA，性能可以超越了 SwinT 中的 LSA 与动态滤波器。

将该模块替换掉 Swin Transformer 中的前三个 stage 中的 LSA 模块即可。ELSA 的关键模块为Hadamard 注意力与 Ghost 头模块，表达式如下所示：

相比 Dot product，Hadamard product ( Pytorch 中的 A*B )可以有效地提取高阶信息，而 ghost 模块是受启发于 GhostNet，可以在有限的容量下提取更丰富的通道特征。作者在分类、目标检测和语义分割任务上都进行了验证。
ELSA 的解决办法看起来过于复杂，整个分析过程也感觉有点复杂，可能会存在更简单的改进策略。

(3) 稀疏注意力（局部紧密相关和远程稀疏）和 Performer 一样，NLP 领域内也有很多对自注意力模块的改进中是引入局部注意力的，典型的例如 Sparse Transformers，其出发点是实际场景中或许用不到真正的全局注意力，只要提供稍微等价的全局效果就行，不一定真的理论上达到全局。基于此可以实现局部紧密相关和远程稀疏的 Sparse Transformers，后续改进版本也非常多，例如 LongFormer。同样的，本文对 NLP 领域发展不展开描述。

3.1.2.4 FFN 模块

FFN 模块比较简单，主要是进行特征变换，ViT 中代码如下所示：

class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, dropout = 0.):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim),
            nn.Dropout(dropout)
        )
    def forward(self, x):
        return self.net(x)

FFN 的改进方向和前面类似，也是希望引入 CNN 局部特征提取能力，加速收敛、提升性能，主要可以归纳为如下所示：

在引入 Conv 增强局部信息特征信息提取方面，也有不少论文进行了尝试。
LocalViT 中在其他模块都不改动情况下，仅仅在 FFN 中引入 1x1 卷积和带 padding 的 3x3 DW 卷积来增强局部特征提取能力，实验结果能够带来不少的性能提升，如下所示：

(a) 是简单的卷积 forward 模块，(b) 是反转残差块，(c) 是本文所提结构，仅仅需要通过替换 FFN 中间的 MLP 层为卷积层。同样采用类似做法的有 PVTv2 和  HRFormer 等等。
CeiT 中也提出了非常类似的结构，其出发点也是希望引入 CNN 来加强局部特征能力，结构如下所示：

主要包括线性投影、恢复空间结构、3x3 DW 卷积、flatten 和线性投影。
在 3.1.2.2 位置编码模块中说过， PVTv2 在 FFN 中引入了零填充的逐深度通道卷积来自动感知位置信息，从而可以省掉位置编码。从 FFN 角度来看，除了有省掉位置编码的好处外， 还能够加强局部特征提取能力，相比 LocalViT 和 CeiT ，直接抛弃位置编码的做法更加彻底，更加合理。

3.1.2.5 Norm 位置改动

Norm 通常是 Layer Norm，按照该模型放在自注意力和 FFN 模块的前面还是后面，可以分成 pre norm 和 post norm 方式，如下所示：

绝大部分模型例如 ViT、Swin Transformer 等都是 pre norm 方式，如下所示：

class Transformer(nn.Module):
    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):
        super().__init__()
        self.layers = nn.ModuleList([])
        for _ in range(depth):
            self.layers.append(nn.ModuleList([
                PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout)),
                PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))
            ]))
    def forward(self, x):
        for attn, ff in self.layers:
            x = attn(x) + x
            x = ff(x) + x

但是 Swin Transformer v2 却改成了 post norm。Swin Transformer V2 主要探讨如何 scale up 视觉模型，并最终在 4 个数据集上重新刷到了新的 SOTA。将视觉 Transformer 的模型参数调大，作者认为会面临以下几个挑战：

• 增大视觉模型时候可能会带来很大的训练不稳定性
  
• 在很多需要高分辨率的下游任务上，还没有很好地探索出对低分辨率下训练好的模型迁移到更大scale 模型上的方法
  
• GPU 内存代价太大
  

作者发现，将 Swin Transformer 模型从 small size 增大到 large size 后，网络深层的激活值会变得很大，与浅层特征的激活值有很大的 gap，如下图所示，可以看出来随着模型 scale 的增大，这个现象会变得很严重。

作者发现使用 post-norm 操作后，上面所观察到的问题可以得到很明显的改善，并且为了更进一步稳定 largest Swin V2 的训练，在每 6 个 transformer block 后还额外加了一层 layer normalization。也就是说在大模型 Transformer 中使用 post Norm 能够更加稳定训练过程。

3.1.2.6 分类预测头模块

在 ViT 中通过追加额外一个 Class Token，将该 Token 对应的编码器输出输入到 MLP 分类头(实际上是一个线性投影层)进行分类，如下所示：

self.mlp_head = nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, num_classes)
        )
# x 是 patch embedding 输出        
cls_tokens = repeat(self.cls_token, '() n d -> b n d', b = b)
x = torch.cat((cls_tokens, x), dim=1)  
...   
# x 是最后一个 Transformer 输出     
self.mlp_head(x[:, 0])     

为何不能和我们常规的图像分类一样，直接聚合所有特征，而需要单独引入一个  Class Token ？这个问题自然有人进一步探索，经过简单总结，可以归纳为如下结构图：

从目前来看主要包括 2 种做法：Class Token 和常规分类一样直接聚合特征。
Class Token 方面的改进典型代码是  CeiT ，作者的出发点是 ViT 仅仅使用最后一个  Transformer 输出做分类，但是实验观察不同的 Transformer 层输出特性不一样，如果能够聚合多个输出层的特征进行分类，也许会更好，具体做法如下所示：

首先收集不同 Transformer 层的 Class Token 对应的输出，然后通过一个 Layer-wise Class-token Attention (LCA) 模块来通过注意力自适应聚合所需要的 token。LCA 可以看出一个简单的 Transformer 模块，其包括多头注意力模块 MSA 和前向网络 FFN，在聚合信息后，取输出序列的最后一个输出进行分类训练。为了节省计算量， MSA 只计算最后一个 Class Token 和其他 Class Token 之间的单向注意力，而没有必要算全局。
随着视觉 Transformer 的快速发展，后人发现其实直接采用 Avg Pool 聚合特征性能也很不错，甚至会更好，典型的代表是 Swin Transformer 和 CPVT-GAP，这种做法无需引入额外的 Class Token，和我们 CNN 中常用的分类套路一致，如下所示：

3.1.2.7 其他

这里说的其他包括两个部分内容，如下所示：

3.1.2.7 其他

这里说的其他包括两个部分内容，如下所示：

3.1.2.7.1 多尺度特征图输出

在 CNN 时代的各种下游任务（例如目标检测、语义分割）中，多分辨率多尺度特征已经被广泛证明非常重要。不同尺度特征可以提供不同的感受野，适合提取不同物体尺度的特征。然而 ViT 仅仅是为图像分类而设计，无法很好地应用于下游任务，这严重制约了视觉 Transformer 的广泛应用，故迫切需要一种能够类似 ResNet 在不同 stage 输出不同尺度的金字塔特征做法。

ViT 要输出多尺度特征图，最常见做法是 Patch Merging。Patch Mergeing 含义是对不同窗口的 Patch 进行合并，在目前主流的 PVT、Twins、Swin Transformer 和 ResT 中都有广泛的应用，以 PVT 为例详细说明，结构图如下所示：

假设图片大小是 (H, W, 3)，暂时不考虑 batch。

1. 考虑将图片切割为 HW/(4X4) 个块，每个块像素大小是 4x4x3, 此处 stride=4。
   
2. 将每个 4x4x3 像素块展开，变成 1 维向量，然后经过线性投影层，输出维度变成 C1，此时特征图 shape 是 (HW/(4X4), C1) 即每个像素块现在变成了长度为 C1 的向量，这两个步骤合并称为 Patch Embedding。
   
3. 将上一步输出序列和位置编码相加，输入到编码器中，输出序列长度不变。
   
4. 将这个输出序列恢复成空间结构，其 shape 是 (H/4, W/4, C1)，此时特征图相比原始图片就下采样了 4x4 倍。
   
5. 在下一个 stage 中改变 stride 数目，然后重复 1-4 步骤就又可以缩减对应 sxs 倍，假设设置 4 个 stage 的 stride 为 [4, 2, 2, 2]，那么 4 个 stage 输出的 stride 就是 [4, 8, 16, 32]，这个就和 ResNet 输出 stride 完全对齐。
   

CSWin Transformer 则更加彻底，直接将上述 stride 下采样过程替换为 kernel 为 3x3，stride为 2 的卷积，做法和 ResNet 中下采样一致。我们将上述改变 stride 导致序列长度变少从而缩减空间特征图的过程统称为 Patch Merging。
除了上述这种相对朴素的做法，还有一些其他做法。例如 MViT ，其不存在专门的 Patch Merging 模块，而是在注意力模块中同时嵌入下采样功能，如下所示：

只要在每个 stage 中改进 Pool 模块的 stride 就可以控制实现 ResNet 一样的多尺度输出，从而实现多分辨率金字塔特征输出。

ViTAE 则直接引入空洞卷积，然后设置不同的空洞率来实现下采样从而输出不同尺度的特征图。

3.1.2.7.2 训练深层 Transformer

探讨如何训练更深的 Transformer 典型算法是 CaiT 和 DeepViT。前面的诸多 ViT 改进都是在编码层为 6 的基础上进行设计的，是否可以类似 CNN 分类网络设计更深的 Transformer，性能也能出现一致性提升，不至于过早饱和？

(1) CaiT

在 CaiT 算法中，作者从 Transformer 架构和优化关系会相互影响相互作用角度出发进行探讨。在 ResNet 中，作者也说过较深的网络能不能顺利的优化才是其成功的关键，而不是在于特征提取能力多强，对应到 ViT 中依然是首先要考虑如何更容易优化，这会涉及到 Normalize 选择、权重初始化、残差块初始化等模块，不能小看这些模块，一个良好的初始化设计都可能避免深模型的过早饱和问题。除了上述影响，作者发现 ViT 的 class token 似乎有点不合理，ViT 在一开始就将 class token 和 patch embedding 输出序列拼接，然后输入给自注意力模块，这个自注意力模块需要同时完成两个任务：

• 引导 attention 过程，帮助得到 attention map。
  
• token 最后输入到 classifier 中，完成分类任务。
  

也就是说，将 class token 和  patch embedding 过早融合，但是两者优化方向不一样，可能存在矛盾。同时说到更顺利的优化，Normalize 的重要性应该是第一个能想到的。基于这两个出发点，作者进行了如下改进：

• 提出 LayerScale，更合理的 Norm 策略使深层 Transformer 易于收敛，并提高精度。
  
• 提出 class-attention layers，class token 和 patch embedding 在最后融合，并且通过 CA 模块来更加高效地将 patch embedding 信息融合到 class embedding 中，从而提升性能。
  

结构图如下所示：

(a) LayerScale

• (a) 是 ViT 做法即先进行 Layer Normalization，再进行 Self-attention 或者 FFN，然后结果与 block 输入相加。
  
• (b) 是 ReZero、Skipinit 和 Fixup 算法的做法，引入了一个可学习的参数 alpha 作用在 residual block 的输出，并移除了 Layer Normalization 操作，这个可学习参数初始化可以从 0 开始，也可以从 1 开始，实验表示都无法很好地解决深层 Transformer 优化问题。
  
• (c) 是一个组合做法，实验发现这种组合做法有一定效果。
  
• (d) 是本文所提 LayerScale，效果最好。
  

LayerScale 的做法是保留 Layer Normalization，并对 Self-attention 或者 FFN 的输出乘上一个对角矩阵，由于其对角矩阵，这个实际上就等效于通道注意力(对不同通道乘上不同的系数)，这些系数的设置比较有讲究，在 18 层之前，它们初始化为 0.1，若网络更深，则在 24 层之前初始化为 10^-5，若网络更深，则在之后更深的网络中初始化为 10^-6，这样设计的原因是希望越深的 block 在一开始的时候更接近恒等映射，在训练的过程中逐渐地学习到模块自身所需要的特征，这个理论和 ResNet 中初始化设计非常类似。
(b) Class-Attention Layers
从上述结构图中可以看出， class token 是在后面融合，前面的 Transformer 层专注于 patch embedding 信息提取。CA 模块做法为：

Q 仅仅来自 class token，而 KV 来自 z=[class token, patch embedding]，上述做法是常规的交叉注意力计算方式。作者实验发现 CA layer 使用 2 层就够了，前面依旧是 N 个正常的 Transformer Block。如果直接看代码可能会更清晰一些：

def forward_features(self, x):
    B = x.shape[0]
    # x 是图片，经过 patch embed 变成序列
    x = self.patch_embed(x)
    # n 个 transformer 模块
    x = x + self.pos_embed
    x = self.pos_drop(x)
    # SA+FFN
    for i , blk in enumerate(self.blocks):
        x = blk(x)
    cls_tokens = self.cls_token.expand(B, -1, -1) 
    # CA+FFN    
    for i , blk in enumerate(self.blocks_token_only):
        cls_tokens = blk(x,cls_tokens)
    x = torch.cat((cls_tokens, x), dim=1)   
    x = self.norm(x)
    return x[:, 0]

上述就是 CaiT 算法主要改进，但是其实要想成功训练出来还需要诸多 trick，作者在论文中也贴了大量对比实验，有兴趣的建议直接阅读原文，本文仅仅描述了最重要的改进而已。更贴心的是，作者还列出了 DeiT-S 到 CaiT-36 的性能提升 trick 表，这个表格同样重要。

(2) DeepViT
CaiT 算法是从整体架构和优化联系角度入手，而 DeepViT 不一样，他通过分析得出深层 Transformer 性能饱和的原因是注意力崩塌，即深层的 Transformer 学到的 attention 非常相似。这意味着随着 ViT 的层次加深，self-attention 模块在生成不同注意力以捕获多样性特征方面变得低效。如下图所示，上面一行是 ViT 训练结果，随着模型深入，自注意力层学到的自注意力权重图趋向均匀且都非常类似，下面一类是作者改进后的可视化结果，根据多样性。

DeepViT 中作者也做了很多前期实验，来验证注意力崩塌现象。那么现在核心问题就是要使得各个注意力层所提取的注意力权重图根据多样性，独特性。
为了解决这个问题，作者提出了 2 种解决办法：
(a) 增加 self-attention 模块的嵌入维度

作者认为，增加嵌入维度，可以增加表征容量从而避免注意力崩塌。实验结果如上所示，随着 Embedding Dimension 的增长，相似的 Block 的数量在下降， Acc 在上升，说明这个策略实际上确实是有效的。但增加 Embedding Dimension 也会显著增加计算成本，带来的性能改进往往会减少，且需要更大的数据量来训练，增加了过拟合的风险，这个也不是啥实质性的解决办法。为此作者提出了第二个改进 Re-attention。
(b)  Re-attention

只需要替换 Self-Attention 为 Re-Attention 即可，其结构图如下右图所示：

作者通过实验观察到同一个层的不同 head 之间的相似度比较小，这表明来自同一自注意力层的不同head 关注到了输入 token 的不同方面，基于此作者提出可以把不同的 head 的输出信息融合，然后利用它们再生出一个新的 attention map。具体为引入一个可学习 Linear Transformation，在注意力图生成后乘上 Linear Transformation，从而生成一个新的  attention map，然后进行 Norm 后乘上 V。

因为我们目的是加强同一个自注意力层间不同 head (通常是8，也就是  Linear Transformation 矩阵大小是 8x8) 之间的交流，所以这个可学习 Linear Transformation 是作用在 head 维度，同时 Norm 是 BatchNorm 而不是传统的 LayerNorm。
CaiT 和 DeepViT 都是关注深层 Transformer 出现的过早饱和问题，不过关注的地方不一样，解决办法也完全不同，但是效果类似，这或许说明还有很大的改进空间。

4. 总结

ViT 的核心在于 Attention，但是整个架构也包括多个组件，每个组件都比较关键，有诸多学者对多个组件进行了改进。我们可以简单将 ViT 结构分成 6 个部分：

1. Token 模块，其中可以分成 Image to Token 模块 和 Token to Token 模块，Image to Token 将图片转化为 Token，通常可以分成非重叠 Patch Embedding 和重叠 Patch Embedding，而  Token to Token 用于各个 Transformer 模块间传递 Token，大部分方案都和 Image to Token 做法一样即 Patch Embedding，后续也有论文提出动态窗口划分方式，本质上是利用了图片级别的语义自动生成最有代表性的采样窗口。
   
   
2. 位置编码模块，其中可以分成显式位置编码和隐式位置编码，显式位置编码表示需要手动设置位置编码，包括绝对位置编码和相对位置编码，而隐式位置编码一般是指的利用网络生成自适应内容的位置编码向量，其提出的主要目的是为了解决显式位置编码中所遇到的当图片尺寸变化时位置编码插值带来的性能下降的问题。
   
   
3. 注意力模块，早期的自注意力模块都是全局注意力，计算量巨大，因此在图片领域会针对性设计减少全局注意力，典型做法是降低 KV 空间维度，但是这种做法没有解决根本问题，因此 Swin Transformer 中提出了局部窗口自注意力层，自注意力计算仅仅在每个窗口内单独计算，不再存在上述问题。
   
   
4. FFN 模块，其改进方向大部分是引入 DW 卷积增强局部特征提取能力，实验也证明了其高效性。
   
   
5. Normalization 模块位置，一般是 pre norm。
   
   
6. 分类预测模块，通常有两种做法，额外引入 Class Token 和采用常规分类做法引入全局池化模块进行信息聚合。
   

随着研究的不断深入，大家发现 Attention 可能不是最重要的，进而提出了 MLP-based 和 ConvMixer-based 类算法，这些算法都是为了说明自注意力模块可以采用 MLP 或者 Conv 层代替，这说明 Transformer 的成功可能来自整个架构设计。MLP-based 和 ConvMixer-based 部分将会在下一篇文章中进行说明。

文章来源：【OpenMMLab】

2022-01-26 18:00