【Transformer】Augmented Shortcuts for Vision Transformers

文章目录

一、背景
二、动机
三、方法
- 3.1 增强残差连接
- 3.2 使用循环映射进行高效实现
四、效果

论文链接： 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 模块可以表示如下：

Tli(.)T_{li}(.)Tli(.) 是第 lll 个 layer 的第 iii 个增强残差连接
Θli\Theta_{li}Θli 是增强残差连接的参数

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

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

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

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

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

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

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

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

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

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

原始的参数矩阵 Θ\ThetaΘ 被切分成 b2b^2b2 个子矩阵 Cij∈Rd′×d′C^{ij} \in R^{d' \times d'}Cij∈Rd′×d′

d′d'd′ 是子矩阵的尺寸，bd′=dbd'=dbd′=d
每个 CijC^{ij}Cij 都是通过循环 d′d'd′ 维向量 cij=[c1ij,ci2ij,...,cd′ij]c^{ij}=[c_1^{ij}, c_i2^{ij},..., c_{d'}^{ij}]cij=[c1ij,ci2ij,...,cd′ij] 中的元素得到的：

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

四、效果

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

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