Neural Architecture Search (NAS) has shown promising performance in the automatic design of vision transformers (ViT) exceeding 1G FLOPs. However, designing lightweight and low-latency ViT models for diverse mobile devices remains a big challenge. In this work, we propose ElasticViT, a two-stage NAS approach that trains a high-quality ViT supernet over a very large search space that supports a wide range of mobile devices, and then searches an optimal sub-network (subnet) for direct deployment. However, prior supernet training methods that rely on uniform sampling suffer from the gradient conflict issue: the sampled subnets can have vastly different model sizes (e.g., 50M vs. 2G FLOPs), leading to different optimization directions and inferior performance. To address this challenge, we propose two novel sampling techniques: complexity-aware sampling and performance-aware sampling. Complexity-aware sampling limits the FLOPs difference among the subnets sampled across adjacent training steps, while covering different-sized subnets in the search space. Performance-aware sampling further selects subnets that have good accuracy, which can reduce gradient conflicts and improve supernet quality. Our discovered models, ElasticViT models, achieve top-1 accuracy from 67.2% to 80.0% on ImageNet from 60M to 800M FLOPs without extra retraining, outperforming all prior CNNs and ViTs in terms of accuracy and latency. Our tiny and small models are also the first ViT models that surpass state-of-the-art CNNs with significantly lower latency on mobile devices. For instance, ElasticViT-S1 runs 2.62x faster than EfficientNet-B0 with 0.1% higher accuracy.

翻译：神经架构搜索（NAS）在超过1G FLOPs的视觉变压器（ViT）自动设计中展现出良好性能。然而，为多样化移动设备设计轻量级低延迟ViT模型仍是重大挑战。本文提出ElasticViT——一种两阶段NAS方法：首先在支持广泛移动设备的极大搜索空间上训练高质量ViT超网，随后搜索最优子网络（子网）直接部署。但现有依赖均匀采样的超网训练方法面临梯度冲突问题：采样子网络在模型规模（如50M vs. 2G FLOPs）上差异显著，导致优化方向分歧与性能下降。为应对此挑战，我们提出两种新型采样技术：复杂度感知采样与性能感知采样。复杂度感知采样限制相邻训练步骤间采样子网络的FLOPs差异，同时覆盖搜索空间中不同规模的子网络；性能感知采样则进一步选择具有良好精度的子网络，从而减少梯度冲突并提升超网质量。经本文方法发现的ElasticViT模型，在60M至800M FLOPs的ImageNet数据集上无需额外重训练即可实现67.2%至80.0%的Top-1准确率，在精度与延迟指标上全面超越此前所有CNN与ViT模型。其中，微型与小尺寸模型更是首批以显著更低延迟在移动设备上超越最先进CNN的ViT模型。例如，ElasticViT-S1的推理速度比EfficientNet-B0快2.62倍，同时精度高出0.1%。