[论文解读] Multi-Exit Vision Transformer for Dynamic Inference

Arian Bakhtiarnia, Qi Zhang+1 more|arXiv (Cornell University)|Jun 29, 2021

Advanced Neural Network Applications被引用 8

一句话总结

本文提出七种新型的视觉变换器（Vision Transformers）早停分支架构，以实现动态推理，使简单样本能够获得更快的预测结果，同时保持高精度。在图像分类与回归任务上的大量实验表明，ViT-EE 和 MLP-EE 分支在端到端训练下实现了速度与精度的最佳平衡，其中 ViT-EE 在 CIFAR-100 上的早停准确率达到 97.33%，相比完整推理，FLOPS 最多降低 67%。

ABSTRACT

Deep neural networks can be converted to multi-exit architectures by inserting early exit branches after some of their intermediate layers. This allows their inference process to become dynamic, which is useful for time critical IoT applications with stringent latency requirements, but with time-variant communication and computation resources. In particular, in edge computing systems and IoT networks where the exact computation time budget is variable and not known beforehand. Vision Transformer is a recently proposed architecture which has since found many applications across various domains of computer vision. In this work, we propose seven different architectures for early exit branches that can be used for dynamic inference in Vision Transformer backbones. Through extensive experiments involving both classification and regression problems, we show that each one of our proposed architectures could prove useful in the trade-off between accuracy and speed.

研究动机与目标

• 为延迟敏感的边缘和物联网应用提供视觉变换器的动态推理能力，支持可变的计算预算。
• 设计并评估多种早停分支架构，在降低推理时间的同时保持性能。
• 研究架构选择（如注意力机制、归一化方式以及与主干层的相似性）对早停性能的影响。
• 比较三种训练策略（端到端、逐层、分类器级）在多出口 ViT 中的表现，识别最优配置。

提出的方法

• 提出七种不同的早停分支架构：MLP-EE、CNN-Ignore-EE、CNN-Add-EE、CNN-Project-EE、ViT-EE、MLP-Mixer-EE 和 ResMLP-EE，每种架构具有不同的结构特征。
• 将这些分支集成到视觉变换器主干网络的中间层，实现在完整前向传播完成前的早期预测。
• 采用三种训练策略：端到端（联合优化）、逐层（渐进微调）和分类器级（分支训练期间冻结主干）。
• 使用 FLOPS 和准确率指标，评估不同退出位置和数据集下速度与性能之间的权衡。
• 在图像分类（CIFAR-10、CIFAR-100、Fashion-MNIST）和回归（DISCO 人群计数）任务上评估模型。
• 分析架构影响因素，如感受野、注意力类型、归一化方式以及与主干层的架构相似性对退出性能的影响。

实验结果

研究问题

• RQ1 在视觉变换器中，哪种早停分支架构能实现最佳的精度-速度权衡？
• RQ2 不同的训练策略（端到端、逐层、分类器级）如何影响多出口 ViT 模型的性能？
• RQ3 早停分支与主干层之间的架构相似性在多大程度上能提升早停性能？
• RQ4 注意力机制、归一化方式和感受野大小等架构组件如何影响早停准确率？
• RQ5 多出口 ViT 架构是否能在早期层实现高性能，同时显著降低计算成本？

主要发现

• ViT-EE 分支在所有架构中达到最高准确率，在 CIFAR-100 的最终退出点准确率达到 97.33%，且主干参数退化最小。
• MLP-EE 和 ViT-EE 分支展现出最理想的精度-速度权衡，相比完整推理，FLOPS 最多降低 67%，同时保持高准确率。
• 端到端训练策略优于逐层和分类器级策略，尤其在保持主干性能和实现有效多出口学习方面表现更优。
• 与主干架构差异较大的 CNN 基分支（如 CNN-Ignore-EE）性能较差，而架构相似性更高的分支（如 ViT-EE）表现更优，凸显了架构一致性的关键作用。
• 位于较后层的分支（如 Early@9）准确率高于早期分支，但前提是架构与训练策略高度匹配。
• 逐层训练策略因初始微调后预训练权重丢失，导致后期准确率难以提升，限制了其实际应用价值。