D R NAS: D IRICHLET N EURAL A RCHITECTURE S EARCH
A BSTRACT
This paper proposes a novel differentiable architecture search method by formulating it into a distribution learning problem.
We treat the continuously relaxed architecture mixing weight as random variables, modeled by Dirichlet distribution.
With recently developed pathwise derivatives, the Dirichlet parameters can be easily optimized with gradient-based optimizer in an end-to-end manner.
This formulation improves the generalization ability and induces stochasticity that naturally encourages exploration in the search space.
Furthermore, to alleviate the large memory consumption of differentiable NAS, we propose a simple yet effective progressive learning scheme that enables searching directly on large-scale tasks, eliminating the gap between search and evaluation phases.
Extensive experiments demonstrate the effectiveness of our method.
Speciﬁcally, we obtain a test error of 2.46% for CIFAR-10, 23.7% for ImageNet under the mobile setting.
On NASBench-201, we also achieve state-of-the-art results on all three datasets and provide insights for the effective design of neural architecture search algorithms.
I NTRODUCTION
Recently
Neural Architecture Search (NAS) has attracted lots of attentions for its potential to democratize deep learning.
For a practical end-to-end deep learning platform, NAS plays a crucial role in discovering task-speciﬁc architecture depending on users’ conﬁgurations (e g., dataset, evaluation metric, etc.).
Pioneers in this ﬁeld develop prototypes based on reinforcement learning (Zoph & Le, 2017), evolutionary algorithms (Real et al., 2019) and Bayesian optimization (Liu et al., 2018).
These works usually incur large computation overheads, which make them impractical to use.
More recent algorithms signiﬁcantly reduce the search cost including one-shot methods (Pham et al., 2018; Bender et al., 2018), a continuous relaxation of the space (Liu et al., 2019) and network morphisms (Cai et al., 2018).
In particular, Liu et al.
(2019) proposes a differentiable NAS framework - DARTS, converting the categorical operation selection problem into learning a continuous architecture mixing weight.
They formulate a bi-level optimization objective, allowing the architecture search to be efﬁciently performed by a gradient-based optimizer.
While current differentiable NAS methods achieve encouraging results, they still have shortcomings that hinder their real-world applications.
Firstly, several works have cast doubt on the stability and generalization of these differentiable NAS methods (Chen & Hsieh, 2020; Zela et al., 2020a).
They discover that directly optimizing the architecture mixing weight is prone to overﬁtting the validation set and often leads to distorted structures, e.g., searched architectures dominated by parameter-free operations.
Secondly, there exist disparities between the search and evaluation phases, where proxy tasks are usually employed during search with smaller datasets or shallower and narrower networks, due to the large memory consumption of differentiable NAS.
In this paper, we propose an effective approach that addresses the aforementioned shortcomings named Dirichlet Neural Architecture Search (DrNAS).
Inspired by the fact that directly optimizing the architecture mixing weight is equivalent to performing point estimation (MLE/MAP) from a probabilistic perspective, we formulate the differentiable NAS as a distribution learning problem instead, which naturally induces stochasticity and encourages exploration.
Making use of the probability simplex property of the Dirichlet samples, DrNAS models the architecture mixing weight as random variables sampled from a parameterized Dirichlet distribution.
Optimizing the Dirichlet objective can thus be done efﬁciently in an end-to-end fashion, by employing the pathwise derivative estimators to compute the gradient of the distribution (Martin Jankowiak, 2018).
A straightforward optimization, however, turns out to be problematic due to the uncontrolled variance of the Dirichlet, i.e., too much variance leads to training instability and too little variance suffers from insufﬁcient exploration.
In light of that, we apply an additional distance regularizer directly on the Dirichlet concentration parameter to strike a balance between the exploration and the exploitation.
We further derive a theoretical bound showing that the constrained distributional objective promotes stability and generalization of architecture search by implicitly controlling the Hessian of the validation error.
Furthermore, to enable a direct search on large-scale tasks, we propose a progressive learning scheme, eliminating the gap between the search and evaluation phases.
Based on partial channel connection (Xu et al., 2020), we maintain a task-speciﬁc super-network of the same depth and number of channels as the evaluation phase throughout searching.
To prevent loss of information and instability induced by partial connection, we divide the search phase into multiple stages and progressively increase the channel fraction via network transformation (Chen et al., 2016).
Meanwhile, we prune the operation space according to the learnt distribution to maintain the memory efﬁciency.
We conduct extensive experiments on different datasets and search spaces to demonstrate DrNAS’s effectiveness.
Based on the DARTS search space (Liu et al., 2019), we achieve an average error rate of 2.46% on CIFAR-10, which ranks top amongst NAS methods.
Furthermore, DrNAS achieves superior performance on large-scale tasks such as ImageNet.
It obtains a top-1/5 error of 23.7%/7.1%, surpassing the previous state-of-the-art (24.0%/7.3%) under the mobile setting.
On NAS-Bench201 (Dong & Yang, 2020), we also set new state-of-the-art results on all three datasets with low variance.
Our code is available at https://github.com/xiangning-chen/DrNAS .
T HE P ROPOSED A PPROACH
In this section, we ﬁrst brieﬂy review differentiable NAS setups and generalize the formulation to motivate distribution learning.
We then layout our proposed DrNAS and describe its optimization in section 2.2.
In section 2.3, we provide a generalization result by showing that our method implicitly regularizes the Hessian norm over the architecture parameter.
The progressive architecture learning method that enables direct search is then described in section 2.4.
2. P RELIMINARIES : D IFFERENTIABLE A RCHITECTURE S EARCH
Cell-Based Search Space
The cell-based search space is constructed by replications of normal and reduction cells (Zoph et al., 2018;
Liu et al., 2019).
A normal cell keeps the spatial resolution while a reduction cell halves it but doubles the number of channels.
Every cell is represented by a DAG with
N nodes and E edges, where every node represents a latent representation x i and every edge ( i, j ) is associated with an operations o
( i,j ) (e.g., max pooling or convolution ) selected from a predeﬁned candidate space O .
The output of a node is a summation of all input ﬂows, i.e., x j = (cid:80) i<
j o ( i,j ) ( x i ) , and a concatenation of intermediate node outputs, i.e., concat ( x 2 , ..., x N − 1 ) , composes the cell output, where the ﬁrst two input nodes x 0 and x 1 are ﬁxed to be the outputs of previous two cells.
Gradient-Based Search via Continuous Relaxation To enable gradient-based optimization, Liu et al.
(2019) apply a continuous relaxation to the discrete space.
Concretely, the information passed from node i to node j is computed by a weighted sum of all operations alone the edge, forming a mixed-operation ˆ
o
( i,j ) ( x )
=
(cid:80) o
∈O θ ( i,j ) o o ( x ) .
The operation mixing weight θ ( i,j ) is deﬁned over the probability simplex and its magnitude represents the strength of each operation.
Therefore, the architecture search can be cast as selecting the operation associated with the highest mixing weight for each edge.
To prevent abuse of terminology, we refer to θ as the architecture/operation mixing weight, and concentration parameter β in DrNAS as the architecture parameter throughout the paper.
Bilevel-Optimization with Simplex Constraints With continuous relaxation, the network weight w and operation mixing weight θ can be jointly optimized by solving a constraint bi-level optimization problem:
[Equation]
where the simplex constraint (cid:80) |O| o
=1 θ ( i,j ) o
= 1 can be either solved explicitly via Lagrangian function (Li et al., 2020), or eliminated by substitution method (e.g., θ = Softmax ( α ) , α ∈
R |O|×| E | ) (Liu et al., 2019).
In the next section we describe how this generalized formulation motivates our method.
2. D IFFERENTIABLE A RCHITECTURE S EARCH AS D ISTRIBUTION L EARNING Learning a Distribution over Operation Mixing Weight
Previous differentiable architecture search methods view the operation mixing weight θ as learnable parameters that can be directly optimized (Liu et al., 2019; Xu et al., 2020; Li et al., 2020).
This has been shown to cause θ to overﬁt the validation set and thus induce large generalization error (Zela et al., 2020a;b; Chen &
Hsieh, 2020).
We recognize that this treatment is equivalent to performing point estimation (e.g.,
MLE/MAP) of θ in probabilistic view, which is inherently prone to overﬁtting (Bishop, 2016; Gelman et al., 2004).
Furthermore, directly optimizing θ lacks sufﬁcient exploration in the search space, and thus cause the search algorithm to commit to suboptimal paths in the DAG that converges faster at the beginning but plateaus quickly (Shu et al., 2020).
Based on these insights, we formulate the differentiable architecture search as a distribution learning problem.
The operation mixing weight θ is treated as random variables sampled from a learnable distribution.
Formally, let q ( θ | β ) denote the distribution of θ parameterized by β .
The bi-level objective is then given by:
[Equation]
where d ( · , · ) is a distance function.
Since θ lies on the probability simplex, we select Dirichlet distribution to model its behavior, i.e., q ( θ | β ) ∼ Dir ( β ) , where β represents the Dirichlet concentration parameter.
Dirichlet distribution is a widely used distribution over the probability simplex (Joo et al., 2019; David M. Blei, 2003; Lee et al., 2020; Kessler et al., 2019), and it enjoys nice properties that enables gradient-based training (Martin Jankowiak, 2018).
The concentration parameter β controls the sampling behavior of Dirichlet distribution and is crucial in balancing exploration and exploitation during the search phase.
Let β o denote the concentration parameter assign to operation o .
When β o (cid:28) 1 for most o = 1 ∼ |O| , Dirichlet tends to produce sparse samples with high variance, reducing the training stability; when β o (cid:29) 1 for most o = 1 ∼ |O| , the samples will be dense with low variance, leading to insufﬁcient exploration.
Therefore, we add a penalty term in the objective (2) to regularize the distance between β and the anchor ˆ β = 1 , which corresponds to the symmetric Dirichlet.
In section 2.3, we also derive a theoretical bound showing that our formulation additionally promotes stability and generalization of the architecture search by implicitly regularizing the Hessian of validation loss w.r.t.
architecture parameters.
Learning Dirichlet Parameters via Pathwise Derivative Estimator Optimizing objective (2) with gradient-based methods requires back-propagation through stochastic nodes of Dirichlet samples.
The commonly used reparameterization trick does not apply to Dirichlet distribution, therefore we approximate the gradient of Dirichlet samples via pathwise derivative estimators (Martin Jankowiak, 2018)
[Equation]
where F Beta and f Beta denote the CDF and PDF of beta distribution respectively, δ ij is the indicator function, and β tot is the sum of concentrations.
F Beta is the iregularised incomplete beta function, for which its gradient can be computed by simple numerical approximation.
We refer to (Martin Jankowiak, 2018) for the complete derivations.
Joint Optimization of Model Weight and Architecture Parameter With pathwise derivative estimator, the model weight w and concentration β can be jointly optimized with gradient descent.
Concretely, we draw a sample θ ∼ Dir ( β ) for every forward pass, and the gradients can be obtained easily through backpropagation.
Following DARTS (Liu et al., 2019), we approximate w ∗ in the lower level objective of equation 2 with one step of gradient descent, and run alternative updates between w ∗ and β .
Selecting the Best Architecture At the end of the search phase, a learnt distribution of operation mixing weight is obtained.
We then select the best operation for each edge by the most likely operation in expectation:
[Equation]
In the Dirichlet case, the expectation term is simply the Dirichlet mean
[Equation]
Note that under the distribution learning framework, we are able to sample a wide range of architectures from the learnt distribution.
This property alone has many potentials.
For example, in practical settings where both accuracy and latency are concerned, the learnt distribution can be used to ﬁnd architectures under resource restrictions in a post search phase.
We leave these extensions to future work.
2. T HE IMPLICIT R EGULARIZATION ON H ESSIAN
It has been observed that the generalization error of differentiable NAS is highly related to the dominant eigenvalue of the Hessian of validation loss w.r.t.
architecture parameter.
Several recent works report that the large dominant eigenvalue of ∇ 2 θ ˜ L val ( w, θ ) in DARTS results in poor generalization performance (Zela et al., 2020a; Chen & Hsieh, 2020).
Our objective (2) is the Lagrangian function of the following constraint objective:
[Equation]
Here we derive an approximated lower bound based on (5), which demonstrates that our method implicitly controls this Hessian matrix.
Proposition 1 Let d ( β, ˆ β ) =
(cid:107) β − ˆ β (cid:107) 2 ≤ δ and ˆ β = 1 in the bi-level formulation (5).
Let µ denote the mean under the Laplacian approximation of Dirichlet.
If ∇ 2 µ ˜ L val ( w ∗ , µ ) is Positive Semi-deﬁnite, the upper-level objective can be approximated bounded by:
[Equation]
with:
[Equation]
This proposition is driven by the Laplacian approximation to the Dirichlet distribution (MacKay, 1998; Akash Srivastava, 2017).
The lower bound (6) indicates that minimizing the expected validation loss controls the trace norm of the Hessian matrix.
Empirically, we observe that DrNAS always maintains the dominant eigenvalue of Hessian at a low level (Appendix A.4).
The detailed proof can be found in Appendix A.1.
2. P ROGRESSIVE A RCHITECTURE L EARNING
The GPU memory consumption of differentiable NAS methods grows linearly with the size of operation candidate space.
Therefore, they usually use a easier proxy task such as training with a smaller dataset, or searching with fewer layers and number of channels (Cai et al., 2019).
For instance, the architecture search is performed on 8 cells and 16 initial channels in DARTS (Liu et al., 2019).
But during evaluation, the network has 20 cells and 36 initial channels.
Such gap makes it hard to derive an optimal architecture for the target task (Cai et al., 2019).
PC-DARTS (Xu et al., 2020) proposes a partial channel connection to reduce the memory overheads of differentiable NAS, where they only send a random subset of channels to the mixed-operation while directly bypassing the rest channels in a shortcut.
However, their method causes loss of information and makes the selection of operation unstable since the sampled subsets may vary widely across iterations.
This drawback is ampliﬁed when combining with the proposed method since we learn the architecture distribution from Dirichlet samples, which already injects certain stochasticity.
As shown in Table 1, when directly applying partial channel connection with distribution learning, the test accuracy of the searched architecture decreases over 3% and 18% on CIFAR-10 and CIFARrespectively if we send only 1/8 channels to the mixed-operation.
To alleviate such information loss and instability problem while being memory-efﬁcient, we propose a progressive learning scheme which gradually increases the fraction of channels that are forwarded to the mixed-operation and meanwhile prunes the operation space based on the learnt distribution.
We split the search process into consecutive stages and construct a task-speciﬁc super-network with the same depth and number of channels as the evaluation phase at the initial stage.
Then after each stage, we increase the partial channel fraction, which means that the super-network in the next stage will be wider, i.e., have more convolution channels, and in turn preserve more information.
This is achieved by enlarging every convolution weight with a random mapping function similar to Net2Net (Chen et al., 2016).
The mapping function g : { 1 , 2 , . . . , q } → { 1 , 2 , . . . , n } with q > n is deﬁned as
[Equation]
random sample from { 1 ,
[Equation]
. , n
[Equation]
n
[Equation]
To widen layer l , we replace its convolution weight W ( l ) ∈
R Out ×
In × H × W with a new weight U ( l ) .
[Equation]
, (8) where Out, In, H, W denote the number of output and input channels, ﬁlter height and width respectively.
Intuitively, we copy W ( l ) directly into U ( l ) and fulﬁll the rest part by choosing randomly as deﬁned in g .
Unlike Net2Net, we do not divide U ( l ) by a replication factor here because the information ﬂow on each edge has the same scale no matter the partial fraction is.
After widening the super-network, we reduce the operation space by pruning out less important operations according to the Dirichlet concentration parameter β learnt from the previous stage, maintaining a consistent memory consumption.
As illustrated in Table 1, the proposed progressive architecture learning scheme effectively discovers high accuracy architectures and retains a low GPU memory overhead
[Table]
D ISCUSSIONS AND R ELATIONSHIP TO P RIOR W ORK
[Table]
Early methods in NAS usually include a full training and evaluation procedure every iteration as the inner loop to guide the consecutive search (Zoph & Le, 2017; Zoph et al., 2018; Real et al., 2019).
Consequently, their computational overheads are beyond acceptance for practical usage, especially on large-scale tasks.
Differentiable NAS Recently, many works are proposed to improve the efﬁciency of NAS (Pham et al., 2018;
Cai et al., 2018; Liu et al., 2019; Bender et al., 2018;
Yao et al., 2020b;a; Mei et al., 2020).
Amongst them,
DARTS (Liu et al., 2019) proposes a differentiable NAS framework, which introduces a continuous architecture parameter that relaxes the discrete search space.
Despite being efﬁcient, DARTS only optimizes a single point on the simplex every search epoch, which has no guarantee to generalize well after the discretization during evaluation.
So its stability and generalization have been widely challenged (Li & Talwalkar, 2019; Zela et al., 2020a; Chen & Hsieh, 2020; Wang et al., 2021).
Following DARTS,
SNAS (Xie et al., 2019) and GDAS (Dong & Yang, 2019) leverage the gumbel-softmax trick to learn the exact architecture parameter.
However, their reparameterization is motivated from reinforcement learning perspective, which is an approximation with softmax rather than an architecture distribution.
Besides, their methods require tuning of temperature schedule (Yan et al., 2017; Caglar Gulcehre, 2017).
GDAS linearly decreases the temperature from 10 to 1 while SNAS anneals it from 1 to 0.03.
In comparison, the proposed method can automatically learn the architecture distribution without the requirement of handcrafted scheduling.
BayesNAS (Zhou et al., 2019) applies Bayesian Learning in
NAS.
Speciﬁcally, they cast NAS as model compression problem and use Bayes Neural Network as the super-network, which is difﬁcult to optimize and requires oversimpliﬁed approximation.
While our method considers the stochasticity in architecture mixing weight, as it is directly related to the generalization of differentiable NAS algorithms (Zela et al., 2020a; Chen & Hsieh, 2020).
Memory overhead When dealing with the large memory consumption of differentiable NAS, previous works mainly restrain the number of paths sampled during the search phase.
For instance,
ProxylessNAS (Cai et al., 2019) employs binary gates and samples two paths every search epoch.
PARSEC (Casale et al., 2019) samples discrete architectures according to a categorical distribution to save memory.
Similarly, GDAS (Dong & Yang, 2019) and DSNAS (Hu et al., 2020) both enforce a discrete constraint after the gumbel-softmax reparametrization.
However, such discretization manifests premature convergence and cause search instability (Zela et al., 2020b; Zhang et al., 2020).
Our experiments in section 4.3 also empirically demonstrate this phenomenon.
As an alternative, PC-DARTS (Xu et al., 2020) proposes a partial channel connection, where only a portion of channels is sent to the mixed-operation.
However, partial connection can cause loss of information as shown in section 2.4 and PC-DARTS searches on a shallower network with less channels, suffering the search and evaluation gap.
Our solution, by progressively pruning the operation space and meanwhile widening the network, searches in a task-speciﬁc manner and achieves superior accuracy on challenging datasets like ImageNet (+2.8% over BayesNAS, +2.3% over GDAS, +2.3% over
PARSEC, +2.0% over DSNAS, +1.2% over ProxylessNAS, and +0.5% over PC-DARTS).
E XPERIMENTS
In this section, we evaluate our proposed DrNAS on two search spaces: the CNN search space in
DARTS (Liu et al., 2019) and NAS-Bench-201 (Dong & Yang, 2020).
For DARTS space, we conduct experiments on both CIFAR-10 and ImageNet in section 4.1 and 4.2 respectively.
For NAS-Bench201, we test all 3 supported datasets (CIFAR-10, CIFAR-100, ImageNet-16-120 (Chrabaszcz et al., 2017)) in section 4.3.
Furthermore, we empirically study the dynamics of exploration and exploitation throughout the search process in section 4.4.
4. R ESULTS ON CIFAR-
Architecture Space
For both search and evaluation phases, we stack 20 cells to compose the network and set the initial channel number as 36.
We place the reduction cells at the 1/3 and 2/3 of the network and each cell consists of N = 6 nodes.
Search Settings We equally divide the 50K training images into two parts, one is used for optimizing the network weights by momentum SGD and the other for learning the Dirichlet architecture distribution by an Adam optimizer.
Since Dirichlet concentration β must be positive, we apply the shifted exponential linear mapping β
= ELU ( η ) + 1 and optimize over η instead.
We use l 2 norm to constrain the distance between η and the anchor ˆ η = 0 .
The η is initialized by standard Gaussian with scale 0.001, and λ in (2) is set to 0.001.
The ablation study in Appendix A.3 reveals the effectiveness of our anchor regularizer, and DrNAS is insensitive to a wide range of λ .
These settings are consistent for all experiments.
For progressive architecture learning, the whole search process consists ofstages, each with 25 iterations.
In the ﬁrst stage, we set the partial channel parameter K as 6 to ﬁt the super-network into a single GTX 1080Ti
GPU with 11GB memory, i.e., only 1/6 features are sampled on each edge.
For the second stage, we prune half candidates and meanwhile widen the network twice, i.e., the operation space size reduces from 8 to 4 and K becomes 3.
Retrain Settings
The evaluation phase uses the entire 50K training set to train the network from scratch for 600 epochs.
The network weight is optimized by an SGD optimizer with a cosine
[Table]
annealing learning rate initialized as 0.025, a momentum of 0.9, and a weight decay of 3 × 10 − 4 .
To allow a fair comparison with previous work, we also employ cutout regularization with length 16, drop-path (Zoph et al., 2018) with probability 0.3 and an auxiliary tower of weight 0.4.
Results Table 2 summarizes the performance of DrNAS compared with other popular NAS methods, and we also visualize the searched cells in Appendix A.2.
DrNAS achieves an average test error of 2.46%, ranking top amongst recent NAS results.
ProxylessNAS is the only method that achieves lower test error than us, but it searches on a different space with a much longer search time and has larger model size.
We also perform experiments to assign proper credit to the two parts of our proposed algorithm, i.e., Dirichlet architecture distribution and progressive learning scheme.
When searching on a proxy task with 8 stacked cells and 16 initial channels as the convention (Liu et al., 2019; Xu et al., 2020), we achieve a test error of 2.54% that surpasses most baselines.
Our progressive learning algorithm eliminates the gap between the proxy and target tasks, which further reduces the test error.
Consequently, both of the two parts contribute a lot to our performance gains.
4. R ESULTS ON I MAGE N ET
Architecture Space
The network architecture for ImageNet is slightly different from that for
CIFAR-10 in that we stack 14 cells and set the initial channel number as 48.
We also ﬁrst downscale the spatial resolution from 224 × 224 to 28 × 28 with three convolution layers of stride 2 following previous works (Xu et al., 2020; Chen et al., 2019).
The other settings are the same with section 4.1.
Search Settings Following PC-DARTS (Xu et al., 2020), we randomly sample 10% and 2.5% images from the 1.3M training set to alternatively learn network weight and Dirichlet architecture distribution by a momentum SGD and an Adam optimizer respectively.
We use 8 RTX 2080 Ti
GPUs for both search and evaluation, and the setup of progressive pruning is the same with that on
CIFAR-10, i.e., 2 stages with operation space size shrinking from 8 to 4, and the partial channel K reduces from 6 to 3.
Retrain Settings For architecture evaluation, we train the network for 250 epochs by an SGD optimizer with a momentum of 0.9, a weight decay of 3 × 10 − 5 , and a linearly decayed learning rate initialized as 0.5.
We also use label smoothing and an auxiliary tower of weight 0.4 during training.
The learning rate warm-up is employed for the ﬁrst 5 epochs following previous works (Chen et al., 2019; Xu et al., 2020).
[Table]
Results As shown in Table 3, we achieve a top-1/5 test error of 23.7%/7.1%, outperforming all compared baselines and achieving state-of-the-art performance in the ImageNet mobile setting.
The searched cells are visualized in Appendix A.2.
Similar to section 4.1, we also report the result achieved with 8 cells and 16 initial channels, which is a common setup for the proxy task on
ImageNet (Xu et al., 2020).
The obtained 24.2% top-1 accuracy is already highly competitive, which demonstrates the effectiveness of the architecture distribution learning on large-scale tasks.
Then our progressive learning scheme further increases the top-1/5 accuracy for 0.5%/0.2%.
Therefore, learning in a task-speciﬁc manner is essential to discover better architectures.
4. R ESULTS ON NAS-B ENCH -
Recently, some researchers doubt that the expert knowledge applied to the evaluation protocol plays an important role in the impressive results achieved by leading NAS methods (Yang et al., 2020; Li & Talwalkar, 2019).
So to further verify the effectiveness of DrNAS, we perform experiments on
NAS-Bench-201 (Dong & Yang, 2020), where architecture performance can be directly obtained by querying in the database.
NAS-Bench-201 provides support for 3 dataset (CIFAR-10, CIFAR-100,
ImageNet-16-120 (Chrabaszcz et al., 2017)) and has a uniﬁed cell-based search space containing 15,625 architectures.
We refer to their paper (Dong & Yang, 2020) for details of the space.
Our experiments are performed in a task-speciﬁc manner, i.e., the search and evaluation are based on the same dataset.
The hyperparameters for all compared methods are set as their default and for DrNAS, we use the same search settings with section 4.1.
We run every method 4 independent times with different random seeds and report the mean and standard deviation in Table 4.
As shown, we achieve the best accuracy on all 3 datasets.
On CIFAR-100, we even achieve the global optimal.
Speciﬁcally, DrNAS outperforms DARTS, GDAS, DSNAS, PC-DARTS, and SNAS by 103.8%, 35.9%, 30.4%, 6.4%, and 4.3% on average.
We notice that the two methods (GDAS and
DSNAS) that enforce a discrete constraint, i.e., only sample a single path every search iteration, perform undesirable especially on CIFAR-100.
In comparison, SNAS, employing a similar Gumbelsoftmax trick but without the discretization, performs much better.
Consequently, a discrete constraint during search can reduce the GPU memory consumption but empirically suffers instability.
In comparison, we develop the progressive learning scheme on top of the architecture distribution learning, enjoying both memory efﬁciency and strong search performance.
[Table]
4. E MPIRICAL S TUDY ON E XPLORATION V . S . E XPLOITATION
We further conduct an empirical study on the dynamics of exploration and exploitation in the search phase of DrNAS on NAS-Bench-201.
After every search epoch, We sample 100 θ s from the learned
Dirichlet distribution and take the arg max to obtain 100 discrete architectures.
We then plot the range of their accuracy along with the architecture selected by Dirichlet mean (solid line in Figure 1).
Note that in our algorithm, we simply derive the architecture according to the Dirichlet mean as described in Section 2.2.
As shown in Figure 1, the accuracy range of the sampled architectures starts very wide but narrows gradually during the search phase.
It indicates that DrNAS learns to encourage exploration in the search space at the early stages and then gradually reduces it towards the end as the algorithm becomes more and more conﬁdent of the current choice.
Moreover, the performance of our architectures can consistently match the best performance of the sampled architectures, indicating the effectiveness of DrNAS.
[Figure]
the 100 sampled architectures.
Note that the solid line is our derived architecture according to the Dirichlet mean as described in Section 2.2.
C ONCLUSION
In this paper, we propose Dirichlet Neural Architecture Search (DrNAS).
We formulate the differentiable NAS as a constraint distribution learning problem, which explicitly models the stochasticity in the architecture mixing weight and balances exploration and exploitation in the search space.
The proposed method can be optimized efﬁciently via gradient-based algorithm, and possesses theoretical beneﬁt to improve the generalization ability.
Furthermore, we propose a progressive learning scheme to eliminate the search and evaluation gap.
DrNAS consistently achieves strong performance across several image classiﬁcation tasks, which reveals its potential to play a crucial role in future end-to-end deep learning platform.
A CKNOWLEDGEMENT
This work is supported in part by NSF under IIS-1901527, IIS-2008173, IIS-2048280 and by Army
Research Laboratory under agreement number W911NF-20-2-0158.
A A PPENDIX A. P ROOF OF P ROPOSITION
Preliminaries: Before the development of Pathwise Derivative Estimator, Laplace Approximate with Softmax basis has been extensively used to approximate the Dirichlet Distribution (MacKay, 1998; Akash Srivastava, 2017).
The approximated Dirichlet distribution is:
[Equation]
Where θ ( h ) is the softmax-transformed h , h follows multivariate normal distribution, and g ( · ) is an arbitrary density to ensure integrability (Akash Srivastava, 2017).
The mean µ and diagonal covariance matrix Σ of h depends on the Dirichlet concentration parameter β :
[Equation]
It can be directly obtained from (10) that the Dirichlet mean
[Equation]
µ ) .
Sampling from the approximated distribution can be down by ﬁrst sampling from h and then applying Softmax function to obtain θ .
We will leverage the fact that this approximation supports explicit reparameterization to derive our proof.
Proof: Apply the above Laplace Approximation to Dirichlet distribution, the unconstrained upperlevel objective in (5) can then be written as:
[Equation]
[Equation] [Equation]
In our full objective
[Equation]
we constrain the Euclidean distance between learnt Dirichlet concentration and ﬁxed prior concentration
|| β − 1 || 2 ≤ δ .
The covariance matrix Σ of approximated softmax Gaussian can be bounded as:
[Equation]
Then (11) becomes:
[Equation] [Equation]
The last line holds when ∇ 2 µ ˜ L val ( w ∗ , µ ) is positive semi-deﬁnite.
In Appendix A.4 we provide an empirical justiﬁcation for this implicit regularization effect of DrNAS.
A. S EARCHED A RCHITECTURES
We visualize the searched normal and reduction cells in Figure 2 and 3, which is directly searched on
CIFAR-10 and ImageNet respectively.
[Figure] [Figure]
A. A BLATION S TUDY ON A NCHOR R EGULARIZER P ARAMETER λ
Table 5 shows the accuracy of the searched architecture using different value of λ while keeping all other settings the same.
Using anchor regularizer?
for a wide range of value can boost the accuracy and DrNAS performs quite stable under different λ s.
[Table]
E MPIRICAL S TUDY ON THE H ESSIAN R EGULARIZATION E FFECT
We track the anytime Hessian norm on NAS-Bench-201 in Figure 4.
The result is obtained by averaging from 4 independent runs.
We observe that the largest eigenvalue expands about 10 times when searching by DARTS for 100 epochs.
In comparison, DrNAS always maintains the Hessian norm at a low level, which is in agreement with our theoretical analysis in section 2.3.
Figure 5 shows the regularization effect under various λ s. As we can see, DrNAS can keep hessian norm at a low level for a wide range of λ s, which is in accordance to the relatively stable performance in Table 5.
Moreover, we compare DrNAS with DARTS and R-DARTS on 4 simpliﬁed space proposed in (Zela et al., 2020a) and record the endpoint dominant eigenvalue.
The ﬁrst space S1 contains 2 popular operators per edge based on DARTS search result.
For S2, S3, and S4, the operation sets are { 3 ×separable convolution , skip connection } , { 3 × 3 separable convolution , skip connection , zero } , and { 3 × 3 separable convolution , noise } respectively.
As shown in Table 6, DrNAS consistently outperforms DARTS and R-DARTS.
The endpoint eigenvalues for DrNAS are 0.0392, 0.0390, 0.0286, 0.0389 respectively.
Figure 5 shows the Hessian norm trajectory under various λ .
A. C ONNECTION TO V ARIATIONAL I NFERENCE
In this section, we draw a connection between DrNAS and Variational Inference (David M. Blei, 2016).
We use w , θ , and β to denote the model weight, operation mixing weight, and Dirichlet concentration parameters respectively, following the main text.
The true posterior distribution can be written as p ( θ | w, D ) , where D = { x n , y n } Nn =1 is the dataset.
Let q ( θ | β ) denote the variational approximation of the true posterior; and assume that q ( θ | β ) follows Dirichlet distribution.
We follow Joo et al.
(2019) to assume a symmetric Dirichlet distribution for the prior p ( θ ) as well,
[Table]
i.e., p ( θ ) =
Dir ( 1 ) .
The goal is to minimize the KL divergence between the true posterior and the approximated form, i.e., min β KL ( q ( θ | β ) ||
p
( θ | w, D )) .
It can be shown that this objective is equivalent to maximizing the evidence lower bound as below (David M. Blei, 2016):
[Equation]
The upper level objective of the bilevel optimization under variational inference framework is then given as:
[Equation]
Note that eq.
(23) resembles eq.
(2) if we use the negative log likelihood as the loss function and replace d ( · , · ) with KL divergence.
In practice, we ﬁnd that using a simple l2 distance regularization works well across datasets and search spaces.