深度学习之SuperViT

这篇文章是今年发表的一篇Transformer类文章。文章作者提出了ViT的增强版——SuperViT，其核心有2个：①多尺度输入；②Token-keeping-rate机制。作者通过一系列实验证明SuperViT可以在计算效率和正确率上实现比Swin-T更好的表现力！

参考文档：
①源码
②Transformer学习(四)—DeiT
③各类Transformer都得稍逊一筹，LV-ViT：探索多个用于提升ViT性能的高效Trick

Abstract

首先需要指明的一点是，SuperViT比ViT强在哪？

1. 计算效率。SuperViT通过使用较低分辨率的输入图像来将token控制在较小的范围内，从而降低相似度计算次数；此外，作者将一些含有重要信息较少的token直接丢弃，从而进一步降低每一种尺度下的注意力计算量。
2. 准确率。SuperViT使用Token-Keeping-rate机制，也可以理解为丢弃不相关的token的方法。这种方式可以将注意力集中在信息更加重要的地方；此外，多尺度的输入产生了不同尺度的相关性与感受野，从而可以提取更加丰富的特征信息。

---

本文提出了一种新的优化ViT的模型——SuperViT，其具有以下特点：

1. Vision Transformer具有较弱的归纳偏置、计算量和token数量成正比等缺陷。近年来，为了解决这些问题，陆续出现了DeiT、LV-ViT、T2T-ViT、PVT、Swin-T等用于改进ViT模型的方法。SuperViT也是出一这一目的，简单来说，SuperViT通过使用多尺度的输入形式以及Token-Keeping-rate机制来实现计算效率的提升与目标识别准确度的增加。
2. 具体而言，多尺度输入(Multi-size)可以产生多种感受野的token，其中高分辨率图像具有较小的感受野但是其具有丰富的局部信息；低分辨率图像则具有较大的感受野，但是天然会损失一些细节信息。这种多种分辨率结合的方式可以提取到较大和较小范围内的特征信息。此外较小分辨率的输入在产生token的时候，必定会输出较少数量的token，这样一来便可以减少相似度计算量。Token-Keeping-rate机制根据注意力的大小来丢弃一些不重要的token，从而为下一次迭代中Transformer的计算量带来计算效率和准确性的提升；对于同一个输入尺度，作者设置了多个keeping rate来选择不同粗细的重要token。
3. SuperViT在ImageNet上证明了其具有较大幅度降低计算资源以及提升识别准确率的特点。

1 Introduction

1. Vision Transformer的兴起始于2020年的ViT一文，ViT最大的缺陷在于其计算量偏大，对于Transformer来说，token越多计算复杂度越高，故ViT很难推广到实时场景下使用。
2. 对于一幅图像来说，token序列长度和patch-size是成反比的，因此一个思路就是想办法增大patch-size，但是众所周知Transformer在较小尺寸的patch上可以实现较好的表现力，这是因为其可以提取更丰富的局部信息并且更加接近卷积的归纳偏置，这也可以减少训练的难度。现存的一些Vision-Transformer方法都只采用固定的patch-size，而使用多种尺寸的patch-size还未有人尝试过，显然多尺寸(multi-size)带来的最大好处在于通过将大尺寸和小尺寸结合，可以利用大patch实现较低计算量和较小patch实现性能的提升。
3. 图像中存在着一些冗余区域或者说含有较多无用信息的区域——背景等。近年来有许多模型就争对这一问题选择去丢弃这些区域，虽然提升了计算效率，但他们都带来了性能的降低。此外大多都采用固定的丢弃率从让模型变得不那么灵活，且固定的丢弃率并不能较好的实现表现力和计算效率的trade-off。

---

基于上述的现状，作者提出SuperViT模型，这里的“super”表示使用多种尺度patch-size的同时采用多种keeping-rate来较好的实现计算效率和性能的平衡。在具体实施过程中：

1. SuperViT输入$m$种不同尺寸的图像(通过上、下采样实现)，然后使用固定的patch-size(源码是$16\times 16$)去产生token，这样就相当于是不同的patch-size去分割同一张图像了。
2. 而每一个尺度都配有$n$种keep-rate，这样就一共有$m\times n$种组合，每一种组合都去做inference，最后将所有梯度结合并更新网络，具体如何更新见3.3节的Training Objective。此外，每一种组合都会进行$N$次的MSA计算，源码中采用间断式地丢弃token，比如$N=12$，则第3、6、9次保持$keeping\_rate=1$，其余情况都进行丢弃。从最后结果来看，这种间断式的丢弃证明是可行的。
3. SuperViT是首次在前向过程中拥有多种复杂度的尝试。

Note：

1. SuperViT获取token的方式是使用patch-size为16的卷积核，步长stride为16的Conv2d来做的。

2 Related Work

略

3 Methodology

3.1 Overview

