[论文解读] Adaptive Token Sampling For Efficient Vision Transformers

Mohsen Fayyaz, Soroush Abbasi Koohpayegani+6 more|arXiv (Cornell University)|Nov 30, 2021

Cell Image Analysis Techniques被引用 4

一句话总结

本文提出自适应标记采样（ATS），一种可微、无参数的模块，可根据输入图像内容动态减少视觉变换器中的标记数量。通过利用分类标记的注意力权重对显著性补丁进行评分并采用逆变换采样，ATS 在 ImageNet、Kinetics-400 和 Kinetics-600 上实现高达 2 倍的 GFLOPs 减少，且准确率无下降，使预训练模型可即插即用，实现效率提升。

ABSTRACT

While state-of-the-art vision transformer models achieve promising results in image classification, they are computationally expensive and require many GFLOPs. Although the GFLOPs of a vision transformer can be decreased by reducing the number of tokens in the network, there is no setting that is optimal for all input images. In this work, we therefore introduce a differentiable parameter-free Adaptive Token Sampler (ATS) module, which can be plugged into any existing vision transformer architecture. ATS empowers vision transformers by scoring and adaptively sampling significant tokens. As a result, the number of tokens is not constant anymore and varies for each input image. By integrating ATS as an additional layer within the current transformer blocks, we can convert them into much more efficient vision transformers with an adaptive number of tokens. Since ATS is a parameter-free module, it can be added to the off-the-shelf pre-trained vision transformers as a plug and play module, thus reducing their GFLOPs without any additional training. Moreover, due to its differentiable design, one can also train a vision transformer equipped with ATS. We evaluate the efficiency of our module in both image and video classification tasks by adding it to multiple SOTA vision transformers. Our proposed module improves the SOTA by reducing their computational costs (GFLOPs) by 2X, while preserving their accuracy on the ImageNet, Kinetics-400, and Kinetics-600 datasets.

研究动机与目标

• 为解决视觉变换器计算成本高的问题，其计算复杂度随序列长度呈平方级增长。
• 实现基于内容的动态、自适应的标记剪枝，根据输入图像不同而变化，避免固定比例的剪枝方法。
• 设计一种无参数、可微的模块，可无缝集成到预训练模型中，无需微调或重新训练。
• 在图像和视频分类基准测试中，减少 GFLOPs 的同时保持或提升准确率。
• 为边缘设备提供无需架构微调的即插即用高效推理解决方案。

提出的方法

• ATS 使用分类标记的注意力权重，为所有输入标记计算显著性得分。
• 对这些得分的累积分布应用逆变换采样，以选择一组显著的标记。
• 通过注意力权重与值的加权组合，对所选标记进行软下采样，以保留特征表示。
• 该模块具有可微性且无参数，可无缝集成到任意视觉变换器模块中，无需额外训练。
• 可作为即插即用层插入现有预训练模型中，或与主干网络联合训练。
• 该方法支持在多个阶段逐步剪枝，每张图像的标记数量根据其复杂度动态变化。

实验结果

研究问题

• RQ1 能否设计一种无参数、可微的模块，根据输入内容动态减少视觉变换器中的标记数量？
• RQ2 与固定比例剪枝相比，自适应标记采样是否能在降低计算成本的同时保持准确率？
• RQ3 ATS 是否可作为即插即用模块应用于预训练视觉变换器，而无需微调？
• RQ4 在多个阶段中逐步进行标记采样，对效率和准确率有何影响？
• RQ5 基于注意力的显著性评分是否在效率和泛化能力方面优于可学习的评分网络？

主要发现

• ATS 在 DeiT-S 及其他 SOTA 视觉变换器中，将 GFLOPs 最多减少 2 倍，同时在 ImageNet 上保持或提升 top-1 准确率。
• DeiT-S+ATS 模型在仅 2.9 GFLOPs 的情况下达到 79.7% 的 top-1 准确率，优于相同 FLOP 水平下的基线模型 DeiT-S（78.9%）。
• 将 ATS 集成到后期阶段（如第 3 阶段）可获得更高的准确率（78.5%），而早期阶段（73.1%）的准确率较低，原因在于后期注意力得分更可靠。
• 多阶段集成 ATS（第 3 至第 11 阶段）在 2.6 GFLOPs 下实现 79.2% 的 top-1 准确率，证明了逐步提升效率的潜力。
• 当应用于 EViT-DeiT-S 时，ATS 进一步将 GFLOPs 从 3.0 降低至 2.5，且准确率无下降，证明其兼容性与优于静态剪枝的优越性。
• 可视化结果表明，模型对复杂、杂乱的图像采样更多标记，对简单图像采样更少，且对显著区域的关注更集中。