Subnet-Aware Dynamic Supernet Training for Neural Architecture Search

2025-03-13

AI Summary

In N-shot neural architecture search (NAS), static supernet training induces unfairness among subnetworks—favoring low-complexity architectures—and amplifies momentum-based gradient noise. To address these issues, we propose a subnet-aware dynamic training paradigm. Our key contributions are: (1) Complexity-Aware Learning Rate scheduling (CaLR), which adaptively modulates learning rates based on subnetwork computational complexity to mitigate optimization bias; and (2) Momentum Separation (MS), which maintains independent momentum buffers for subnetworks of differing complexities to suppress gradient noise propagation. The method seamlessly integrates with weight-sharing supernets and multi-subnet joint training, incurring negligible computational overhead. Extensive experiments on NAS-Bench-201, the MobileNet search space, and CIFAR/ImageNet benchmarks demonstrate significant improvements in search performance. Moreover, our approach is plug-and-play compatible with mainstream N-shot NAS frameworks.

Technology Category

• Search and Optimization: Learning to Search
• Machine Learning: Neuro-Symbolic Learning
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Economics, Online Markets and Human Computation: Incentives in network design for Web infrastructures and ecosystems

Abstract

N-shot neural architecture search (NAS) exploits a supernet containing all candidate subnets for a given search space. The subnets are typically trained with a static training strategy (e.g., using the same learning rate (LR) scheduler and optimizer for all subnets). This, however, does not consider that individual subnets have distinct characteristics, leading to two problems: (1) The supernet training is biased towards the low-complexity subnets (unfairness); (2) the momentum update in the supernet is noisy (noisy momentum). We present a dynamic supernet training technique to address these problems by adjusting the training strategy adaptive to the subnets. Specifically, we introduce a complexity-aware LR scheduler (CaLR) that controls the decay ratio of LR adaptive to the complexities of subnets, which alleviates the unfairness problem. We also present a momentum separation technique (MS). It groups the subnets with similar structural characteristics and uses a separate momentum for each group, avoiding the noisy momentum problem. Our approach can be applicable to various N-shot NAS methods with marginal cost, while improving the search performance drastically. We validate the effectiveness of our approach on various search spaces (e.g., NAS-Bench-201, Mobilenet spaces) and datasets (e.g., CIFAR-10/100, ImageNet).

Problem

• Addresses unfairness in supernet training towards low-complexity subnets.
• Reduces noisy momentum updates in supernet training.
• Introduces adaptive training strategies for individual subnet characteristics.

Innovation

• Dynamic supernet training adapts to subnet characteristics
• Complexity-aware LR scheduler adjusts learning rates
• Momentum separation groups subnets to reduce noise