AAAI 2023 | CF-ViT：由粗到细的两阶段动态推理架构

2022 年 11 月 24 日 PaperWeekly

©作者 | 陈锰钊

单位 | 厦门大学

研究方向 | 模型压缩与加速

此文章介绍刚被 AAAI 2023 会议接收的动态 ViT 论文 CF-ViT: A General Coarse-to-Fine Method for Vision Transformer。代码、训练日志、预训练模型皆已在 GitHub 上开源。

论文题目：

CF-ViT: A General Coarse-to-Fine Method for Vision Transformer

论文链接：

https://arxiv.org/abs/2203.03821

代码链接：

https://github.com/ChenMnZ/CF-ViT

TL;DR（太长不看版）

Vision Transformers (ViTs) 的核心操作是 self-attention，而 self-attention 的计算复杂度和输入 token 数目的平方成正比，因此压缩 ViT 计算量最直接的思路就是减少推理时的 token 数目，也即减少图像划分的 patch 数目。

论文通过两阶段自适应推理减少推理过程的 token 数目：第一阶段将图像划分为粗粒度（大尺寸）patch，目的是使用较少的计算量识别“简单”样本；第二阶段将第一阶段中信息量高的粗粒度 patch 进一步划分成细粒度（小尺寸）patch，目的是使用较少的计算量识别“困难”样本。

论文还设计了全局注意力用于识别信息量高的粗粒度 patch，以及特征复用机制用于增大两阶段推理的模型容量。 在不影响 Top-1 准确率的情况下，该方法在 ImageNet 上将 LV-ViT-S 的 FLOPs 降低 53%，GPU 上实测推理速度也加快了 2 倍。

▲ 图1. CF-ViT两阶段推理示例

研究动机及简介

近年来，Vision Transforms (ViTs) 受到了学术界的广泛关注，被应用到各式各样的视觉任务中。ViT 将图片切割成固定数目的 patch，之后将每个 patch 线性映射成 token，输入一系列堆叠的 Transformer Block。其中，Transformer Block 包含两个组件：Multi Head Self-Attention (MHSA) 和 Feed Forward Network (FFN)。

假设图像被编码成个维 token，则 Transformer Block 的计算复杂度为：

如上式，即 MHSA 的计算复杂度和 token 数的平方成正比，FFN 的计算复杂度和 token 数成正比。 因此，一种简单的高效推理思路是通过降低推理时的 token 数目来对 ViTs 的计算量进行压缩。

那么问题来了，如何降低 token 数目呢？token 对应的是图像中的 patch，也即图像中的某个区域，减少 token 数目的本质在于挖掘图像的空间冗余性，以减少冗余区域的 token 编码数量。而图像的空间冗余性根据粒度，其实可以划分为两种：

1. 图像之间的空间冗余性： 根据图像识别难度，可以将图像划分成简单图像和困难图像。如下图所示：对于简单图片，使用较低计算量就能够正确识别。对于困难图片，可能只有使用较高的计算量才能实现正确识别。

▲ 图2. 图像之间的空间冗余性

2. 图像内部的空间冗余性： 根据图像的语义信息，可以将图像划分成前景区域和背景区域，模型能否成功识别主要依赖于前景区域。因此，我们可以认为前景部位是重要区域，背景区域是不重要区域。假设计算量是可自适应调度的，显然，对于重要区域，应该分配更多的计算量；而对于不重要区域，则使用较少的计算量。

▲ 图3. 图像内部的空间冗余性

文中探索了上述两种空间冗余性在 ViT 上的特殊表现形式 ：

1. 图像之间的空间冗余性表现为大部分图像编码成少量 token 即可分类成功。 如表 1 所示，196 个 token 计算量是 49 个 token 的 4.18 倍，但带来的性能提升只有 6.6%，这表明，大部分图片较为简单，例如 ImageNet-1k 验证集中 73.2% 的图片用 49 个 token 编码就能成功分类。若是推理过程中对于简单图片使用较少 token 数量进行编码，则可以在保证性能的同时减少计算量。

▲ 表1. DeiT-S使用不同token编码数目时的Top-1准确率和计算量

2. 图像内部的空间冗余性表现为只将重要的 token 输入模型，就可以获得较高准确率。 我们采用预训练好的 DeiT-S 模型作为 baseline，其将图像编码成 196 个 token。

