Stochastic Gradient Estimator for Differentiable NAS
Abstract
Neural architecture search (NAS) has recently attracted more attention due to its ability to design deep neural networks automatically.
Differentiable NAS methods have predominated due to their search eﬀiciency.
However, differentiable NAS methods consistently adopt approximate gradient-based methods to solve bilevel optimization problems.
While second derivative approximation optimizes Jacobian or/and Hessian vector computation, it is imprecise and time-consuming in practice.
In this paper, we revisit the hypergradient of bilevel optimization problems in NAS, then propose a new optimizer based on a stochastic gradient estimator(SGE) for the computation of the Jacobian matrix in the hypergradient.
The SGE is adaptable to previous differentiable NAS methods and more accurate than previously direct gradient approximation method.
In the experiments on commonly differentiable NAS benchmarks, the proposed SGE-NAS algorithm outperforms the baseline algorithm.
The test result demonstrates that the proposed SGE-NAS can effectively reduce search time and find the model with higher classification performance.
1. Introduction
Neural Architecture Search (NAS) is an effective method on automating the process of neural network design, with achieving remarkable success on various deep learning applications [1, 4, 10, 12].
A recently proposed method Differentiable ARchiTecture Search (DARTS)
[6] adopts the continuous relaxation to convert the operation selection problem into the continuous magnitude optimization for a set of candidate operations which makes it possible solve the bilevel optimization problem in NAS via continuous method [8]
Gradient-based methods are widely used in bilevel optimization problems for NAS [2, 6, 14, 15].
Hypergradient in bilevel optimization is the gradient of the outer variable and depends on the minimal solution of the inner function [8].
Previous differentiable NAS methods [6] commonly replace hypergradient by a simple approximation scheme.
The architecture parameters α and network weights parameters w are alternately updated on the validation and training datasets via gradient descent, with approximating either the previous step forward or the current w ∗ .
Only two forward passes to the weights and two backward passes to α are needed to compute the hypergradient, thus the complexity is partially reduced [6].
It also lays a hidden trouble for the common performance collapse problem of differentiable NAS [15].
However, due to the high dimension of the neural network parameter α , gradient-based methods are still computationally expensive in practice after approximate processing.
Gradient Descent
Algorithm
Estimator
Black Box
Algorithm
SGE-NAS * * , ( ) 1 1 ( ( ), ) ( , ; ) , ( ; D
D , ) Q N
N val k k val k k
F j val k k F k j w j w
Y
Y u Q        =  =  +   ( ) ( ) ( ) ;
N
N N w u w u u      + − =
* * * * * ( ) ( ( ), ) =
( ( ), ) ( ) ( ( ), ) val val val w w w w            +  * ( ( ), ) val w     * * ( ) ( ( ), ) val w w    
To overcome the computational challenges of current gradient-based methods, we employ an estimation approach to optimize the second-order computation￿ namely stochastic gradient estimator(SGE).
An ES(evolution strategy)-based bilevel optimizer is proposed for hyperparameter optimization [3], which approximates the hypergradient, treating it completely as a black-box objective.
However, bilevel optimization as Eq.
3 has a more complex form than the gradient in the standard minimization problem. SGE in our method is better combined with Eq. 3 and only estimates the Jacobian matrix, thus alleviate the error of the two-step gradient approximation and has better performance than the method that uses ES estimation. Overall, our contributions to this work are summarized as followed:
In this paper, SGE-NAS is proposed, a novel differentiable NAS combined with estimator and gradient descent.
SGE-NAS replaces the previous two-step approximation algorithm and improves the search accuracy.
A simple but effective stochastic gradient estimator(SGE) is introduced to optimize the complex Jacobian matrix calculation, which can be easily adapted with previous gradientbased methods to reduce the search cost.
2. Methods 2.1. Hypergradient of NAS
Differentiable NAS stacks the entire neural network by searching the structure of nomal cells and reduction cells [6].
Typically, a cell is defined as a Directed Acyclic Graph (DAG) with N nodes, where each node represents a potential connection and the information between every two nodes is transformed by an edge [13].
Each edge ( i, j ) contains multiple candidate operations, and the discrete operation selection is transformed into a differentiable parameter optimization problem by applying continuous relaxation through a learnable architecture parameter set α to mix the outputs of different operations. ¯ O ( i,j
[Equation]
( x ) =
|O|
X
k = β ( i,j ) k O k ( x ) β ( i,j ) k = exp (cid:16) α
( i,j ) k (cid:17)
P |O| k ′
=1 exp (cid:16) α ( i,j ) k ′ (cid:17) (1) where x and ¯ O are the input and mixed output of an edge, O is the candidate operation set, and β denotes the softmax-activated architecture parameter set.
In this way, we can perform architecture search in a differentiable manner by solving following bilevel optimization objective.
[Equation]
Since the optimization of the inner variable α is too expensive, it is impossible to accurately evaluate the architectural gradient.
In the common-used gradient-based bilevel optimizer, the key step is the estimation of the hypergradient.
Different from the approximation processing step widely used in gradient-based NAS[9].
The hypergradient of the optimization objective in Eq.
2 is expanded and analyzed in the following form
[Equation]
where the Jacobian matrix J ∗ ( α ) =
∂w ∗ ( α ) ∂α including the Hessian inverse and the secondorder mixed derivative, and expansed as
[Equation]
[Equation]
Eq.
3 consists of two parts, direct gradient ∇ α L val ( w ∗ ( α ) , α ) and indirect gradient J ∗ ( α ) ⊤ ∇ w ∗ ( α ) L val ( w ∗ ( α ) , α ) .
The direct component can be eﬀiciently computed using existing automatic differentiation techniques from deep learning.
However, the indirect component is much more computationally complex.
2.2. Stochastic gradient estimator
In this section, we propose a new estimation method to solve the computational problem of hypergradient.
An estimation method [3] takes the objective function as a black-box objective and directly applies the ES (Evolution Strategy) method to estimate the hypergradient of the objective function.
Because the hypergradient form is complex and very sensitive, such estimation may lead to a large bias error [5, 7].
Instead of this approach, the gradient ∇ α L val ( w ∗ ( α ) , α ) and ∇ w ∗ ( α ) L val ( w ∗ ( α ) , α ) can be computed eﬀiciently and accurately in
Eq. 3, while the Jacobian matrix J ∗ ( α ) is the main term causing computational diﬀiculties.
We subsume the Hessian matrix into J ∗ ( α ) , then use the stochastic gradient estimator (SGE) to form a Hessian-free method without two-step approximation, which makes the calculation of the hypergradient more accurate and time-saving.
The overall bilevel optimization process is shown in Algorithm 1.
Specifically, our SGE-NAS optimizer has the following two steps:
Firstly, for the inner variable α : The amplitude is set to α k + µu k,j , and all SGD processes are sampled in the same size batch { S 0 , ..., S N − 1 } .
Inspired by [11], b J N ( α k ; u k,j ) can be calculated as
[Equation]
where u is an independent and identically distributed Gaussian vector.
For any vector v , the vector product of the Jacobian can be eﬀiciently computed using the dot product as b J N ( α ; u ) ⊤ v = ⟨ δ ( α ; u ) , v ⟩ u , where δ ( α ; u ) =
y N ( α + µu ) − w N ( α ) µ .
Secondly, for the inner variable w ∗ ( α ) : Sample a new batch D F independent from { S 0 , ..., S N − 1 } , where sampled as an estimator of the stochastic gradient ∇ α L val ( α k , Y Nk ; D F ) and ∇ w ∗ ( α ) L val ( α k , Y Nk ; D F ) .
The hypergradient is then estimated as
[Equation]
Data: inner and outer variables α and w ∗ ( α ) , initializations α 0 and w 0 ∗ ( α ) , lower- and upper-level stepsizes p and q , inner and outer iterations numbers K and N , and number of Gaussian vectors Q .
Result: updated inner variable α k +1 , and hypergradient of bilevel optimization ∇L val ( w ∗ ( α ) , α ) initialization; while k = 0 , 1 , ..., K do set Y 0 k = w 0 ∗ ( α ) ,
Y[-
-]
[Equation]
=
[Equation] j = [Equation] ., Q ; generate u k,j = N (0 , I )
∈ R p , j = 1 , ..., Q ; while t = 1 , 2 , ..., N do draw a sample batch S t − 1 ; update Y tk =
{+Y t − 1 k − p ∇ w ∗ ( α ) L train ( α k , Y t − 1 k ; S t − 1 ) ; while j = 1 , 2 , ..., Q do update Y tk,j = Y t − 1 k,j − p ∇ w ∗ ( α ) L train (cid:16) α k + µu k,j , Y t − 1 k,j ; S t −(cid:17) ; end end calculate δ j = Y Nk,j − Y Nk µ , j = 1 , ..., Q ; draw a sample batch D F ; calculate ∇L val ( w ∗ ( α k ) , α k ) via Eq. 6 ; update α k +
+}
=
α k − q ∇L val ( w ∗ ( α ) , α ) ; end
[Table]
3. Experiment
In this section [Table] we conduct extensive experiments applying our SGE-NAS optimization algorithm based on the following differentiable NAS with SGD-Methods:(1)
SGE-NAS-a:
DARTS [6], (2) SGE-NAS-b: PC-DARTS [14],(3) SGE-NAS-c: Fair DARTS [2], (4) SGE
NAS-d: β -DARTS [15].
Image benchmark datasets CIFAR-10 and CIFAR-100 are used in the experiments.
The search space for the candidate operations includes: 3×3 and 5×5 separable convolutions, 3×3 and 5×5 dilated separable convolutions, 3×3 max pooling, 3×3 average pooling
[Table]
and skip connection.
The hyperparameters for the architecture are kept the same as DARTS,
PC-DARTS, Fair DARTS, and β -DARTS.
SGD optimizer is set with a momentum of 0.
and a weight decay of 3e-4, and the initial learning rate is set to 0.005 with a cosine decay scheduler.
For our SGE optimizer, we set lower- and upper-level stepsizes p =0.01, q =0.01, number of Gaussian vectors Q =10 in Algorithm 1 for searching the best architecture.
Table 1 shows that all SGE-NAS series models achieve comparable or better performance than the baselines in terms of top-1 accuracy.
Among them, models with comparable accuracy on CIFAR-10 are obtained while our search cost is 34.3 % GPU hours less than β -DARTS [15].
Besides, there is an improvement of 1 . 05% ∼ 3 .
50% of top-1 accuracy on
CIFAR-100 with approximately 20% ∼ 34% search time reduction.
After easily adapting to the baselines, our SGE-NAS shows its advantage in reducing search time and improving model accuracy compared to previous gradient-based methods.
4. Conclusion
In this work, we propose SGE-NAS, a novel Neural Architecture Search optimization algorithm based on stochastic gradient estimation, to improve the search speed and reduce the approximation error.
Specifically, we analyze the hypergradient of bilevel optimization in
NAS and propose a stochastic gradient estimator for the Jacobian matrix to relieve the most time-consuming part.
The proposed estimator is combined with the gradient descent of the derivative process to form a new method to solve the NAS bilevel optimization problem.
Extensive experimental results on several benchmark datasets demonstrate the effectiveness and eﬀiciency of the SGE-NAS algorithm, which keeps low search costs and achieves better search accuracy performance.