STO-DARTS: Stochastic Bilevel Optimization for Differentiable Neural Architecture Search

Differentiable bilevel Neural Architecture Search (NAS) has emerged as a powerful approach in automated machine learning (AutoML) for efficiently searching for neural network architectures. However, the existing differentiable methods encounter challenges, such as the risk of becoming trapped in local optima and the computationally expensive Hessian matrix inverse calculation performed when solving the bilevel NAS optimization model. In this paper, a novel-but-efficient stochastic bilevel optimization approach, called STO-DARTS, is proposed for the bilevel NAS optimization problem. Specifically, we design a hypergradient estimate, which is constructed using stochastic gradient descent from the gradient information contained in the Neumann series. This estimate alleviates the issue of local optima traps, enabling searches for exceptional network architectures. To validate the effectiveness and efficiency of the proposed method, two versions of STO-DARTS with different hypergradient estimators are constructed and experimentally tested on different datasets in NAS-Bench-201 and DARTS search spaces. The experimental results show that the proposed STO-DARTS approach achieves competitive performance with that of other state-of-the-art NAS methods in terms of determining effective network architectures. To support our approach, we also provide theoretical analyses.