首先简略回顾下ViT，具体可参考我的另一篇深度学习之ViT。
①设输入为$\mathcal{I}\in {\mathbb{R}}^{H\times W\times C}$，使用$P\times P$的size去分割图像，则获取$N=\frac{HW}{P^2}$个块，每个token的shape为$P\times P\times C$；每一个token经过投影层输出$D$维的向量——$x_i^0\in {\mathbb{R}}^D$，然后这些token进行整合就是注意力模块的输入了，我们一般记为${\mathbb{R}}^{B\times N\times P^{2}C}$($B$为batchsize)：
$$X^0=[x_{cls}^0;x_1^0;\cdots ;x_N^0]+E_{pos}.\text{\hspace{1em}(1)}$$其中$x_{cls}$为类别token，$E_{pos}$为位置编码向量。
②接下来要将输入送进Attention模块，在Vision-Transformer里，一般只用Encoder模块，其包括MHSA和FFN，2者都需要使用LN层和skip-connection。具体表达如下：
$$Y^l=X^{l-1}+MHSA(X^{l-1}).\text{\hspace{1em}(2)}$$$$X^l=Y^l+FFN(Y^l).\text{\hspace{1em}(3)}$$
③由于是多头注意力，所以之后需要进行concat：
$$MHSA(X^{l-1})=Concat[Attention(Q^{l,h}{K}^{l,h})V^{l,h}{]}_{h=1}^H{W}^l.\text{\hspace{1em}(4)}$$在代码中，我们一般会增加一个head的维度，使得$Q、K、V$的shape变成$(B,H,N,D/H)$。
④Attention的计算过程如下，为了不让相似度值落入softmax的两边区域，则除以$\sqrt{D}$来避免相似度值太大。理论上余弦相似度需要$Q、K$进行归一化，但是可能是为了简化计算此外当时Transformer对不进行归一化产生的效果还不错就流传下来都不进行归一化了：
$$Attention(Q^{l,h},K^{l,h})=[a_{cls}^{l,h};a_1^{l,h};\cdots ;a_N^{l,h}]=Softmax(\frac{Q^{l,h}(K^{l,h}{)}^T}{\sqrt{D}}).\text{\hspace{1em}(5)}$$
⑤由于SuperViT是基于目标识别的，因此最后第$L$个Transformer会输出$x_{cls}^L$，并通过FC和softmax进行分类，最后输出类别分布概率$p$：
$$p=\mathcal{T}(X^0)=Softmax(FC(x_{cls}^L)).\text{\hspace{1em}(6)}$$

3.2 Super Vision Transformer

上图就是SuperViT的pipeline，下面提几个点：

1. 由于使用ImageNet，所以输入为$224\times 224\times 3$的格式。
2. 接下去会使用插值(F.interpolate)进行下采样从而获取multi-size的输入格式，不同输入格式通过并行的结构送进网络。
3. 实际设计中，不同输入尺寸的patch-size是相同大小的，变动的是输入的图像，这样产生的效果就相当于不同尺寸的patch。
4. 设$m$种尺寸，$n$种Keeping rate，则一共$m\times n$种情况，每种情况都进行inference，等所有情况都完成后，再直接更新网络参数，这是利用了PyTorch的梯度累加机制。
5. Keeping rate为1下产生的概率分布是和标签进行交叉熵计算的，其余小于1的keeping rate则和同一尺寸下的Keeping rate为1的结果进行KL散度计算，旨在控制好和Keeping rate为1的差距，从而不至于损失太多信息。
6. SuperViT有2个版本，一个是基于DeiT的，另一个是基于LV-ViT的，我自己在复现的时候就只使用了前者。
7. Token-Keeping-rate是使用上一次迭代的注意力值或置信度进行排序，输出前$N^{\prime }<N$个token——$shape=(B,N^{\prime },C)$，拿着这个结果继续进行下一轮迭代。

3.2.1 Multi-size Patch Splitting

一般来说，大尺寸的patch-size可以减少计算量，而小尺寸的patch-size可以有较好的表现力(因为其归纳偏置和卷积更接近)；此外固定patch-size的模型无法有效调节性能和计算效率之间的trade-off。因此SuperViT决定同时使用多种尺度的patch-size来有效调节准确度和计算效率的平衡。

---

设一共并行输入$G$组具有不同尺寸patch-size的图像，patch组记为 { P g × P g } g = 1 G ， P g > P g + 1 \{P_g\times P_g\}_{g=1}^G，P_g > P_{g+1} {Pg​×Pg​}g=1G​，Pg​>Pg+1​，即第$g$组的patch大小为$P_g\times P_g$。因此对于$\mathcal{I}\in {\mathbb{R}}^{H\times W\times C}$的输入，第$g$组获取token的shape为$(B,N,P_g^{2}C)$，其中$N=\frac{HW}{P_g\times P_g}$。
这样每一组获取的token大小都是不一样的，接下去就要统一不同组之间的token的shape。这是为了更好的设计和适应一个网络模型。作者指出使用计算消耗较低的双线性插值来做上下采样，从而统一不同尺寸的patch。之后将每个组内的所有patch($1\to N_g$)进行整合，数学表达式如下：
$$X_g^0=[x_{g,cls}^0;x_{g,1}^0;\cdots ;x_{g,N_g}^0]+E_{g,pos}.\text{\hspace{1em}(7)}$$接下去将$G$组的并行输入 { X g 0 } g = 1 G \{X_g^0\}_{g=1}^{G} {Xg0​}g=1G​共同输入一系列串联的ViT中，即MHSA-FFN。最后接全连接层和Softmax层输出$G$个概率分布—— { p g } g = 1 G \{p_g\}_{g=1}^G {pg​}g=1G​。数学表达式如下：
$$p_g=\mathcal{T}(X_g^0)=Softmax(FC(x_{g,cls}^L)).\text{\hspace{1em}(8)}$$

