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

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

极市导读

本文介绍了一篇清华大学黄高研究团队和华为联合研究的一篇关于动态Transformer 的最新工作，该项工作发现采用定长的token序列表征数据集中所有的图像是一种低效且次优的做法，并提出一种可针对每个样本自适应地使用最合适的token数目进行表征的动态ViT模型。>>加入极市CV技术交流群，走在计算机视觉的最前沿

先验知识

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层的预测，我们可以得到预测结果，并且第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.4. Relation Resue的消融实验

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

4.5. 可视化

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

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).

本文亮点总结

1.Feature Reuse就是一个特征重用的模块，将上游ViT最后一层输出的特征取出，经MLP变换和上采样后，作为上下文信息，以Concat的方式整合入下游模型每一层的MLP模块中。
2.相比于CNN，ViT的一个显著优点是，它的Self-Attention能够聚合整个图像中的信息，从而有效地对长序列信息进行建模。既然不同token之间的关系也是非常重要的，作者进一步对上游ViT的关系矩阵进行了重用，使得下层的ViT能够进行更加准确的关系建模，这一步就是论文中的Relation Reuse。

如果觉得有用，就请分享到朋友圈吧！

△点击卡片关注极市平台，获取最新CV干货
赞 (0)

相关推荐