Day 14: Not All Images are Worth 16 x 16 Words: Dynamic Visition Transformer 动态视觉 Transformer

这篇是华为和清华共同发表的文章，果然清华出品，还是比较硬核！看完第一遍并没有特别了解，也许需要结合代码再看一遍。

总体上来说，也是为了在不影响效果的情况下，降低计算复杂度。

要点总结

之前的 Transformer 结构计算量都很大，因此需要找方法降低它的计算复杂度。

并非所有的图片都需要被分割成特别多的 token 来参与计算，有的图片相对简单，只需要拆成 $4\times 4$ 的 token，就能有很好的准确率，因此提出了动态视觉 Transformer 的模型。模型的目的是自动地根据每张图片来调节 token 的数量，以达到较高的计算效率。

具体来讲，就是用了一连串 token 数量逐渐增加的 transformer。在测试阶段，这些模型从使用较少 token 的开始依次激活。当一个具有足够高确信度的预测产生后，推断过程就会立刻中止。这样做的结果，就是通过对 token 数量的调控，整个计算量是非平均地分布在“简单”和“困难”的样本之间的。

此外，我们还提出了“特征间（feature-wise）”和“关系间（relationship-wise）”复用机制，来减少冗余的计算。其中，前者使得下游模型能在之前提取出来的深层特征上进行训练，后者则使得模型可以利用现有的上游任务中产生的自注意力关系，来学习到更准确的注意力图。

其它不太重要的一些点

DVT 是作为一个通用框架来设计的，像目前很多 SOTA 的模型，比如 ViT, T2T-ViT，都可以直接作为 backbone 来部署。

通过调节早中止（early-termination）的条件，DVT 的计算量可以在训练时动态调节。

Dynamic Vision Transformer

概述

我们提出部署多个 Transformer 结构，而且每个 Transformer 结构使用的 token 数量逐渐增加。这样，在测试中，我们可以挨个激活它们，直到得到一个可以接受的预测准确率。

如果 transformer 分别单独训练，那当下游的 transformer 被激活后，上游模型所有的计算结果都会被抛弃，这样会导致模型非常低效。因此作者才提出特征和关系的复用机制，来解决这个问题。

如上图，对于每张图片，我们首先用比较少的一维 token 来嵌入它。我们可以直接将图片块拉成一维，或者用 tokens-to-token 中的方法。我们通过这样使用比较少的 token，来快速地得到一个预测结果。这么做的好处是效率非常高，因为计算量是和 token 数目成二次方增长的。

接下来，根据评价指标的结果，我们会决定当前的结果是否可以立即提取出来。在本文中，当结果足够可信时，我们会使用早中止（early termination）。

但当结果不够可信时，我们会将原始的图片拆成更多的 token，以通过更复杂的计算来得到更精确的结果。需要注意，此处我们仅仅是 token 的数量变多了，但每个 token 的 embedding 维数是不变的。

此时，一个有着完全相同架构，仅仅是参数不同的 transformer 会被激活。根据设定，这一阶段，我们用更复杂的计算来换取在“困难”样本上更准确的结果。

为了使模型更高效，这一新模型可以复用之前模型中学习到的特征和关系。

同样的，当得到新的预测结果后，我们会再次对它进行评价，看是否中止的条件。如果没有达到，则上面的过程会再重复，直到最后退出（sample exits）或者最后一个 transformer 的推断结束。

训练

训练阶段，我们训练 DVT，在每个出口处（如图，每个有着对应 token 数量的 transformer 的最末端）产生正确的预测结果。最终优化目标为：

$$\mathrm{minimize}\frac{1}{\left|{\mathcal{D}}_{\text{train }}\right|}\sum _{(x,y)\in {\mathcal{D}}_{\text{tain }}}\left[\sum _i{L}_{\mathrm{C}\mathrm{E}}\left({\mathbf{p}}_i,y\right)\right]$$

其中，$(x,y)$ 表示训练集${\mathcal{D}}_{train}$ 中的一个样本及标签。$L_{CE}$ 是交叉熵 cross-entropy。$p_i$ 是第$i$个出口处预测概率经过 softmax 函数后的结果。我们发现这么一个简单的训练目标在实际训练中就有不错的效果。

特征和关系复用（feature and relationship reuse）

设计中的一个主要挑战就是，如何复用先前的计算。也就是，当下游含有更多 token 的 transformer 推断结束后，如果直接将上游中先前得到的计算结果直接抛弃，显得是非常低效的。尽管上游模型中的 token 数量较少，但毕竟它们的目标都是一样的，有着同样有价值的信息。因此，作者提出下面的两种复用机制。二者都能在轻微地增加计算量的情况下，显著地提高测试准确率。

Transformer 回顾

