Evo-ViT: Slow-Fast Token Evolution for Dynamic Vision Transformer

论文链接：https://arxiv.org/pdf/2108.01390v3.pdf

代码链接：https://github.com/YifanXu74/Evo-ViT

一些想法

这个论文和 ICLR 2022 的Expediting vision transformers via token reorganization 非常像，都是把有效性较低的 token 合并不参与 attention 计算以达到加速的目标。不过，这个论文提出了一个 global 的策略，用来评价 token 的有效性，但是 ICLR2022 是局部在层里计算。

Motivation

虽然最近Vit模型取得了巨大的成功，但其巨大的计算成本仍是存在的主要问题。目前高效设计方案主要包括两种：一种是结构化压缩（比如：shift windows 在局部计算注意力，或者PVT里面对attention简化计算），另一种是进行非结构化的网络裁剪（比如：DynamicViT）。但是第一种方法会以相同的优先级处理信息token和占位符token(相对不重要的背景token)，第二种方法则会造成信息结构不完整。

主要方法

1. Token 选择策略

作者提出了global class attention，通过每个token的得分，定义得分最高的前K个token为 information token，其它为 placeholder token，并将 placeholder token保留下来，以进行后续的操作。

cls attention 的计算方法：$A_{cls}=\text{Softmax}\left(\frac{q_{cls}\cdot K^T}{\sqrt{C}}\right)$

全局 class attention 的计算公式为：$A_{cls,g}^k=\alpha \cdot A_{cls,g}^{k-1}+(1-\alpha )\cdot A_{cls}^k$

其中 $A_{cls,g}^k$表示第$k$层的全局 class attention，用来对 k+1 层的token进行筛选；$A_{cls}^k$表示第$k$层的class attention。

慢速-快速双流更新

1） 慢速更新方法： 对于N个 token，通过 global class attention 将其分为$k$ 个 information token $x_{inf}$和 $N-k$ 个placeholder token $x_{ph}$，之后将 placeholder token 使用求和操作合并为一个代表性token $x_{rep}$。代码如下：

#add_token: B*1*C
add_token = merge_tokens(x_drop, global_attn[:, N_:])
#x:B*(N+1)*C
x = torch.cat((x, add_token), dim=1)  # B,N+1,C


def merge_tokens(x_drop, score):
    # score B,N
    # scale
    weight = score / torch.sum(score, dim=1, keepdim=True)
    x_drop = weight.unsqueeze(-1) * x_drop
    return torch.sum(x_drop, dim=1, keepdim=True)

之后进行常规的 transformer 运算：

2） 快速更新策略： 对placeholder token $x_{ph}$ 通过 $x_{rep}$ 的残差进行快速更新：

其中，${\varphi }_{exp}$ 是复制操作。

训练策略

作者使用了一种分层到阶段(layer-to-stage)的训练策略。具体的是在前200个epoches每一层都使用token选择而在后100个epoches只在每一个阶段开始时进行token选择；对于类似DeiT这种transformer结构，作者将每四个block列为一个阶段stage。

实验结果与分析

作者主要是通过将本文的方法加在DeiT以及Levit网络结构上对ImageNet-1k分类数据集进行分类。从实验一是与主流的 token 裁剪方法比较，可以看出虽然Evo-Vit的准确率有非常小的下降，但是吞吐量上升了近60%。

实验二是与主流的VIT模型进行比较，可以看出，在处理高分辨率图像时该方法可以以较低的精度下降来较快加速LeVit的运行。

下面是消融实验，native selection表示使用每层中的class attention 来进行占位符token的删除；structure preservation表示保留占位符token但慢速更新。

下面是可视化结果，可以看出本文方法可以较好的分辨 information token 和 placeholder token