[论文解读] AdaViT: Adaptive Vision Transformers for Efficient Image Recognition

Lingchen Meng, Hengduo Li+5 more|arXiv (Cornell University)|Nov 30, 2021

Advanced Neural Network Applications参考文献 49被引用 15

一句话总结

AdaViT 提出了一种用于视觉变换器的自适应计算框架，通过根据输入图像动态选择图像块、自注意力头和变换器块，以提升推理效率。通过使用 Gumbel-Softmax 联合训练轻量级决策网络与主干网络，AdaViT 在 ImageNet 上仅损失 0.8% 准确率的情况下，实现比最先进模型快 2 倍以上的推理速度。

ABSTRACT

Built on top of self-attention mechanisms, vision transformers have demonstrated remarkable performance on a variety of vision tasks recently. While achieving excellent performance, they still require relatively intensive computational cost that scales up drastically as the numbers of patches, self-attention heads and transformer blocks increase. In this paper, we argue that due to the large variations among images, their need for modeling long-range dependencies between patches differ. To this end, we introduce AdaViT, an adaptive computation framework that learns to derive usage policies on which patches, self-attention heads and transformer blocks to use throughout the backbone on a per-input basis, aiming to improve inference efficiency of vision transformers with a minimal drop of accuracy for image recognition. Optimized jointly with a transformer backbone in an end-to-end manner, a light-weight decision network is attached to the backbone to produce decisions on-the-fly. Extensive experiments on ImageNet demonstrate that our method obtains more than 2x improvement on efficiency compared to state-of-the-art vision transformers with only 0.8% drop of accuracy, achieving good efficiency/accuracy trade-offs conditioned on different computational budgets. We further conduct quantitative and qualitative analysis on learned usage polices and provide more insights on the redundancy in vision transformers.

研究动机与目标

• 解决视觉变换器计算成本过高的问题，其计算成本随输入图像块数量和模型深度呈二次方增长。
• 通过识别并利用视觉变换器中的冗余性，认识到复杂图像比简单图像需要更多计算量。
• 开发一种端到端可训练的框架，学习针对不同输入实例的图像块、注意力头和变换器块的使用策略。
• 通过根据输入复杂度动态分配计算资源，实现在不损失准确率的前提下提升推理效率。
• 提供一种灵活且与模型无关的解决方案，可适用于多种视觉变换器主干网络，且仅需极少的架构修改。

提出的方法

• 在每个变换器块上附加一个轻量级决策网络，用于预测图像块、注意力头和模块的二值使用决策。
• 在训练过程中使用 Gumbel-Softmax 实现离散决策的可微采样，从而支持端到端优化。
• 通过结合交叉熵分类损失和使用成本损失的复合损失函数，联合优化决策网络与视觉变换器主干网络。
• 通过超参数 γ ∈ (0,1] 控制目标计算预算，表示训练期间允许的完整模型计算成本的百分比。
• 实现基于输入的自适应性：简单图像使用更少的图像块/注意力头/模块；复杂图像则使用更多，依据学习到的策略。
• 设计框架时保持与模型无关，可无缝集成到多种视觉变换器架构中，如 DeiT 和 T2T-ViT。

实验结果

研究问题

• RQ1 视觉变换器能否根据输入图像的复杂度动态调整其计算负载？
• RQ2 在不造成显著准确率下降的前提下，可被剪枝的冗余图像块、注意力头和变换器模块的程度有多大？
• RQ3 所学习的自适应计算策略在不同图像类别和复杂度等级下如何变化？
• RQ4 一个端到端可训练的轻量级决策网络能否有效学习以可微方式分配计算资源？
• RQ5 在视觉变换器中使用自适应计算时，推理效率与准确率之间的权衡关系如何？

主要发现

• AdaViT 在 ImageNet 上实现的推理效率比最先进视觉变换器高出 2 倍以上，且仅损失 0.8% 准确率。
• 该方法学会将更多计算资源分配给复杂、杂乱的图像（如“barber shop”、“toyshop”），而为简单、以物体为中心的图像（如“parachute”、“kite”）分配更少计算。
• 图像块选择在深层网络中呈现渐进式计算减少，表明主干网络后端阶段冗余性更高。
• 决策网络学会在后期变换器块中保留更多计算资源，这些模块对最终预测更为关键。
• 定性分析表明，模型在深层网络中聚焦于具有判别性的图像区域，随着网络加深，激活的图像块数量减少。
• AdaViT 具备良好的跨主干泛化能力：当应用于 DeiT-small 和 T2T-ViT 时，其效率与准确率的权衡表现优于标准模型和随机基线变体。