z l ′ = MSA ⁡ ( LN ⁡ ( z l − 1 ) ) + z l − 1 , l ∈ { 1 , … , L } z l = MLP ⁡ ( LN ⁡ ( z l ′ ) ) + z l ′ , l ∈ { 1 , … , L } \begin{array}{ll}\boldsymbol{z}_{l}^{\prime}=\operatorname{MSA}\left(\operatorname{LN}\left(\boldsymbol{z}_{l-1}\right)\right)+\boldsymbol{z}_{l-1}, & l \in\{1, \ldots, L\} \\\boldsymbol{z}_{l}=\operatorname{MLP}\left(\operatorname{LN}\left(\boldsymbol{z}_{l}^{\prime}\right)\right)+\boldsymbol{z}_{l}^{\prime}, & l \in\{1, \ldots, L\}\end{array} zl′​=MSA(LN(zl−1​))+zl−1​,zl​=MLP(LN(zl′​))+zl′​,​l∈{1,…,L}l∈{1,…,L}​

其中${\mathbf{z}}_l\in {\mathbb{R}}^{N\times D}$ 是第 $l$ 个 transformer 层的输出，$N$ 是每个样本的 token 数，它等于 $HW+1$，1 是 classification token。$D$ 是每个 token embedding 的维度

特征复用

我们使用上游 transformer 最后一层中输出的图片 token，比如 ${\mathbf{z}}_L^{\text{up }}$，来学习下游模型中的 layer-wise 嵌入 $E_l$：

$${\mathbf{E}}_l=f_l\left({\mathbf{z}}_L^{\text{up }}\right)\in {\mathbb{R}}^{N\times D^{\prime }}$$

上式中，$f_l:{\mathbb{R}}^{N\times D}\to {\mathbb{R}}^{N\times D^{\prime }}$ 由一系列的操作组成，这些操作的第一步是 LN-MLP $\left({\mathbb{R}}^D\to {\mathbb{R}}^{D^{\prime }}\right)$，作用是引入非线性，以及增加一些更灵活的变换。

接下来，这个图片 token 被 reshape 到原始图片中对应的位置，再被上采样，然后展平，来匹配下游模型中的 token 数量。通常我们使用一个比较小的 $D^{\prime }$ 来获得一个更高效的 $f_l$。

结果就是，上游任务中得到的嵌入 $E_l$ 被输入到下游模型中，提供了一些可以用来识别图片的先验知识。这样，上面的传统 transformer 中的第二个方程，可以被替换成下面的这个：

z l = MLP ⁡ ( LN ⁡ ( Concat ⁡ ( z l ′ , E l ) ) ) + z l ′ , l ∈ { 1 , … , L } z_{l}=\operatorname{MLP}\left(\operatorname{LN}\left(\operatorname{Concat}\left(z_{l}^{\prime}, \mathbf{E}_{l}\right)\right)\right)+z_{l}^{\prime}, \quad l \in\{1, \ldots, L\} zl​=MLP(LN(Concat(zl′​,El​)))+zl′​,l∈{1,…,L}

这里的 $E_l$ 与中间的 token ${\mathbf{z}}_l^{\prime }$ concatenate 到了一起，我们只需要在 LN 和 MLP 层中，把维度从之前的 $D$ 换成现在的 $D+D^{\prime }$ 。

此外，我们并不会复用 classification token 和 pad zero。

上面这个特征复用公式，也可以理解为是增大了模型的深度。

关系复用

下游模型也能读取上游模型中的注意力图，因此我们也可以复用这些信息。

给定一个输入 ${\mathbf{z}}_l$，自注意力是这么计算的。首先，Q, K, V 矩阵分别由下面的线性映射得到，

$${\mathbf{Q}}_l={\mathbf{z}}_l{\mathbf{W}}_l^{\mathrm{Q}},{\mathbf{K}}_l={\mathbf{z}}_l{\mathbf{W}}_l^{\mathrm{K}},{\mathbf{V}}_l={\mathbf{z}}_l{\mathbf{W}}_l^{\mathrm{V}}$$

接着，注意力由以下公式得到：

$$\mathrm{Attention}\left({\mathbf{z}}_l\right)=\mathrm{Softmax}\left({\mathbf{A}}_l\right){\mathbf{V}}_l,{\mathbf{A}}_l={\mathbf{Q}}_l{\mathbf{K}}_l^{\top }/\sqrt{d}$$

上式中 $d$ 是 Q 或者 K 中的隐藏维度，${\mathbf{A}}_l\in {\mathbb{R}}^{N\times N}$ 表示注意力图的数值。我们在上式中忽略了多头，因为并不影响我们的描述。

为了复用关系，我们首先将所有上游模型中产生的注意力数值按如下方式拼接（concatenate）起来

$${\mathbf{A}}^{\text{up }}=\mathrm{Concat}\left({\mathbf{A}}_1^{\text{up }},{\mathbf{A}}_2^{\text{up }},\dots ,{\mathbf{A}}_L^{\text{up }}\right)\in {\mathbb{R}}^{N_{\text{up }}\times N_{\text{up }}\times N_{\text{up }}^{\text{At }}}$$

