【Transformer】13、Augmented Shortcuts for Vision Transformers

论文链接： https://arxiv.org/abs/2106.15941

代码链接：未开源

一、背景

Transformer 已经在计算机视觉领域的某些任务上取得了超越 CNN 的效果。在 Transformer 的结构中，self-attention 和 MLP 模块前后通常会使用残差连接。残差连接最初是在 CNN 的结构中提出，用来缓解梯度消失问题。

二、动机

有些文章指出，没有残差连接的结构的网络效果很差，如表 1 所示，移除残差连接后，随着网络的加深，会使得不同 patch 的特征丧失区别性，如图3a所示。

这些特征无法支撑后续的预测工作。作者把这个现象叫做 feature collapse（特征崩塌）。在加上残差连接之后，能缓解上述现象，如图3 b 所示。

但是，一般的残差连接仅仅是把输入加到输出上，这会限制特征的不同性。于是，本文提出了一种增强残差连接，来提升 vision transformer 中的特征差异性。

三、方法

本文提出的增强残差连接结构如图 1 所示，除了基本的残差连接以外，作者提出了一个与 MSA 并行的模块，该模块有多个参数。为了降低计算成本，作者提出了使用块循环映射的方法来实现增强残差连接。

3.1 增强残差连接

使用 T 个增强残差连接的 MSA 模块可以表示如下：

• $T_{li}(.)$ 是第 $l$ 个 layer 的第 $i$ 个增强残差连接
• ${\Theta }_{li}$ 是增强残差连接的参数

不同于传统的残差连接，增强残差连接能够将输入映射到另外一个特征空间，只要其参数是不同的，那么就可以把输入映射为多个不同的特征，丰富特征空间。

$T_{li}(.)$ 的一个简单的形式就是一系列的线性映射和激活函数的堆叠：

• ${\Theta }_{li}\in R^{d\times d}$ 是权重矩阵
• $\sigma$ 是非线性激活函数（如GELU）
• 上式中，$T_{li}(.)$ 独立的处理每个 patch，能够保留其各自的特征，是对 MSA 的一个补充（MSA 会把不同 patch 的特征进行聚合）

使用 T 个增强残差连接的 MLP 如下：

• 使用多个增强残差连接同样能够增大特征多样性

作者把同时使用 AugMSA 和 AugMLP 的模块叫做 Aug-ViT 模型，该模型的特征有更大的多样性，如图 3c 和图 4 所示。

3.2 使用循环映射进行高效实现

直接使用如公式 8 所示的方法来实现 $T_{li}(.)$ 会涉及到很多矩阵相乘，会很耗资源，所以作者提出了一种 “block-circulant matrices”。

Circulant matrix 是频域中的一种特殊的结构化矩阵，参数和计算量都很少。

一个 Circulant matrix $C\in R^{d^{\prime }\times d^{\prime }}$ 仅有 $d^{\prime }$ 个参数，两个 $C$ 相乘的计算复杂度为 $O(d^{\prime }log{d}^{\prime })$ (使用快速傅里叶变换 FFT）。

原始的参数矩阵 $\Theta$ 被切分成 $b^2$ 个子矩阵 $C^{ij}\in R^{d^{\prime }\times d^{\prime }}$

• $d^{\prime }$ 是子矩阵的尺寸，$b{d}^{\prime }=d$
• 每个 $C^{ij}$ 都是通过循环 $d^{\prime }$ 维向量 $c^{ij}=[c_1^{ij},c_i{2}^{ij},...,c_{d^{\prime }}^{ij}]$ 中的元素得到的：

为了实现投影映射 $T(Z)=\sigma (Z\Theta )$，作者首先将输入 $Z$ 拆分为 $b$ 个切片，然后对每个切片乘以循环矩阵 $C^{ij}$，元定义域中的循环矩阵和向量直接的乘积，等价于傅里叶域中的元素相乘，输出的计算方式为：

四、效果

图 5 展示了不同模型的特征图，没有残差连接的深层的特征图已经没有有用的信息了（a），加上残差连接是（b）的效果，加上 Aug-ViT 是（c）的效果，可以看出（c）的特征最丰富。

总之，残差连接是 vision transformer 的一个非常重要的部分，能够有效避免特征崩溃，而增强残差连接能够将输入特征映射到更丰富的特征空间，让特征更加丰富。