CMT：当CNN遇到Transformer，“迷途知返”

AIWalker

发布于 2021-07-29 11:56:22

1.1K0

发布于 2021-07-29 11:56:22

文章被收录于专栏：AIWalkerAIWalker

标题&作者团队

本文是华为诺亚与悉尼大学在Transformer+CNN架构混合方面的尝试，提出了一种同时具有Transformer长距离建模与CNN局部特征提取能力的CMT。相比之前的各种Transformer变种，本文更倾向于将Transformer的优势集成到CNN中。整体架构采用了ResNet的分阶段架构，Normalization方面采用CNN中常用的BN而非Transformer中的LN，在核心模块CMTBlock方面，内部设计了具有局部特征提取的LPU，在降低计算量方面对K与V进行了特征分辨率的下降，与此同时将MobileNetV2中的逆残差思想引入到FFN中得到了IRFFN。总而言之，CMT代表着CV中的Transformer架构趋势又回到了CNN原先研究路线：即CNN为主，其他领域思想为辅。

Abstract

由于其所具有的长距离依赖建模能力，Vision Transformers已被成功应用到图像识别任务中。然而，其性能与计算量距离优秀的CNN仍存在差距。

为解决上述问题，我们设计了一种新的网络CMT，它不仅由于Transformer，同时优于高性能CNN。所提CMT是一种混合CNN与Transformer的架构，它同时利用率Transformer的长距离建模与CNN的局部特征提取能力。具体来说，所提CMT-S取得了83.5%的top1精度，同时比现有的DeiT快14倍，比EfficientNet快2倍。所提CMT-S同样具有非常好的泛化性能，比如CIFAR10取得了99.2%，CIFAR100上取得了91.7%，Flowers上取得了98.7%，COCO上取得了44.3%mAP。

Method

本文的初衷是构建一种混合网络，它同时利用CNN与Transformer的优势。下图给出了ResNet50、DeiT以及所提CMT的网络架构示意图。

如上图所示，DeiT直接将输入图像拆分为非重叠图像块，图像块的结构信息则通过线性投影方式弱建模。为克服该局限性，我们采用类似ResNet的stem架构，它由三个

3\times 3

卷积构成，但激活函数采用了GELU，而非ResNet的ReLU。

类似经典CNN(如ResNet)架构设计，所提CMT包含四个阶段以生成多尺度特征(这对于稠密预测任务非常重要)。为生成分层表达，在每个阶段开始之前采用

2\times 2

卷积降低特征分辨率并提升通道维度。在每个阶段，堆叠多个CMT模块进行特征变换同时保持特征分辨率不变，每个CMT模块可以同时捕获局部与长距离依赖关系。在模型的尾部，我们采用GAP+FC方式进行分类。

给定输入图像，我们可以得到四个不同分辨率的分层特征，类似于经典的CNN(ResNet, EfficientNet)。所得四个不同分辨率的特征对应的stride分别为4、8、16、32，因此，CMT所得多尺度特征表达可以轻易应用到下游任务(比如检测与分割)中。

CMT Block

所提CMT模块包含一个局部感知单元(Loal Perception Unit, LPU)、一个轻量型多头自注意力模块(Lightweight Multi-Head Self-Attention, LMHSA)以及一个逆残差前馈网络(Inverted Residual Feed-Forward Network, IRFFN)。

Local Perception Unit 旋转与平移是视觉任务中两种常见数据增广方法，这些操作应当不能影响模型最终的结果。然而，Transformer中的绝对位置编码会破坏该不变性。此外，Transformer忽略了块内的局部相关性与结构信息。为缓解该限制，我们提出了局部感知单元以提升局部信息，定义如下：

LPU(X) = DWConv(X) + X

Lightweight Multi-Head Self-Attention 在原始的自注意力模块中，输入X线性变换为Q、K以及V，然后通过如下方式执行自注意力操作：

Attention(Q,K,V) = Softmax(\frac{QK^T}{\sqrt{d_k}})V

为减少计算复杂度，我们采用

k \times k

深度卷积降低K与V的空间尺寸，即

K = DWConv(K), V = DWConv(V)

。此外，类似Swin，我们在自注意力模块中添加了相对位置偏置B：

\text{LightweightAttention}(Q,K,V) = Softmax(\frac{QK^T}{\sqrt{d_k}} + B)V

Inverted Residual Feed-forward Network 原始的FFN仅包含两个全连接层+GeLU，第一个FC用于扩展特征维度，第二个FC用于降低特征维度：

FFN(X) = GELU(XW_1 + b_1)W_2 + b_2

本文所提IRFFN采用了类似MobileNetV2的逆残差模块，包含一个扩展层、深度卷积以及投影层。具体来说，我们改变了短连接的位置以获得更好的性能：

IRFFN(X) = Conv(\mathcal{F}(Conv(X))) \\ \mathcal{F}(X) = DWConv(X) + X

注：

\mathcal{F}(X)

这里在推理阶段可以进行合并，即跳过连接可以移除掉。它的作用则是利用微小的代价提取局部解雇信息。

基于上述所提到的三个成分，CMT模块定义如下：

X = LPU(X) \\ X = LMHSA(LN(X)) + X \\ X = IRFFN(LN(X)) + X

Complexity Analysis

接下来，我们对ViT与CMT的计算量复杂度进行分析。标准的Transformer模块包含MHSA与FFN，假设输入特征尺寸为

n \times d

，那么整体计算复杂度为：

\mathcal{O}(MHSA) = 2nd(d_k + d_v) + n^2(d_k + d_v) \\ \mathcal{O}(FFN) = 2nd^2r

因此，当

d=d_k=d_v,r=4

时，ViT中的Transformer计算复杂度表示如下：

\mathcal{O}(ViTTransformer) = \mathcal(O)(MHSA) + \mathcal{O}(FFN) = 12nd^2 + 2n^2d

采用类似ViT的配置，CMT的计算量复杂度表示如下：

\mathcal{O}(LPU) = 9nd \\ \mathcal{O}(LMHSA) = 2nd^2(1+1/k^2) + 2n^2d/k^2 \\ \mathcal{O}(IRFFN) = 8nd^2 + 36nd \\ \mathcal{O}(CMT) = 10nd^2(1+0.2/k^2) + 2n^2d/k^2 + 45nd

相比ViT，CMT对计算复杂度更友好，可以处理更高分辨率的特征。

Scaling Strategy

受启发于EfficientNet，我们提出了一种适用于Transformer的复合缩放策略：它采用复合系数

\phi

均匀的调整深度、维度以及输入分辨率：

\text{depth}: \alpha^{\phi}, \text{dimension}: \beta^{\phi}, \text{resolution}: \gamma^{\phi} \\ s.t. \alpha \cdot \beta^{1.5} \cdot \gamma^2 \approx 2.5, \alpha \ge 1, \beta \ge 1, \gamma \ge 1

注：在实验中，我们设置

\alpha=1.2, \beta=1.3, \gamma=1.15

。下表给出了不同复杂度的CMT模型配置信息。