3.2.2 Multi-Token Keeping Rate

为了进一步节约计算量以及舍弃图片中信息冗余或无用的区域，作者提出了多种keeping-rate下的token丢弃机制。之前也有过类似的模型，但是它们大多使用固定keeping-rate的方式，这样带来的缺陷就是由于不同图片冗余程度是不一样的，因此固定的丢弃反倒会损害表现力。

---

假设每组下有$M$种keeping-rate—— { η m } m = 1 M , η m > η m + 1 \{\eta_m\}_{m=1}^M,\eta_m > \eta_{m+1} {ηm​}m=1M​,ηm​>ηm+1​，且满足$0<{\eta }_m<1,form>1$，这就意味着$top-({\eta }_m\cdot N)$个token将会被保留下来，这个排序是按照注意力值或者置信度来排列的。

因此第$l$个ViT的输入可表示为：
$$X_g^l=[x_{g,cls}^l;x_{g,1}^l;\cdots ;x_{g,N_g}^l]$$接下去常规操作，将$X_g^l$输入MHSA-FFN中，其中$l\in [1,\cdots ,L]$。

Note：

1. 就目前SuperViT的源码版本来看，其基于注意力的drop机制是有问题的——每个token保留下来的token都是一样的，都是基于第一个token丢弃的结果，这显然是不对的，每个token应该丢弃的情况肯定是不一样的。

最后输入到全连接层和Softmax层获取概率分布就完事了：
$$p_{g,m}=\mathcal{T}(X_g^0|{\eta }_m)=Softmax(FC(x_{g,cls}^L|{\eta }_m)).\text{\hspace{1em}(9)}$$

3.3 Training Objective

理论上目前已有$M\times G$种复杂度，但是如果全上的话会带来较大的训练压力，故作者只采用$2\times 2$种复杂度，分别是：$(P_G\times P_G,{\eta }_1)、(P_G\times P_G,{\eta }_{m_1})、(P_{g_1}\times P_{g_1},{\eta }_1)、(P_{g_1}\times P_{g_1},{\eta }_{m_2})$。
Note：

1. ${\eta }_{m_1}、{\eta }_{m_2}$是从 { η m } m = 2 M \{\eta_m\}_{m=2}^M {ηm​}m=2M​中随机采样而来的。
2. ${\eta }_1=1,P_G=224$。

---

SuperViT的整体训练损失函数如下：
$$\mathcal{L}=\sum _{g=1}^{G}CE(p_{g,1},y)+\sum _{g=1}^G\sum _{m=2}^{M}KL(p_{g,m},p_{g,1}).\text{\hspace{1em}(10)}$$Note：

1. 只有$p_{g,1}$和标签直接做CE交叉熵计算，这是因为$p_{g,1}$是无丢弃的，它可以最大程度拟合标签；而含有token丢弃的$p_{g,m}$则和$p_{g,1}$做KL散度来最小化两者的差距，这么做是为了让$p_{g,1}$去引导含有token丢弃的学习。

4 Experiments

实验部分就不详细介绍了，SuperViT的数据集是ImageNet。我自己在复现的时候，使用miniImagenet，也是$224\times 224$，只不过不是1000分类，而是100类。

下面贴一个SuperViT在计算效率和准确率上优于目前的一些Vision Transformer方法的实验结果：

5 Limitation and Future

SuperViT的Token-Keeping-rate机制虽然聚焦于重要区域，但是其会丢失一些次要信息；此外多尺度的方式会降低分辨率，这也会丢失一些信息。那么对于一些对于细节恢复较高的任务，这种方式用的时候要谨慎选择。

6 Conclusion

1. 文章推出了一种比现有SOTA计算效率更高且有更佳的识别准确度——SuperViT。
2. SuperViT提供了多种复杂度，可以对计算效率和准确性进行调节从而达到较好的trade-off。
3. SuperViT的核心是3个：①多尺度的输入patch-size；②多种token-keeping-rate；③由CE和KL散度组成的损失函数。
4. 多尺度的patch-size可以平衡较大尺度带来的低计算消耗问题、高感受野以及较低尺度带来的性能提升问题；Keeping-rate可以降低计算消耗以及专注于重要信息区域；此外多种keeping-rate可以应对不同图片所带来重要性程度不同的问题。
5. SuperViT适合于目标识别这种探究图片结构信息的任务，但是对于需要细节处理的任务来说值得考究。