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}=\textrm{Softmax}\left(\frac{q_{cls}\cdot K^T}{\sqrt{C}}\right)$

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

其中 $A^k_{cls,g}$ 表示第 $k$ 层的全局 class attention，用来对 k+1 层的token进行筛选； $A^k_{cls}$ 表示第 $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)