英伟达-A-ViT模型| 自适应选择tokens的高效视觉Transformer模型

[189****6672]

英伟达-A-ViT模型| 自适应选择tokens的高效视觉Transformer模型

论文标题：A-ViT: Adaptive Tokens for Efficient Vision Transformer

论文地址：https://arxiv.org/pdf/2112.07658.pdf

代码地址：None

摘要

1. 引入主题：介绍了一种 自动调整不同复杂度图像推理cost 的视觉方法A-ViT。
2. 解决方法： A-ViT通过在推理过程中自动减少视觉Transformer中的tokens数量来实现这一点。视觉Transformer吸引人的体系结构特性使自适应tokens暂停机制能够在不修改网络体系结构或推理硬件的情况下加快推理速度。证明了，A-ViT不需要额外的参数或子网络来进行暂停机制（暂停部分tokens停止训练）。
3. 实验结果： 在图像分类任务（ImageNet1K）中，A-ViT在过滤信息空间特征和减少总体计算量方面具有很高的效率。该方法将DeiT-Tiny和DeiT-Small的吞吐量分别提高了62%和38%，准确度仅下降0.3%，大大优于现有技术。

算法

考虑一张图像$x\in R^{C\times H\times W}$作为模型输入，整个模型的计算过程可以表示如下：

其中编码网络$\mathcal{E}(\cdot )$将图像x转化为序列tokens $t\in {\mathcal{R}}^{K\times E}$，K表示tokens总数，E表示每个token的嵌入维度。

$\mathcal{C}(\cdot )$后处理最后获得的类token，用于最终分类。

L个中间Transformer block $\mathcal{F}(\cdot )$通过自注意操作来transform输入，计算过程如下：

其中$t_{1:K}^l$表示所有K个需要更新的tokens。请注意，每一层之间的tokens数量K是允许改变的。这种tokens停止机制提供了可观的计算增益。

Vision transformer的所有层中的tokens使用相同的特征维度E，与CNN相比，这使得暂停机制设计起来变得很容易，因为CNN需要明确处理不同深度的结构尺寸，例如通道数量。

为了自适应暂停tokens训练，为每个token引入了一个根据输入特征计算的暂停分数来作为暂停概率$h_k^l$：

其中$H(\cdot )$是一个暂停模块。强制每个标记$h_k^l$的暂停分数在[0, 1]范围内，并使用多层暂停分数的累计值来判断是否暂停，当累计停车分数超过$1-ϵ$时，就进行token暂停训练

式中，ε是一个小的正常数。

为了进一步缓解相邻层之间对动态暂停tokens的任何依赖，在暂停tokens进行掩码操作，方法是（i）将token值归零，以及（ii）阻止其关注其他token。

通过在MLP层中分配一个神经元来将$H(\cdot )$整合到现有的视觉Transformer中。因此，不引入任何额外的可学习参数或暂停机制的计算。更具体地说，每个token的嵌入维度E已经提供了足够的信息来适应自适应暂停的学习，使暂停分数计算成为可能

其中$t_{k,e}^l$表示$t_k^l$的第e个特征维度

$\sigma (u)$是一个sigmoid函数

如上所述，β和γ是在应用非线性之前调整嵌入的偏置和缩放参数。请注意，这两个标量参数在所有tokens的所有层上共享。嵌入维度E中只有一个维度用于停止分数计算。根据经验，发现简单选择e=0（第一个维度）表现良好.

为了跟踪各层暂停概率的进度，计算如下：

随后形成的暂停概率为

给定h和r的范围，每层每个token的暂停概率始终有界于0，1。整个损失为

视觉Transformer使用一个特殊的类token $t_k$来生成分类预测，我们将其表示为$t_c$。该token 与其他输入tokens 类似，在所有层中更新。应用平均场公式（先前状态的暂停概率加权平均值）来形成输出token和相关任务损失，如下所示：

然后，视觉Transformer可以通过最小化以下内容进行训练：

实验

可视化

Comparison to Prior Art

Ablations