如图 4 所示，我们使用 attention map 在输入端对 token 进行筛选：分别塞选出 100 个注意力值最大的 token (绿色) 和 100 个注意力值最小的 token (红色)。可以发现，保留注意力值最大的 100 个 token 性能优势极为明显。其中，attention map 通过先将图片输入模型前向推理一次获得。

▲ 图4. 筛选出100个最重要的token输入模型的Top-1准确率

方法详述

受上述所观察现象的启发，同时考虑图像之间的空间冗余性与图像内部的空间冗余性，论文提出了一个由粗到细的两阶段动态推理架构 (Coarse-to-Fine Vision Transformer, CF-ViT)。

▲ 图5. Coarse-to-Fine Vision Transformer, CF-ViT

首先介绍 CF-ViT 的推理过程。如图 5 所示：

1. 第一阶段是利用图像间空间冗余性的粗粒度推理，其目的是使用较少计算量获得”简单“图片的正确识别结果。 具体做法是对图像使用粗粒度 patch 分割，分割后的结果编码成 token 输入 ViT 模型，以低成本获得粗粒度推理结果。紧接着判断粗粒度推理结果是否可信。此处 follow 此前 early-exiting 工作，采用预测置信度与一个预设阈值进行比较。若预测置信度大于阈值，则认为结果可信，结束推理；否则认为结果不可信，进入第二阶段推理。

2. 第二阶段是利用图像内空间冗余性的细粒度推理，其目的是使用较少计算量获得”困难“图片的正确识别结果。 具体做法是识别图像中重要的粗粒度 patch，然后将每一个重要的粗粒度 patch 进行进一步的细粒度切割，获取对应的四个细粒度 patch。最后将新的 patch 分割结果编码成 token 输入 ViT 模型以获得细粒度推理结果。

值得注意的是，两个阶段中用到的 ViT 模型是共享参数的，因此减少了多阶段推理的参数负担。同时，为了能够更好的协同两个推理阶段，论文设计了两个协同模块，以下进行简单介绍，详见 paper：

1. 第一个模块是重要区域识别（Informative Region Identification） , 其主要通过统计全局类注意力图来衡量每一个粗粒度 patch 的重要性，使得第二阶段可以轻量化无负担的筛选图像中的重要区域，然后进一步进行细粒度切割。

2. 第二个模块是特征复用（Feature Reuse） ，其将第一阶段提取到的特征融合进第二阶段的推理过程。具体做法是对第一阶段提取到的特征进行线性变换和维度对齐，然后直接和第二阶段的图像编码相加，以此实现跨粒度特征注入，从而进一步提高第二阶段的模型容量。

在训练时，我们使用真实标签监督细粒度推理结果，使用细粒度推理结果监督粗粒度推理结果：

实验结果

实验中，论文基于 DeiT-S 和 LV-ViT-S 测试了所提出的方法。如表 2 所示，在不影响模型准确率的前提下，CF-ViT 分别分别将 DeiT-S 与 LV-ViT-S 的计算量降低 61% 和 53%。

值得注意的是，表示所有的样本都会输入细粒度推理阶段，此种情形下 CF-ViT 分别将 DeiT-S 和 LV-ViT-S 的 Top-1 准确率提升和，且使用更少的计算量。同时，我们也测试了 CF-ViT 在 ImageNet 上的实际推理速度：基于 Nvidia A100 GPU, batch size , CF-ViT 可以显著提高推理速度。

▲ 表2. CF-ViT和Backbone的比较

表 3 展示了 CF-ViT 和只利用图像内部空间冗余性的 token 剪枝方法的对比结果。

▲ 表3. 和现有的token剪枝方法对比

图 5 展示了 CF-ViT 和现有只利用图像间空间壳余性的 early-exiting 方法的对比结果。值得注意的是，CF-ViT 不仅性能优于 DVT，且得益于两个推理阶段共享模型参数，其参数量仅为 DVT 的。

▲ 图5. 和现有的early-exiting方法对比

图 6 展示了在粗粒度推理阶段被正确分类和在细粒度推理阶段被正确分类的样本，其中细粒度推理阶段没有灰色表示的是模型识别出的重要区域。可以发现，在第一阶段就被正确识别的往往是”简单“样本，其特征是目标物体较为显著，背景简单；而在第二阶段才被正确识别的往往是”困难“样本，目标物体较小，且背景较为杂乱。同时，可以发现所提出的重要区域识别模块往往能够成功定位到目标物体，符合预期。