英伟达提出AdaViT：出于DeiT而又快于DeiT

©作者 |小欣

01 简介

自从Vision Transformer(ViT)被提出以来，Transformer已经成为一种流行的神经网络架构，并广泛的应用到计算机视觉领域之中，如图像分类、目标检测、图像生成和语义分割等任务。

但是，Transformer目前最主流的模式仍然是ViT通过将图像分割成一系列有序的Patches来形成Token，并通过在Token之间建模来解决CV底层问题。

目前使用ViT处理图像的计算成本仍然很高，主要是原因是因为Token之间交互计算带来了二次方的计算复杂度。因此，在大量的计算和内存资源中，在服务端或边缘计算端部署ViT是一个具有挑战性的任务。

AdaViT主要研究如何根据输入图像的复杂度自动调整ViT的计算量。目前几乎所有主流的ViT在推理过程中都有一个独立于输入的固定开销。然而，预测任务的难度随着输入图像的复杂性而变化。

例如，从具有同质背景的图像中对汽车和人进行分类是相对简单的；而在复杂的背景下区分不同品种的狗则更具挑战性。即使是在一张图像中，包含物体详细特征的Patch比那些来自背景的Patch提供的信息要多得多。

针对神经网络的输入依赖问题在以前的工作中已经得到了研究。Graves提出ACT，将神经网络模块的输出作为一个mean-field模型定义的Halt分布。

ACT将离散的暂停问题松弛为一个连续的优化问题，使总计算的上限最小化；最近，随机方法也被应用于解决这个问题，利用几何模型的输出分布，让一些层在训练前期便停止更新。

fignov等人提出了一种ACT的空间扩展，它沿着空间而不是残差层暂停卷积操作。

这种方法并不会带来更快的推理，因为高性能硬件仍然依赖于密集计算；AdaViT作者遵循ACT提出了一种基于输入的ViT自适应推理机制，其中对残差层中的所有Token同时停止计算。作者还观察到，这种方法仅仅减少了计算的一小部分，但造成了不小的精度损失。

为了解决上述精度问题，作者提出了AdaViT，一种空间自适应推理机制，它在不同深度暂停不同Token的计算，以动态的方式保留有价值的Token参与计算。

与卷积特征映射中的逐点ACT不同，AdaViT可以直接由高性能硬件支持，因为停止的Token可以有效地从底层计算中移除。

此外，在不引入额外参数的情况下，可以利用模型中已有的参数来学习整个自适应机制。

02 AdaViT方法

假如让一张图像

作为输入，并通过下式进行预测，得到输出结果y：

其中，编码网络

将x的图像patch标记为token

，K为token的总数，E为每个token的嵌入维数。C是class token的后处理，而L中间Transformer则通过Self-Attention对输入进行转换。

<