[论文解读] AdaViT: Adaptive Tokens for Efficient Vision Transformer

Hongxu Yin, Arash Vahdat+4 more|arXiv (Cornell University)|Dec 14, 2021

CCD and CMOS Imaging Sensors被引用 16

一句话总结

A-ViT 为视觉变换器引入了一种自适应标记计算机制，该机制在推理过程中根据输入图像的复杂度动态减少处理的标记数量。通过重用现有网络参数来学习每个标记的停止概率，并应用分布先验正则化，A-ViT 在仅损失 0.3⁠% 准确率的情况下，使 DeiT-Tiny 的吞吐量提升 62⁠%，DeiT-Small 的吞吐量提升 38⁠%，实现了无需架构修改的高效、硬件原生推理。

ABSTRACT

We introduce A-ViT, a method that adaptively adjusts the inference cost of vision transformer (ViT) for images of different complexity. A-ViT achieves this by automatically reducing the number of tokens in vision transformers that are processed in the network as inference proceeds. We reformulate Adaptive Computation Time (ACT) for this task, extending halting to discard redundant spatial tokens. The appealing architectural properties of vision transformers enables our adaptive token reduction mechanism to speed up inference without modifying the network architecture or inference hardware. We demonstrate that A-ViT requires no extra parameters or sub-network for halting, as we base the learning of adaptive halting on the original network parameters. We further introduce distributional prior regularization that stabilizes training compared to prior ACT approaches. On the image classification task (ImageNet1K), we show that our proposed A-ViT yields high efficacy in filtering informative spatial features and cutting down on the overall compute. The proposed method improves the throughput of DeiT-Tiny by 62% and DeiT-Small by 38% with only 0.3% accuracy drop, outperforming prior art by a large margin. Project page at https://a-vit.github.io/

研究动机与目标

• 在视觉变换器中实现输入相关的推理，根据图像复杂度降低计算成本。
• 设计一种机制，在不增加参数或修改网络架构的前提下，自适应地停止对信息量较少的空间标记的处理。
• 通过一种新颖的分布先验，稳定训练过程，并引导停止深度趋向目标计算预算。
• 在标准硬件上实现显著的吞吐量提升，同时保持模型准确率。

提出的方法

• A-ViT 将自适应计算时间（ACT）扩展至标记级别，使每个标记在满足停止条件后可被丢弃。
• 每个标记的停止概率通过每个变换器块最后全连接层导出的单个神经元计算得出，重用现有参数，无需额外参数或计算。
• 引入分布先验正则化，以引导期望停止深度趋向目标，提升训练稳定性和收敛性。
• 该方法实现了动态、空间自适应的计算，仅对信息丰富的标记进行深度处理，从而降低 FLOPs 和推理时间。
• 停止机制通过 ponder 损失和分布先验进行端到端训练，无需辅助网络或微调。

实验结果

研究问题

• RQ1 在不修改架构的前提下，能否在标记级别实现自适应计算以提升视觉变换器的推理效率？
• RQ2 如何仅利用现有网络参数高效学习停止决策？
• RQ3 分布先验能否稳定训练并改善向目标计算预算的收敛？
• RQ4 自适应标记停止与语义图像内容及人类感知的相关性有多大？

主要发现

• A-ViT 在 ImageNet1K 上仅造成 0.3⁠% 的 top-1 准确率下降，使 DeiT-Tiny 的吞吐量提升 62⁠%，DeiT-Small 的吞吐量提升 38⁠%。
• 与传统 ACT 相比，该方法因具备细粒度、空间自适应的标记剪枝，使 FLOPs 减少 25⁠%。
• 分布先验正则化使模型更快收敛至目标停止深度，并相比仅使用 ponder 损失，带来 6.4⁠% 的准确率提升。
• 仅使用单个向量元素进行停止计算，DeiT-Tiny 仅造成 0.08⁠% 的准确率下降，DeiT-Small 仅造成 0.04⁠% 的下降，验证了极低的参数开销。
• 添加辅助网络用于停止学习仅带来 0.06⁠% 的准确率增益，但导致 12.6⁠% 的吞吐量损失，因此不切实际。