[论文解读] Searching for Efficient Multi-Stage Vision Transformers

Yi-Lun Liao, Sertaç Karaman+1 more|arXiv (Cornell University)|Sep 1, 2021

CCD and CMOS Imaging Sensors参考文献 66被引用 12

一句话总结

该论文提出ViT-ResNAS，一种高效的多阶段视觉Transformer架构，通过引入残差空间下采样机制，在各阶段减少序列长度并提升特征通道维度，同时结合一种新颖的基于多架构采样的权重共享神经架构搜索（NAS）方法。该方法在ImageNet上实现了SOTA的准确率-MACs与准确率-吞吐量权衡，优于DeiT、PVT和PiT基线模型，且FLOPs显著更低、吞吐量更高。

ABSTRACT

Vision Transformer (ViT) demonstrates that Transformer for natural language processing can be applied to computer vision tasks and result in comparable performance to convolutional neural networks (CNN), which have been studied and adopted in computer vision for years. This naturally raises the question of how the performance of ViT can be advanced with design techniques of CNN. To this end, we propose to incorporate two techniques and present ViT-ResNAS, an efficient multi-stage ViT architecture designed with neural architecture search (NAS). First, we propose residual spatial reduction to decrease sequence lengths for deeper layers and utilize a multi-stage architecture. When reducing lengths, we add skip connections to improve performance and stabilize training deeper networks. Second, we propose weight-sharing NAS with multi-architectural sampling. We enlarge a network and utilize its sub-networks to define a search space. A super-network covering all sub-networks is then trained for fast evaluation of their performance. To efficiently train the super-network, we propose to sample and train multiple sub-networks with one forward-backward pass. After that, evolutionary search is performed to discover high-performance network architectures. Experiments on ImageNet demonstrate that ViT-ResNAS achieves better accuracy-MACs and accuracy-throughput trade-offs than the original DeiT and other strong baselines of ViT. Code is available at https://github.com/yilunliao/vit-search.

研究动机与目标

• 通过借鉴卷积神经网络（CNNs）的设计原则，特别是空间下采样与多阶段架构，提升视觉Transformer（ViT）的效率与性能。
• 通过在空间下采样过程中引入残差连接，解决深层ViT模型中计算成本高与训练不稳定的挑战。
• 开发一种针对多阶段ViT的更高效神经架构搜索（NAS）框架，实现快速且准确的架构发现。
• 在ImageNet基准上，实现优于现有ViT模型（包括DeiT、PVT和PiT）的准确率-计算量权衡。

提出的方法

• 提出残差空间下采样：一种多阶段ViT设计，其中每个阶段减少空间分辨率并增加通道数，同时通过跳跃连接稳定训练并提升性能。
• 引入基于超网络的NAS框架，采用多架构采样策略，即在单次前向-反向传播中同时采样并训练多个子网络，以加速超网络训练。
• 通过在基础ViT-Res架构上扩展深度与宽度构建搜索空间，实现对多样化子网络的高效探索。
• 在训练好的超网络上采用进化搜索，发现高性能架构而无需重新训练。
• 在搜索空间的所有子网络中共享权重，通过超网络的参数实现性能估计，无需对每个子网络单独训练。
• 利用架构的批量无关特性，提升NAS流程中的训练效率与搜索质量。

实验结果

研究问题

• RQ1 残差空间下采样是否能通过实现类似CNN的多阶段架构，提升视觉Transformer的准确率与效率？
• RQ2 如何使神经架构搜索在计算成本更低的前提下，更高效且更有效地应用于多阶段视觉Transformer？
• RQ3 与标准NAS方法相比，基于多架构采样的权重共享NAS是否能在发现高性能ViT架构方面表现更优？
• RQ4 与当前SOTA ViT模型相比，所提方法在准确率-MACs与准确率-吞吐量权衡方面提升程度如何？
• RQ5 在空间下采样过程中集成残差连接是否能稳定深层ViT模型的训练并提升其性能？

主要发现

• ViT-ResNAS-Tiny在仅多消耗0.5G MACs且吞吐量略低的情况下，Top-1准确率比DeiT-Tiny高出8.6%。
• ViT-ResNAS-Small在单张Titan RTX GPU上，以6.3倍更少的MACs与3.7倍更高的吞吐量，达到与DeiT-Base相当的准确率。
• 与PiT-S相比，ViT-ResNAS-Tiny在相似准确率下，MACs减少1.6倍，吞吐量提升1.6倍。
• ViT-ResNAS-Medium在MACs减少2.8倍、吞吐量提升2.4倍的情况下，准确率比PiT-B高出0.4%。
• 在低MACs场景（约2.0G）下，ViT-ResNAS-Tiny在准确率-MACs权衡上优于ViL-Tiny；在高MACs场景（约5.0G）下，其性能与ViL-Small相当。
• 所提出的基于多架构采样的权重共享NAS方法，实现了更快的超网络训练，并生成了优于标准NAS方法的高性能架构。