Multi-Scale Densenet续作？搞定Transformer降采样，清华联合华为开源动态ViT！

作者丨小马
编辑丨极市平台

先验知识

Transformer最近在CV领域展现出了不错的效果，Vision Transformer（ViT）的大致流程可分为两步：

1）因为Self-Attention（SA）的计算复杂度是和输入特征的大小呈平方关系的，所以如果直接将224x224的图片输入到Transformer中，会导致计算量的“爆炸”。因此，ViT的第一步是将图片转换成更小的token（比如16x16），然后将这些token进行flatten后输入到Transformer中。

2）利用Transformer进行视觉特征的提取并分类。

在第一步中，图片转换成token的操作相当于是一个降采样的操作，降采样的倍数越大（获得的token越少），损失的信息越大，ViT的计算速度也就越快。降采样的倍数越小（获得的token越多），损失的信息越小，ViT的计算速度也就越慢。在以前的ViT中，没有考虑图片的内容，将所有图片都转换成固定大小的token，就会导致对于“简单图片”，这样的“降采样”过于精细，造成计算资源的浪费；对于“复杂”图片，这样的降采样细粒度不足，造成精度的损失。因此，本文作者就提出了动态的ViT来解决这个问题。

Noting：感兴趣的同学可以再去看看黄高老师的Multi-scale DenseNet[1]，我个人觉得，这篇文章的Motivation和Multi-scale DenseNet很像。

Multi-scale DenseNet的Motivation这样的：对于分类网络的测试而言，有些输入图像是网络容易分类的（也就是简单样本，通常是主体比较明确，背景比较清晰），有些输入图像是网络难分类的（也就是困难样本，通常是主体被遮挡会很小，背景比较复杂）。以softmax输出为例，假设一个2分类的例子，如果一张图像属于两个类的概率分别是0.01和0.99，和一张图像属于两个类的概率分别是0.51和0.49，显然相比之下前者更容易分类。那么从这个点出发，作者就想到了能不能对于简单图像仅采用浅层的速度较快（对应到这篇文章就是用更少的token）的网络来分类，而对于难分类图像再采用深层的（对应到这篇文章就是更多的token）速度慢一点的网络来分类。

1. 论文和代码地址

Not All Images are Worth 16x16 Words: Dynamic Vision Transformers with Adaptive Sequence Length

论文地址：https://arxiv.org/abs/2105.15075

代码地址：https://github.com/blackfeather-wang/Dynamic-Vision-Transformer

2. Motivation

首先作者做了一个实验（Table1），用14x14的token能达到76.7%的准确率；用4x4的token就能达到60.8%的准确率，这说明ImageNet中很多的图片用4x4的token就能识别准确了（也就是有很多图片都是简单样本）。用14x14token计算量是4x4的token计算量的8.5倍，因此全部图片都用14x14的token在计算上其实是一种浪费。

因此，作者就提出了先用少的token进行预测，如果置信度能够大于某一个阈值，就直接输出结果；否则就用更多的token进行预测，如果置信度能够大于某一个阈值，就直接输出结果…一直循环，直到最后一层ViT（用最多的token进行预测）。

3. 方法

3.1.1. Overview

首先我们来看整个网络的结果，非常简单，由一系列串联的ViT组成，每个结构都是接收了不同token数量表示的图片。在训练的时候，是从上层到下层，全部结构都训练；而在测试的时候，每经过一个ViT，都会对当前输出进行一次验证，如果prediction大于阈值，那么就直接输出类别；否则就进入下一层ViT进行预测。

训练时的损失函数如下：

除了串联的ViT，这个结构中三个部分值得注意1）Feature Reuse；2）Relationship Reuse；3）Exit（Adaptive Inference）。下面我们详细讲一下这三个结构的具体实现：

3.1.2. Feature Reuse

下层的ViT如果只接受更多token表示的图片信息，而将上层通过ViT后的特征进行丢弃，这就意味着上面ViT的计算结果对于下面的ViT计算毫无用处，这就造成了计算资源的浪费。因此Feature Reuse就是一个特征重用的模块，将上游ViT最后一层输出的特征取出，经MLP变换和上采样后，作为上下文信息，以Concat的方式整合入下游模型每一层的MLP模块中。

具体实现如上图的灰色框中所示的一样，首先将上层ViT的特征通过LN-MLP进行非线性变换并压缩维度：

然后，将上面的到的上下文信息El以concat的方式与当前层ViT的特征进行融合，再进行LN-MLP：

3.1.3. Relation Reuse

相比于CNN，ViT的一个显著优点是，它的Self-Attention能够聚合整个图像中的信息，从而有效地对长序列信息进行建模。既然不同token之间的关系也是非常重要的，作者进一步对上游ViT的关系矩阵进行了重用，使得下层的ViT能够进行更加准确的关系建模，这一步就是论文中的Relation Reuse。

Relation Reuse的操作如上图的灰色框所示：

对于每一层ViT，如果不进行Relation Relation，那么Self-Attention的计算过程为：

现在我们需要将上层的attention map进行重用，第一步，我们首先将上层所有的attention map进行concat：

然后我们将这些attention map通过一个转换网络rl()（结构如下图所示），使得大小和维度与当前层的attention对应；然后将转换后的上层attention map与当前层的attention进行相加，得到relation reuse的attention map。

3.1.4. Exit（Adaptive Inference）

对于第i层的预测，我们可以得到预测结果$p_i$，并且第i层会有阈值${\eta }_i$，如果$p_i>{\eta }_i$，那么就可以执行exit操作，输出预测结果；否则进入下一层ViT。

在FLOPs<B的情况下，我们希望能到更高的准确率，这就变成了一个优化的问题：

作者采用遗传算法来解决这个问题。

4.实验

4.1. 在ImageNet上的实验结果

上面2张图时分别用T2T-ViT和DeiT作为ViT结构的实验，可是看出在相同的准确率下，本文的结构确实可以加速很多。

下表展示了，在相同的速度下，本文准确率更高；在相同准确率下，本文速度更快。

4.2. CIFAR-10/100上的结果

从上表可以看出，在小数据集上，在相同准确率下，本文预测速度更快。

4.3. Feature Reuse的消融实验

可以看出，使用所有上层的Feature信息，比不使用和只使用浅层或者一半的效果更好。

4.3. Relation Resue的消融实验

可以看出，使用所有上层的Relation信息，比不使用和只使用浅层或者一半的效果更好。

4.4. 可视化

从图中可以看出，简单样本是那种背景简单，主体占据图片的绝大部分区域；困难样本则相反。

5. 总结

1、至少在我看来，这篇文章的思想真的跟黄高老师以前的那篇Multi-scale DenseNet的思路真的motivation非常像，当然在实现上还是有所不同的。说明一个好的idea永远不会过时。

2、这篇文章虽然测试的时候能根据图片自适应划分token，但是一次要训练多个ViT，那么训练时间就会大大延长。如果不同token数量的ViT能够参数共享，并且在训练之前就能够通过一个网络自动判断出最合适的token划分方法，那就是真正的input adaptive的动态ViT了。

参考文献

[1]. Huang, Gao, et al. “Multi-scale dense convolutional networks for efficient prediction.” arXiv preprint arXiv:1703.09844 2 (2017).