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

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

Advanced Neural Network Applications参考文献 32被引用 8

一句话总结

本文提出自适应标记采样（ATS），这是一种无参数、可微分的模块，可根据输入图像的复杂度动态减少视觉变换器中的标记数量。通过为每张图像评分并采样关键标记，ATS在保持准确率的同时，使ImageNet上的GFLOPs降低37%，为现有模型提供即插即用的效率提升。

ABSTRACT

While state-of-the-art vision transformer models achieve promising results for image classification, they are computationally very 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 Sampling (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 anymore static but it varies for each input image. By integrating ATS as an additional layer within 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 off-the-shelf pretrained vision transformers as a plug-and-play module, thus reducing their GFLOPs without any additional training. However, due to its differentiable design, one can also train a vision transformer equipped with ATS. We evaluate our module on the ImageNet dataset by adding it to multiple state-of-the-art vision transformers. Our evaluations show that the proposed module improves the state-of-the-art by reducing the computational cost (GFLOPs) by 37% while preserving the accuracy.

研究动机与目标

• 为解决视觉变换器推理所需的高计算成本（需大量GFLOPs）的问题。
• 克服视觉变换器中固定标记数量的局限性，该局限性对多样化输入图像而言并非最优。
• 开发一种无参数、可微分的模块，实现在每张输入图像上自适应地减少标记数量。
• 允许无缝集成到现有视觉变换器中，无需微调，同时仍支持端到端微调。
• 在不牺牲分类准确率的前提下，实现显著的计算效率提升。

提出的方法

• ATS引入一种可微分的评分机制，用于对分类标记注意力图中各标记的重要性进行排序。
• 该模块根据得分自适应地仅采样最重要的标记，从而减少每张图像的总标记数。
• ATS作为额外层插入标准变换器块中，取代原有的固定标记处理流程。
• 该模块无参数，可直接部署于预训练视觉变换器上，无需微调。
• 由于其可微分特性，ATS支持端到端训练，从而可对标记采样过程进行优化。
• 该方法兼容任何视觉变换器架构，可作为即插即用模块应用。

实验结果

研究问题

• RQ1 能否设计一种无参数、可微分的模块，根据输入图像复杂度动态减少视觉变换器中的标记数量？
• RQ2 自适应标记采样是否能在不造成准确率下降的前提下实现显著的计算节省？
• RQ3 所提出的模块能否在不重新训练的前提下直接部署于现有预训练视觉变换器上？
• RQ4 与原始模型相比，采用ATS的视觉变换器在GFLOPs和准确率方面的性能表现如何？
• RQ5 对包含ATS的完整模型进行端到端训练，能否进一步提升效率和准确率？

主要发现

• 在ImageNet数据集上，ATS使视觉变换器的计算成本降低了37%（以GFLOPs衡量）。
• 配备ATS的视觉变换器的准确率与原始模型相比保持不变。
• 该模块可直接作为即插即用组件应用于现有预训练视觉变换器，无需任何额外训练。
• ATS的可微分设计支持端到端微调，从而可进一步优化采样过程。
• 自适应标记采样机制成功识别并保留了每张图像中最具有信息量的标记，显著提升了效率。
• 该方法在显著降低推理复杂度的同时，仍保持了最先进性能。