文 @ 000007
号外号外:awesome-vit 上新啦,欢迎大家 Star Star Star ~
https://github.com/open-mmlab/awesome-vitgithub.com/open-mmlab/awesome-vit
在 Vision Transformer 必读系列之图像分类综述(一):概述 一文中对 Vision Transformer 在图像分类中的发展进行了概述性总结。
在 Vision Transformer 必读系列之图像分类综述(二): Attention-based 一文中对 Vision Transformer 的 Attention-based 部分进行详细说明。
本文则对剩余部分进行说明。
ViT 进展汇总思维导图如下图所示:
注意:文中涉及到的思维导图,可以通过 https://github.com/open-mmlab/awesome-vit 下载。
在 Vision Transformer 大行其道碾压万物的同时,也有人在尝试非注意力的 Transformer 架构(如果没有注意力模块,那还能称为 Transformer 吗)。这是一个好的现象,总有人要去开拓新方向。相比 Attention-based 结构,MLP-based 顾名思义就是不需要注意力了,将 Transformer 内部的注意力计算模块简单替换为 MLP 全连接结构,也可以达到同样性能。典型代表是 MLP-Mixer 和后续的 ResMLP。
虽然 CNN 的卷积操作和 Vision Transformer 注意力在各个架构中都足以获得良好的性能,但它们都不是必需的,如果替换为本文设计的 MLP 结构依然可以取得一致性性能。
Mixer Layer 中整体结构和 Transformer 编码器类似,只不过内部不存在自注意力模块,而是使用两个不同类型的 MLP 代替,其分别是 channel-mixing MLPs 和 token-mixing MLPs。channel-mixing MLPs 用于在通道 C 方向特征混合,从上图中的 Channels (每个通道颜色一样) 可以明显看出其做法,而 token-mixing MLPs 用于在不同 patch 块间进行特征混合,其作用于 patch 方向。
在极端情况下,上述两个 Mixer Layer 可以看出使用 1×1 卷积进行通道混合,并使用全感受野的和参数共享的单通道深度卷积进行 patch 混合。反之则不然,因为典型的 CNN 不是 Mixer 的特例。 此外卷积比 MLP 中的普通矩阵乘法更复杂,因为它需要对矩阵乘法和/或专门的实现进行额外的成本降低。代码如下所示,非常简洁。
# 代码是先进行 token mixing 再进行 channel mixing
class MixerBlock(nn.Module):
"""Mixer block layer."""
tokens_mlp_dim: int
channels_mlp_dim: int
@nn.compact
def __call__(self, x):
# (b, patch, c)
y = nn.LayerNorm()(x)
# 交互为 (b, c, patch)
y = jnp.swapaxes(y, 1, 2)
# MlpBlock 作用于 patch 维度,实现 token mixing
y = MlpBlock(self.tokens_mlp_dim, name='token_mixing')(y)
# 交换回来
y = jnp.swapaxes(y, 1, 2)
x = x + y
y = nn.LayerNorm()(x)
# MlpBlock 作用于 C 维度,实现 channel mixing
return x + MlpBlock(self.channels_mlp_dim, name='channel_mixing')(y)
class MlpBlock(nn.Module):
mlp_dim: int
@nn.compact
def __call__(self, x):
y = nn.Dense(self.mlp_dim)(x)
y = nn.gelu(y)
return nn.Dense(x.shape[-1])(y)
结果如下所示,可以看出性能和 ViT 非常接近。
几乎在同时,ResMLP 也沿着这条思路也进行了一些尝试。示意图如下所示:
从上图来看,几乎和 MLP-Mixer 一样,最核心的两个 MLP 层也是分成跨 patch 交互 MLP 层和跨通道 MLP 层。最后输出也是采用 avg pool 进行聚合后分类。
同时作者观察到如下现象:
MLP-Based 类算法相比 ViT 算法,有如下好处:
通过 MLP-Mixer 和 ResMLP 大家逐渐意识到 ViT 成功的关键可能并不是注意力机制,这也间接说明了目前大家对视觉 Transformer 架构理解度还是不够,还有很多研究空间。
众所周知,MLP 一个非常大的弊端是无法自适应图片尺寸,这对下游密集预测任务不友好,MLP-Mixer 和 ResMLP 都存在无法方便用于下游任务的问题,基于这个缺点,CycleMLP 对 MLP 引入周期采样功能,使其具备了自适应图片尺寸的功能,大大提升了 MLP-based 类算法的实用性。其核心做法如下所示:
将 FC 作用于 Channel 通道即 Channel FC 层可以实现自适应图片尺寸功能,因为其信息聚合维度是 C,而这个维度本身是不会随着图片尺寸而改变,将 FC 作用于 Spatial 维度即 Spatial FC 层无法实现自适应图片尺寸功能,HW 维度会随着图片大小而改变。MLP-Mixer 和 ResMLP 为了聚合信息在自注意力层都会包含 Spatial FC 层和 Channel FC 层。
Spatial FC 层的主要作用是进行 patch 或者序列之间的信息交互,比较关键,无法简单的移除。可以从下面对比实验看下:
为了能够在移除 Spatial FC 层但依然保持 patch 或者序列之间的信息交互能力,论文提出一种循环采样 FC 层 Cycle FC,其属于局部窗口计算机制,周期性地在空间维度进行有序采样,和可变形卷积做法非常类似,实际上代码确实是直接采用可变形卷积实现的,由于比较清晰就不再分析计算过程了。作者绘制了非常详细的可视化图:
不再将自注意力层输入看出序列 (N,C) 格式,而是视为图片特征 (H, W, C) 格式,此时 Cycle FC 采样的空间维度可以选择 H 或者 W,如果设置 S_H=3 S_W=1,则如图 c 所示,采样方向是 H 方向,W 方向不进行聚合。而 Cycle MLP 模块(下图的 Spatial Proj 模块)实际上由 1x7 的 Cycle FC 层、 7x1 的 Cycle FC 层和 Channel FC 并联,然后相加构成
代码如下所示:
class CycleMLP(nn.Module):
def __init__(self, dim, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
super().__init__()
self.mlp_c = nn.Linear(dim, dim, bias=qkv_bias)
# W 方向采样
self.sfc_h = CycleFC(dim, dim, (1, 3), 1, 0)
# H 方向采样
self.sfc_w = CycleFC(dim, dim, (3, 1), 1, 0)
# 通道方向
self.reweight = Mlp(dim, dim // 4, dim * 3)
self.proj = nn.Linear(dim, dim)
self.proj_drop = nn.Dropout(proj_drop)
def forward(self, x):
B, H, W, C = x.shape
h = self.sfc_h(x.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)
w = self.sfc_w(x.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)
c = self.mlp_c(x)
a = (h + w + c).permute(0, 3, 1, 2).flatten(2).mean(2)
a = self.reweight(a).reshape(B, C, 3).permute(2, 0, 1).softmax(dim=0).unsqueeze(2).unsqueeze(2)
x = h * a[0] + w * a[1] + c * a[2]
x = self.proj(x)
x = self.proj_drop(x)
return x
作者也对采样过程进行深入分析,对比了另外 2 种采样策略。
空洞 stepsize 采样即每隔 stepsize 个位置再进行采样。
实验表明空洞采样效果比随机采样好,但是不如循环采样。
因为 FC 也可以用 Conv 层代替,作者又和 Conv 1x3 和 3x1 核进行比较,Conv 模式相比 CycleMLP 属于密集采样模式,效果居然比稀疏的 CycleMLP 差,作者分析原因是:密集采样模式会引入额外的参数量,为了快速比较将 epoch 相应的缩短到 100 epoch,在这种设置下,可能密集采样会引入额外的冗余参数,不利于优化。
在构建好 CycleMLP 后,将带 CycleMLP FC 层自注意力模块应用于 PVT 和 Swin 中即可无缝的应用于各种密集下游任务,其详细的结构图如下所示,所构建的网络性能优于目前主流的 Swin Transformer。
ConvMixer 的含义是:
因为 Channel Mixer MLP 层和 1x1 卷积完全等价,所以这里所说的 ConvMixer-based 是强调 Spatial Mixer 层模块可以替换为 DW 卷积(关于全局自注意力和 DW 卷积的区别和联系,在 3.4.3 中会详细说明)。ConvMixer-based 的典型代表是 ConvMixer,其结构图如下所示:
在 MLP Mixer 模型中说过, ViT 的编码器层核心结构可以分成全局空间自注意力层(用于 token 和 token 间信息交互),通道混合 MLP 层(用于每个 token 内的 channle 间信息交互)。ConvMixer 也是沿用了同样的分解策略,只不过替换模块不再是 Spatial Mixer MLP 和 Channel Mixer MLP,而是全部替换为了卷积,分别是 DW Conv 和 1x1 Point Conv。使用 Depthwise Convolution(逐通道卷积) 来提取 token 间的相关信息,类似 MLP Mixer 中的 token-mixing MLP,使用 Pointwise Convolution(1x1 卷积) 来提取 channel 间的相关信息,类似 MLP Mixer 中的 channel-mixing MLP,然后将两种卷积交替执行,混合两个维度的信息。性能如下所示:
本论文核心结论不是强调 DW 卷积的重要性,而是说明 ViT 这种架构的成功不在于是使用了自注意力模块还是 Spatial Mixer MLP,只要有相应的代替结构,性能其实都差不多,我们可能要关注整个 Transformer 架构而不是注意力等模块。这个问题在 3.4.1 中也有详细说明。
前面所提出的 MLP-Mixer 和 ResMLP 已经证明了 ViT 成功的关键可能并不是注意力机制,而是来自其他地方或者说整体架构。基于这个出发点,有大量学者对整个架构进行深入研究,试图从更高维度来理解 Transformer,典型算法如下所示:
基于 ResMLP 的自注意力不是必备的而是可以通过 MLP 代替的结论, MetaFormer 作者提出了更进一步的架构抽象,如下所示:
其核心观点和 ResMLP 一致,即 Transformer 模型中自注意力模块不是最核心的(并不是说可以直接去掉),Transformer 的成功来源其整体架构,同时可以将 Transformer 的 Attention 模块和 ResMLP 的 Spatial MLP 层统称为 Token Mixer,进而提出了 MetaFormer 通用结构,Meta 的含义代表 Token Mixer 是一种统称,只要能够实现 Token Mixer 功能的模型都属于 MetaFormer 范畴,例如你也可以将 Token Mixer 换成 3x3 DW 卷积。为了验证这个架构的可行性,作者将 Token Mixer 替换为最简单的无参数的 Pooling 算子,其性能如右边图所示,令人意外的是 PoolFormer 性能居然比 DeiT 和 ResMLP 更好,详细结构如下所示。
由于架构非常简单,不再进行详细描述。以下是基于 PoolFormer 的一些 Ablation。
Pooling 模块代码如下:
注意有个减号,原因是 Pooling 模块后面还有一个残差连接线,Pooling 模块中先减掉 x 才能构成可能的恒等变换 x+ (Pool(x)-x) = x,而通常的 Conv 操作是可以通过调整初始化核参数来实现 x+ Conv(x)=x 的。
An Empirical Study of CNN, Transformer, and MLP 对现在火热的 CNN、Transformer 和 MLP 之争进行了一些经验性研究,可以带给我们一些启发。
为了能够公平比较,作者首先设计了一个通用的 SPACH 架构,在该框架下引入了 mixing block 概念,其结构如下所示,上面是单阶段 SPACH 表示 ViT 这种直筒结构,多阶段 SPACH 表示 ResNet 这种下采样结构,SPACH-MS 结构可以用于高分辨率图片输入场合,该结构和大部分网络例如 PVT 一样。
内部的 Spatial Mixing 可以即插即用的替换为任意 CNN、Transformer 和 MLP ,如下所示
作者对上述三种 Spatial Mixing 模块和单阶段、多阶段 SPACH 模块进行实验,得出了如下经验性结论:
可以说上述 4 条结论已经在各个算法中得到了充分验证,是非常可信的,这也能够启发后续相关网络设计。
Demystifying Local Vision Transformer: Sparse Connectivity, Weight Sharing, and Dynamic Weight 深入分析了注意力和 Conv 的关系,并进行了大量的实验验证,是一篇值得学习的论文。
Conv 是众所周知的局部提取算子,而 Swin Transformer 中虽然没有借助 Conv ,但是提出了类似功能的局部注意力模块 W-MSA,本文试图深入分析局部注意力机制和 Conv (尤其是 depth-wise convolution )的区别和联系,其关系如下所示:
这里需要明确 ViT 中我们所说的全局自注意力实际上是全局空间自注意力,是逐通道算的,通道间没有信息混合,通道信息混合是通过后续的 FFN 模块实现。
通过上图可以清晰的发现局部注意力计算可以采用 3x3 DW 卷积等效替换,并且可以总结出局部注意力的三条特性:
DW Conv 具备前两条特性,但是其不是动态权重的,其权重是一旦训练好就固定了,不会随着图片内容不同而改变。下面是总结汇总表格:
作者还进一步描述了 MLP、Conv 等层的联系,如下所示:
从 2d 全连接层 MLP 出发
从 2d 全连接层 MLP 出发
如果同时具备稀疏连接,仅仅在空间维度信息聚合则为逐深度方向可分离 DW-S Conv。
基于上述等价分析,作者基于 Swin Transformer 提出了 DWNet。
对 Swin Transformer 的改动包括:
其性能如下所示:
可以看出,如此简单的模型性能和 Swin Transformer 差不多,这也和 ConvMixer 相互印证了。
如果说 MetaFormer 还有 Transformer 的影子,那么 ConvNeXt 就是一个更彻底的去 Transformer 的例子了。其核心出发点是纯粹的 Conv 堆叠性能能不能超过 Transformer? ConvNeXt 对 Swin Transformer 进行了逐模块分解,并且将其应用于 ResNet 上,通过不断对比两者差异,作者将 ResNet 改造为 ConvNeXt,性能最终超越 Swin Transformer,证明了纯粹的 Conv 堆叠性能能够超过 Transformer,这也间接说明 Transformer 架构和优化策略的优异性,而不是所谓的 Attention。
本文对 Vision Transformer 发展进行了详细的综述,由于内容比较多,下面进行简要总结。
在 Attention is all you need 论文中第一次提出 Transformer 架构,ViT 将其第一次成功推广到了视觉任务领域,DeiT 针对 ViT 需要超大数据预训练问题提出了蒸馏学习训练方式,在解决大数据预训练问题的同时也进一步提高了性能,推动了 ViT 在视觉中的进一步应用的可能。
ViT 的核心在于 Attention,但是整个架构也包括多个组件,每个组件都比较关键,有诸多学者对多个组件进行了改进。我们可以简单将 ViT 结构分成 6 个部分:
在此之后,我们还分析了如何输出多尺度特征图,方便用于下游密集任务以及如何训练更深的 Transformer。
随着研究的深入,大家发现 attention 可能不是最重要的,进而提出了 MLP-based 和 ConvMixer-based 类算法,这些算法都是为了说明自注意力模块可以采用 MLP 或者 Conv 层代替,这说明 Transformer 的成功可能来自整个架构设计。
进一步,MetaFormer、An Empirical Study of CNN, Transformer, and MLP 和 Demystifying Local Vision Transformer 等论文都详细说明和验证了上面的说法,并都提出了自己各自的看法,这或许也指明了 ViT 的下一步研究方向。
从目前 Vision Transformer 进展来看,目前 CNN 和 attention 的边界已经越来越模糊了,相互可以等价替换,也可以相互增强,谁强谁弱尚未可知,或许 Conv、MLP 和 attention 混合才是最终能够实际落地的出路吧!
至此图像分类综述就完成了,主要包括三篇文章:
对 Vision Transformer 系列内容感兴趣的朋友,不要忘记 star 啦~
https://github.com/open-mmlab/awesome-vigithub.com/open-mmlab/awesome-vit
后续我们会对训练 trick 等进行整理,会同步到上述 repository 中,敬请期待!