上式中 $N_{up}$ 和 $N_{up}^{att}$ 分别表示上游模型中 token 的数量和所有注意力图的力量。通常情况下，$N_{up}^{att}=N^{H}L$ ，其中 $N^H$ 表示多头注意力中头的个数，$L$ 表示层数。这样一来，下游的 transformer 通过同时利用自身的 token 和 $A^{up}$，来学习注意力图。我们便可将上面的注意力公式换成下面的：

$$\text{ Attention }\left({\mathbf{z}}_l\right)=\mathrm{Softmax}\left({\mathbf{A}}_l+r_l\left({\mathbf{A}}^{\mathbf{u}\mathbf{p}}\right)\right){\mathbf{V}}_l,{\mathbf{A}}_l={\mathbf{Q}}_l{\mathbf{K}}_l^{\top }/\sqrt{d},$$

上式中 $r_l(\ast )$ 是一个转换网络，它将 $A^{up}$ 提供的信息集中到一起，来改善下游中的注意力值 $A_l$。
$r_l$ 网络的结构如上图4，包括一个用来提供非线性的 MLP，和一个用来匹配注意力图尺寸的上采样操作。当采用多头注意力时，MLP 的输入维度设成头的个数。

需要注意的是，普通的上采样操作不能直接运用在 $r_l$ 中，而需要采用像下图中的上采样过程：

上图是将一个 $HW\times HW(H=W=2)$ （也即 4x4）的注意力图，上采样到 $H^{\prime }{W}^{\prime }\times H^{\prime }{W}^{\prime }\left(H^{\prime }=W^{\prime }=3\right)$ （也即 9x9）的过程。由于 $HW$ 的每一行和每一列，都对应着 $H\times W$ 个图片 token，首先将它的行和列 reshape 成 $H\times W$，再将它们缩放至 $H^{\prime }\times W^{\prime }$，然后再拉平成 $H^{\prime }{W}^{\prime }$ 个向量。（不太明白为什么这样效果会更好）

$(HW\times HW)$ —> reshape —> $(H\times W\times HW)$ —> scale —> $(H^{\prime }\times W^{\prime }\times HW)$ —> flatten —> $(HW\times H^{\prime }{W}^{\prime })$

自适应性推断

如上所述，DVT 是逐渐增加每个样本中的 token 个数的，且会执行提早中止，这样一来，“容易的”和“困难的”样本就会用不同的 token 个数来代表，且处理时的计算复杂度也会不同，这样就提高了整体的计算效率。具体地说，在产生 softmax 预测 $p_i$ 的第 $i$ 个出口处，$p_i$ 的最大值 $ma{x}_j{p}_{ij}$ 会和一个阈值 ${\eta }_i$ 进行比较，只有当 $ma{x}_j{p}_{ij}>{\eta }_i$ 时，推断才结束，同时我们将 $p_i$ 输出。否则的话，图片将会被拆成更多的 token，并激活下游的 transformer。此外，在最后一个 transformer 处，我们将 threshold 的值设为 0。

{ η 1 , η 2 , … } \left\{\eta_{1}, \eta_{2}, \ldots\right\} {η1​,η2​,…} 的取值是在 validation set 上决定的。我们假设有一个“给定预算下的批量分类”（即预算是固定的，不能超出预算）的设定，规定 DVT 需要在一个给定的计算量预算下，识别出 validation 数据集上的图片。让 Acc ⁡ ( D v a l , { η 1 , η 2 , … } ) \operatorname{Acc}\left(\mathcal{D}_{\mathrm{val}},\left\{\eta_{1}, \eta_{2}, \ldots\right\}\right) Acc(Dval​,{η1​,η2​,…}) 和 F L O P s ( D v a l , { η 1 , η 2 , … } ) \mathrm{FLOPs}\left(\mathcal{D}_{\mathrm{val}},\left\{\eta_{1}, \eta_{2}, \ldots\right\}\right) FLOPs(Dval​,{η1​,η2​,…}) 分别代表当设置的阈值为 { η 1 , η 2 , … } \left\{\eta_{1}, \eta_{2}, \ldots\right\} {η1​,η2​,…} 时，在验证集上的准确率和计算量开销，最优的阈值可以通过求解以下的优化来得到：

maximize ⁡ η 1 , η 2 , … Acc ⁡ ( D val  , { η 1 , η 2 , … } ) ,  s.t.  FLOPs ⁡ ( D val  , { η 1 , η 2 , … } ) ≤ B \underset{\eta_{1}, \eta_{2}, \ldots}{\operatorname{maximize}} \operatorname{Acc}\left(\mathcal{D}_{\text {val }},\left\{\eta_{1}, \eta_{2}, \ldots\right\}\right), \quad \text { s.t. } \operatorname{FLOPs}\left(\mathcal{D}_{\text {val }},\left\{\eta_{1}, \eta_{2}, \ldots\right\}\right) \leq B η1​,η2​,…maximize​Acc(Dval ​,{η1​,η2​,…}), s.t. FLOPs(Dval ​,{η1​,η2​,…})≤B

由于上式不可微，因此需要通过 Darrell Whitley. A genetic algorithm tutorial. Statistics and computing, 4(2):65–85, 1994 本文中的遗传算法来求解。