ProtoNAS: Accelerating Differentiable Neural Architecture Search via Prototypical Data Pruning

Abstract

In this paper, we take a close look at the efficient differentiable Neural Architecture Search (NAS) problem through the lens of data pruning. A key insight is that the core samples in the pruned dataset can be interpreted as support vectors for networks to make informed decisions. Therefore, we introduce ProtoNAS, an innovative prototypical data pruning method that dynamically and explicitly measures the prototypical distance of each data sample to its class-specific clustering centroid (i.e., prototype). These prototypes are dynamically updated on-the-fly throughout the search phase, enabling flexible data pruning that adjusts with the architectures and network parameters. Compared to conventional cross-entropy-based or Euclidean-based distances, the prototypical distance not only considers the distance between a sample and its class-specific prototype, but also the relative distances to the prototypes of other classes in the embedding space, facilitating more precise score assignment. Furthermore, beyond sample pruning, we introduce a class pruning policy aimed at masking the classes with the lowest average scores. This additional policy is designed to further expedite the NAS without compromising the performance of the final searched architecture. Extensive experiments conducted on NAS-Bench-201 and NAS-Bench-1Shot1 search spaces demonstrate that ProtoNAS can achieve a significant 90% pruning rate with minimal performance trade-offs or even improved performance when applied to DARTS and its variants, such as PC-DARTS and $\beta$-DARTS. Compared to State-Of-The-Art (SOTA) data pruning techniques, ProtoNAS achieves a superior balance between search performance and search cost.