# ShiftNAS: Improving One-shot NAS via Probability Shift

arXiv: https://arxiv.org/abs/2307.08300 · HTML: https://mlsanity.com/paper/2307.08300/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 45 references that are in the map.
- First posted on arXiv 17 Jul 2023.

## Builds on (most cited first, 30 of 30)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks](https://mlsanity.com/paper/1905.11946/) (2019)
- [MobileNetV2: Inverted Residuals and Linear Bottlenecks](https://mlsanity.com/paper/1801.04381/) (2018)
- [Training data-efficient image transformers & distillation through attention](https://mlsanity.com/paper/2012.12877/) (2020)
- [MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications](https://mlsanity.com/paper/1704.04861/) (2017)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Searching for MobileNetV3](https://mlsanity.com/paper/1905.02244/) (2019)
- [DARTS: Differentiable Architecture Search](https://mlsanity.com/paper/1806.09055/) (2018)
- [Neural Architecture Search with Reinforcement Learning](https://mlsanity.com/paper/1611.01578/) (2016)
- [Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNet](https://mlsanity.com/paper/2101.11986/) (2021)
- [Efficient Neural Architecture Search via Parameter Sharing](https://mlsanity.com/paper/1802.03268/) (2018)
- [Once-for-All: Train One Network and Specialize it for Efficient Deployment](https://mlsanity.com/paper/1908.09791/) (2019)
- [FBNet: Hardware-Aware Efficient ConvNet Design via Differentiable Neural Architecture Search](https://mlsanity.com/paper/1812.03443/) (2018)
- [Single Path One-Shot Neural Architecture Search with Uniform Sampling](https://mlsanity.com/paper/1904.00420/) (2019)
- [Bottleneck Transformers for Visual Recognition](https://mlsanity.com/paper/2101.11605/) (2021)
- [Slimmable Neural Networks](https://mlsanity.com/paper/1812.08928/) (2018)
- [Random Search and Reproducibility for Neural Architecture Search](https://mlsanity.com/paper/1902.07638/) (2019)
- [BigNAS: Scaling Up Neural Architecture Search with Big Single-Stage Models](https://mlsanity.com/paper/2003.11142/) (2020)
- [AutoFormer: Searching Transformers for Visual Recognition](https://mlsanity.com/paper/2107.00651/) (2021)
- [FairNAS: Rethinking Evaluation Fairness of Weight Sharing Neural Architecture Search](https://mlsanity.com/paper/1907.01845/) (2019)
- [SegViT: Semantic Segmentation with Plain Vision Transformers](https://mlsanity.com/paper/2210.05844/) (2022)
- [Blockwisely Supervised Neural Architecture Search with Knowledge Distillation](https://mlsanity.com/paper/1911.13053/) (2019)
- [DSNAS: Direct Neural Architecture Search without Parameter Retraining](https://mlsanity.com/paper/2002.09128/) (2020)
- [GreedyNAS: Towards Fast One-Shot NAS with Greedy Supernet](https://mlsanity.com/paper/2003.11236/) (2020)
- [AttentiveNAS: Improving Neural Architecture Search via Attentive Sampling](https://mlsanity.com/paper/2011.09011/) (2020)
- [BossNAS: Exploring Hybrid CNN-transformers with Block-wisely Self-supervised Neural Architecture Search](https://mlsanity.com/paper/2103.12424/) (2021)
- [Network Augmentation for Tiny Deep Learning](https://mlsanity.com/paper/2110.08890/) (2021)
- [FocusFormer: Focusing on What We Need via Architecture Sampler](https://mlsanity.com/paper/2208.10861/) (2022)

## Built on by (most cited first, 4 of 4)

- [Subnet-Aware Dynamic Supernet Training for Neural Architecture Search](https://mlsanity.com/paper/2503.10740/) (2025)
- [Mixed-precision Supernet Training from Vision Foundation Models using Low Rank Adapter](https://mlsanity.com/paper/2403.20080/) (2024)
- [Vision Transformer Neural Architecture Search for Out-of-Distribution Generalization: Benchmark and Insights](https://mlsanity.com/paper/2501.03782/) (2025)
- [GrowTAS: Progressive Expansion from Small to Large Subnets for Efficient ViT Architecture Search](https://mlsanity.com/paper/2512.12296/) (2025)

## Abstract (arXiv)

One-shot Neural architecture search (One-shot NAS) has been proposed as a time-efficient approach to obtain optimal subnet architectures and weights under different complexity cases by training only once. However, the subnet performance obtained by weight sharing is often inferior to the performance achieved by retraining. In this paper, we investigate the performance gap and attribute it to the use of uniform sampling, which is a common approach in supernet training. Uniform sampling concentrates training resources on subnets with intermediate computational resources, which are sampled with high probability. However, subnets with different complexity regions require different optimal training strategies for optimal performance. To address the problem of uniform sampling, we propose ShiftNAS, a method that can adjust the sampling probability based on the complexity of subnets. We achieve this by evaluating the performance variation of subnets with different complexity and designing an architecture generator that can accurately and efficiently provide subnets with the desired complexity. Both the sampling probability and the architecture generator can be trained end-to-end in a gradient-based manner. With ShiftNAS, we can directly obtain the optimal model architecture and parameters for a given computational complexity. We evaluate our approach on multiple visual network models, including convolutional neural networks (CNNs) and vision transformers (ViTs), and demonstrate that ShiftNAS is model-agnostic. Experimental results on ImageNet show that ShiftNAS can improve the performance of one-shot NAS without additional consumption. Source codes are available at https://github.com/bestfleer/ShiftNAS.