【Transformer】21、AdaViT: Adaptive Tokens for Efficient Vision Transformer

一、背景

Transformer 在多个任务上都取得了亮眼的表现，在计算机视觉中，一般是对输入图像切分成多个 patch，然后计算 patch 之间的自注意力实现下游任务。

但由于自注意力机制的计算量是和输入图像大小呈平方关系的，所以，在边端设备上使用 Transformer 成为了一个问题。

作者认为不同的输入图像对网络来说，预测难度是不同的。如一个车和一个人在干净的背景中，那么就很好识别。如果是多个不同的动物在复杂的背景中，那么就较难识别。

基于此，作者实现了一个网络结构，根据输入的难度，来动态的调节 token 的个数来控制 transformer 的计算复杂度。

二、方法

vision transformer 的过程如下：

• $ϵ(.)$： encoding network，把输入图像编码成 positioned token
• $C(.)$：class token 的后处理
• $L$：transformer block
• $F(.)$：self-attention

为了动态杀掉 tokens，作者为每个 token 引入了一个 input-dependent halting score：

• $H(.)$ 是 halting module
• $k$ 是 token 索引，$l$ 是层

• $t_{k,e}^l$ 是 $t_k^l$ 的第 $e$ 维
• $\sigma$ 是 logistic sigmoid 函数
• $\beta$ 和 $\gamma$ 是非线性操作之前使用的平移和缩放系数

为了根据 layer 来追踪 halting probabilities，每个 token 会计算一个补充参数：

halting probabilities 如下：

ponder loss ：每个 token 的 ponder loss 会平均。

分类任务的损失为：

halting score distribution 分布为：

所以使用 KL 散度来衡量真实和预测的分布偏差：

则总损失为：

三、效果

从图 3 可以看出， adaptive 选择 token 能够对高度突出和巨变的区域产生强响应，通常和类别相关。

1、Token 颜色深度分布：

在图中绘制 token 的颜色，如图 4 所示，其实是一个以图像中心为中心的 2D 类高斯分布，这也说明 ImageNet 的大多数样本都是在中间的。很多计算量都来自于中间区域，边缘参与计算的很少。

2、Halting score distribution：

如图 5 绘制了每个图像的每个 layer 的 halting score。

随机采样了 5k 验证集，在前几个 layer，halting score 随着 layer 的加深而增大，后面慢慢减小。

3、难样本和简单样本

图 6 展示了难例和简单例和其各自所需的计算量。

简单的例子可以被正确分类，AdaViT 处理的也比难例快。

4、类别敏感性

起初非常确信或非常不确信的样本被 adaptive 影响的很小，adaptive 推理能够提升形状明显的类别，如独立的家具或动物。