【ICLR2022】Expediting vision transformers via token reorganization

【ICLR2022】Not all patches are what you need: Expediting vision transformers via token reorganization

文章链接:https://arxiv.org/abs/2202.07800

代码链接:https://github.com/youweiliang/evit

一些个人想法

  • 作者在第5、9、13层应用该方法,这个貌似没有给出原因
  • 图1用的不是特别好,没看明白具体要说明什么
  • 该方法和DynamicVIT是最相关的,最近有一系列相关的方法提出,值得关注
  • 审稿人指出这个方法难以和swin等多尺度方法结合,因此 the real gain could be very limited

Abstract

之前的VIT将图像块作为token,来建立多头自注意力。但是与传统CNN相比,VIT的 global self-attention between image tokens and long-range dependency 使得模型收敛较慢。目前很少有人关注VIT加速,因为和CNN的差异,CNN加速的一些方法(蒸馏和修建等操作)无法直接应用于VIT。在本文中,作者提出在推理时,重新对token进行组织,实现方法如下:

  • 保留关注的图像tokens和融合不关注的图像tokens来重组tokens,以加快后续MHSA和FFN计算,通过这种方式,随着网络的深入,逐渐减少图像tokens的数量,以降低计算成本
  • 在相同数量的输入图像token情况下,该方法减少了MHSA和FFD计算量,实现高效推理

主要方法

在VIT中,最终是通过 [CLS] 来进行分类,和其他tokens的交互通过encoder的注意力机制完成:

\[\textbf{x}_{class}=\textrm{Softmax}\left(\frac{q_{class}\cdot V}{\sqrt{d}}\right)V=\textbf{a}\cdot V \]

这里的\(q_{class}\)表示[CLS],组合系数\(\textbf{a}\)是[CLS] 对所有 tokens 的注意力值。因此,可以用该注意力值来表示 tokens 的重要性。

因此作者使用 [CLS] 来识别 tokens 的重要程度,基于这些参数来移除注意力值最小的 token,但从下面的实验中可以看出,这会严重影响分类精度,所以作者在训练阶段加入了图像 token 重组。

preview

作者的总体思路如下图所示:保留\(k\)个注意力值最大的 tokens,叫做 attentive tokens,然后把其它的 tokens 合并为一个新的 token。

preview

对于 inattentive token,虽然包含较少的信息,作者认为它们仍然是有助于预测结果的,同时,作者是选择保留了固定数量的tokens,所以当图像中的对象较大时,移除与之相关的部分,会对性能造成负面影响。综上,作者将不重要的tokens进行融合(加权平均),融合后将继续向后传播 。

实验分析

实现细节

本文的 token identification module 被添加到 DeiT-S 和 DeiT-B 的第4、7、10层中,以及 LV -ViT-S 的第5、9、13层。另外作者采用了 warm up 的策略来进行 tokens 的标记。让 attention tokens 的保留率由1逐渐降低到目标值

用高分辨图像训练

这里因为对tokens进行了融合减少,所以在维持相同计算量的情况下,允许更多地tokens输入,因此该方法可以输入更高分辨率图像。作者将标准输入图像 224×224 调整为 256×256

preview

可视化

为了进一步研究 EViT 的可解释性,作者进行了可视化实验,可以看到基本上重要的目标区域都得到了保留。

preview

posted @ 2022-03-18 08:00  高峰OUC  阅读(154)  评论(0编辑  收藏  举报