©作者 |小欣
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对输入进行转换。
考虑第l层的Transformer Block，它通过以下方式转换第l-1层的所有token:
其中
表示所有K个待更新的token。
ViT为各个层中的所有token使用一致的特征维度E。这使得学习以联合方式监视所有层的全局Halt机制变得很容易。
这也使得Transformer的Halt设计比CNN的Halt设计更容易，后者需要明确处理不同的架构尺寸，例如不同深度的通道数量。
为了自适应Halt tokens，这里为每个token引入一个依赖于输入的Halt Score，作为Token k在第l层的Halt概率
为:
其中H为Halt模块。与ACT类似，将每个token的Halt Score强制在[0,1]范围内，并在推理深入到更深层次时使用累加进行Halt token。
为此，在累计Halt Score超过
时进行token halt:
其中
是一个比较小的正常数，这里允许在一层之后Halt。
同时，为了进一步减轻对相邻层之间动态Halt Token的依赖，如果一旦它被Halt，则将为所有剩余深度
mask为一个标记
，主要是通过以下方式：
(i) 将token值归零；
(ii)阻止其对其他token的注意，屏蔽其对
的影响。
这里定义
来强制所有Token在最后一层Halt。
Token mask使训练的计算成本与原始ViT的训练成本相似。然而，在推理时从计算中删除Halt Token以获得的实际加速。
通过在MLP层中分配单个神经元来H合并到现有的ViT Block中。因此，不需要引入任何额外的可学习参数。
更具体地说，作者观察到每个token的嵌入维数E有足够的容量来适应自适应Halt学习，使得Halt score计算为:
β 和 γ 是在应用非线性之前调整嵌入的平移和缩放参数。注意，这两个标量是所有token在所有层中共享的。
这里仅仅使用嵌入维数E中的一个项进行Halt计算。从经验上看，观察到e=0的简单选择表现良好，而不同的指标并不会改变原来的表现。
通过实验观察到Halt机制在没有引入额外的参数或子网络的情况下效果超过2个标量参数β和γ所带来的效果。
为了进一步跟踪各层Halt概率，每个token计算余数如下:
随后形成一个Halt概率如下:
给定h和r的范围，每个token在每一层的Halt概率总是限制为
。
建议早期Halt的整体损失是通过辅助变量r (reminder)表示的:
对每个token的
进行平均。ViT使用一种特殊的class token
来产生分类预测
这个token类似于其他输入token，在所有层中都进行更新。并通过采用mean-field公式将输出token生成为
，并将相关的任务损失表示为:
其中
表示思考损失。算法1描述了整个计算流程，图2描述了用于视觉解释的相关Halt机制。
在这个阶段，目标函数鼓励在不同深度下考虑不同token时进行精度-效率权衡，从而实现自适应控制。
图2 AdaViT案例
式10中的一个关键因素是
,它平衡了目标应用程序的Halt强度和网络性能。
更大的
值带来更大的惩罚，因此学会了更早地Halt Token。尽管在减少计算方面有效果。
但之前在自适应计算方面的工作发现，训练对
的选择很敏感，其值可能不能提供精确-效率权衡的精细控制。
在ViT中也观察到了类似的现象。为了解决上述问题，在
正则化之前引入了一个分布，使得token在目标深度能够均匀地被Halt。
在这种情况下，对于无限数量的输入图像，期望token的深度在分布先验中变化。相似的先验分布最近被证明是稳定收敛在随机思考。
为此，定义一个Halt Score分布:
在网络的每一层上为所有token计算的平均的Halt Score。
并使用这个平均值作为一个估计将Halt概率分布在各层，这里使用KL散度正则化这个分布到一个预定义的先验。形成新的分布先验正则项为:
这里KL指的是Kullback-Leibler散度，还有
表示带有指导Halt层的目标Halt Score分布。
这里用高斯分布的概率密度函数来定义以预期的Halt深度为中心
的bell-shaped分布
。
直观地说，这鼓励每个token的Halt Score之和触发退出条件
。
最终的损失函数表示为:
其中
是一个标量系数，平衡了分布的正则化和其他损失项。
03 实验
图3显示了在使用AdaViTT对ImageNet1K验证集进行推断时自适应控制的token深度。值得注意的是，自适应token halt为高度区分和突出的区域(通常与目标类相关)启用了更长的处理时间。
此外，对相对不相关的token及其相关计算的非常有效的halt。例如，在动物类上的方法保留了眼睛、纹理和目标对象的颜色，并对它们进行深度分析，同时使用较少的图层处理背景(例如，鸟周围的天空，海洋周围的海洋动物)。注意，即使被标记为不重要的背景token仍然积极地参与初始层的分类。
此外，自适应token可以很容易地跟踪重复的目标对象，如图3中最后一行的第一张图所示，甚至可以完全屏蔽不相关的对象(见最后一行的第二张图)。
实验结果显示，AdaViT将DeiT-S的吞吐量提高了62%，DeiT-T的吞吐量提高了38%，而在ImageNet1K上的准确率只下降了0.3%。
私信我领取目标检测与R-CNN/数据分析的应用/电商数据分析/数据分析在医疗领域的应用/NLP学员项目展示/中文NLP的介绍与实际应用/NLP系列直播课/NLP前沿模型训练营等干货学习资源。