FlatNAS: optimizing Flatness in Neural Architecture Search for Out-of-Distribution Robustness

TL;DR

FlatNAS introduces a loss-flatness aware NAS framework that optimizes in-distribution accuracy, OOD robustness to input corruptions, and parameter count under a constrained design space. By combining a new robustness-aware figure of merit with SAM and an NSGA-II search strategy over an OFA-based supernet, FlatNAS discovers architectures that generalize better to corrupted data while remaining parameter-efficient. The method demonstrates improved mean corruption error on CIFAR-10-C and CIFAR-100-C compared with CNAS while maintaining similar accuracy and computational cost, with α≈0.5 providing the best trade-off. Overall, FlatNAS shows that incorporating loss-landscape flatness into NAS can yield intrinsically more robust architectures suitable for real-world deployment under distribution shifts.

Abstract

Neural Architecture Search (NAS) paves the way for the automatic definition of Neural Network (NN) architectures, attracting increasing research attention and offering solutions in various scenarios. This study introduces a novel NAS solution, called Flat Neural Architecture Search (FlatNAS), which explores the interplay between a novel figure of merit based on robustness to weight perturbations and single NN optimization with Sharpness-Aware Minimization (SAM). FlatNAS is the first work in the literature to systematically explore flat regions in the loss landscape of NNs in a NAS procedure, while jointly optimizing their performance on in-distribution data, their out-of-distribution (OOD) robustness, and constraining the number of parameters in their architecture. Differently from current studies primarily concentrating on OOD algorithms, FlatNAS successfully evaluates the impact of NN architectures on OOD robustness, a crucial aspect in real-world applications of machine and deep learning. FlatNAS achieves a good trade-off between performance, OOD generalization, and the number of parameters, by using only in-distribution data in the NAS exploration. The OOD robustness of the NAS-designed models is evaluated by focusing on robustness to input data corruptions, using popular benchmark datasets in the literature.

Figures (5)

• Figure 1: The proposed FlatNAS framework, which is composed of a Search Space, Evaluator, and a Search Strategy module.
• Figure 2: Image corruptions from the CIFAR-10-C dataset. The figure shows the 15 noise types with level of severity set to 3 with image resolution 64.
• Figure 3: Comparison of mean NN corruption errors on corrupted datasets in function of the perturbation intensity, for FlatNAS at different $\alpha$ values and CNAS. (Upper panel: NNs trained on CIFAR-10 and evaluated on CIFAR-10-C; lower panel: same as upper panel but for CIFAR-100.)
• Figure 4: Final generalization profiles on CIFAR-10-C of NNs trained on CIFAR-10, in function of the perturbation intensity and for each corruption type. Dotted lines correspond to the best NN model found by CNAS and trained with SGD, dashed lines correspond to the latter NN model trained with SAM, and full lines correspond to the best model found by FlatNAS (the dot symbol corresponds to $\alpha=0.1$, the square to $\alpha=0.5$ and the triangle to $\alpha=0.9$).
• Figure 5: Same as Fig. \ref{['fig:OOD10']} but for CIFAR-100-C.