HEP-NAS: Towards Efficient Few-shot Neural Architecture Search via Hierarchical Edge Partitioning

TL;DR

This work targets the accuracy gap in neural architecture search by addressing co-adaptation in one-shot methods and inefficiencies in edge-wise few-shot partitioning. It introduces HEP-NAS, a hierarchy-wise edge partitioning algorithm that groups edges by their shared end node and searches for optimal operation combinations within each hierarchy, coupled with gradual search-space reduction and a novel search-space mutual distillation (SMD) mechanism to stabilize training. The approach yields state-of-the-art results across DARTS and NAS-Bench-201 benchmarks, achieving CIFAR-10 accuracy of 97.56% and ImageNet transfer accuracy of 76.4%, while reducing search cost relative to competing methods. Overall, HEP-NAS offers a scalable, accurate NAS framework that mitigates co-adaptation, accelerates convergence, and generalizes well across datasets and search spaces.

Abstract

One-shot methods have significantly advanced the field of neural architecture search (NAS) by adopting weight-sharing strategy to reduce search costs. However, the accuracy of performance estimation can be compromised by co-adaptation. Few-shot methods divide the entire supernet into individual sub-supernets by splitting edge by edge to alleviate this issue, yet neglect relationships among edges and result in performance degradation on huge search space. In this paper, we introduce HEP-NAS, a hierarchy-wise partition algorithm designed to further enhance accuracy. To begin with, HEP-NAS treats edges sharing the same end node as a hierarchy, permuting and splitting edges within the same hierarchy to directly search for the optimal operation combination for each intermediate node. This approach aligns more closely with the ultimate goal of NAS. Furthermore, HEP-NAS selects the most promising sub-supernet after each segmentation, progressively narrowing the search space in which the optimal architecture may exist. To improve performance evaluation of sub-supernets, HEP-NAS employs search space mutual distillation, stabilizing the training process and accelerating the convergence of each individual sub-supernet. Within a given budget, HEP-NAS enables the splitting of all edges and gradually searches for architectures with higher accuracy. Experimental results across various datasets and search spaces demonstrate the superiority of HEP-NAS compared to state-of-the-art methods.

Figures (7)

• Figure 1: Overall illustration of HEP-NAS. A hierarchy-wise partition strategy is utilized to create sub-supernets. Subsequently, SMD is employed to expedite and stabilize the training process of these sub-supernets. The next step involves selecting the sub-supernet with the highest accuracy on the validation dataset, replacing the previous optimal one, and proceeding to split the next hierarchy until all hierarchies have been partitioned.
• Figure 2: Performance comparation of HEP-NAS with various NAS methods on ImageNet.
• Figure 3: Remaining search space size and accuracy distribution of sub-supernets after each segmentation stage.
• Figure 4: Normal cell on CIFAR-10.
• Figure 5: Reduction cell on CIFAR-10.