SI-VDNAS: Semi-Implicit Variational Dropout for Hierarchical One-shot Neural Architecture Search

1 Introduction

Deep neural networks (DNNs) have achieved impressive performance in the past decade. However, manually designed DNNs would suffer from unexpected performance loss and tedious design in specific tasks. As an important branch of automatic machine learning (AutoML), neural architecture search (NAS) allows people to get rid of cumbersome design of DNNs. Existing NAS approaches based on evolutionary algorithms or reinforcement learning suffer from excessive calculations and enormous time cost for finding a suitable network structure. Recently, gradient-based approaches have greatly improved the practicality of NAS with a reduced time cost and amount of calculations in search. One-shot NAS is popular in gradient-based approaches, in which each search is evaluated only once to further reduce the time cost. In fact, one-shot NAS transforms the architecture search into a supernet simplification problem to improve the search speed.

One-shot NAS commonly uses the weights of candidate operations [Liu et al., 2019a] as the criterion for selecting final operations and simplifying the super-net. However, gradient descent algorithm has to be carefully designed in these methods, as skip-connect operations tend to be preserved with a large number of epochs and lead to evident degradation of evaluation performance. This problem can be relieved with additional hyperparameters to the gradient [Bi et al., 2019], sigmoid function for selecting operations [Chu et al., 2019] and violent early stopping mechanism [Liang et al., 2019]. These methods improve weight-based selection of operations in NAS, but lack interpretability and stability due to the manually designed hyper-parameters and mechanism.

∗Corresponding author

To improve the interpretability of operation selection in NAS, probabilistic methods have been introduced to determine the final operations [Xie et al., 2019; Zhou et al., 2019; Zheng et al., 2019]. Architectures obtained by sampling instead of the size of parameter is natural and more automated, but the performance is not always satisfactory. Moreover, Gumbel-softmax function is incorporated with weight-based operation selection to enable multiple sampling for preserving variable operations on one edge [Chang et al., 2019].

In this paper, we propose a novel probabilistic approach for hierarchical and stable one-shot neural architecture search, namely Semi-Implicit Variational Dropout one-shot NAS (SIVDNAS). SI-VDNAS leverages variational dropout based on an annealed semi-implicit automatic relevance detection (ARD) prior to sample the super-net and suppress the overselection of skip-connect operations. Furthermore, hierarchical structure is developed to separately achieve the selection of edges and operations. The benefits of SI-VDNAS are summarized as below.

• SI-VDNAS alleviates the performance loss led by overselection of skip-connect operations with semi-implicit variational dropout.
• SI-VDNAS preserves variable operations and edges in NAS to improve the evaluation performance.

2 Preliminaries

In the rest of this paper, for clarity, we use p(x) for the probability density function (PDF) of x and P(x = xi) for the probability that x = xi.

2.1 One-Shot NAS

One-shot NAS is a sophisticated method, composed of several cells, testing once in super-net without testing all possible structures. In one-shot NAS, a cell is defined as a directed acyclic graph (DAG) of N nodes n1, n2, · · · , nN, where directed edge(i, j) (i < j)(from input nodes to intermediate nodes or previous intermediate nodes to current intermediate node) connects the i-th node ni and j-th node nj. ni is processed by the operations o(·) in the search space O. The cell is assumed to have M input nodes n1I , n2I , · · · , nMI and a single output node nO. The input nodes are defined as the outputs in the previous M cells and the output of the cell is obtained by applying a concatenation operation to all the N -M - 1 intermediate nodes. Recent approaches for oneshot NAS commonly adopt gradient descents with parameter sharing to update each candidate operation in each edge (i, j) in each epoch. Thus, one-shot NAS can keep the search space unchanged and speed up the search, making it possible to search under the mobile settings.

2.2 Variational Dropout

We assume that a neural network has L fully-connected layers, which is consistent with the structure in the super-net of NAS. Let us denote A ∈ RM×I , W ∈ RI×O and B ∈ RM×O the input matrix, weight matrix and output matrix for each layer, respectively. Gaussian dropout adds multiplicative noise ξ ∈ RM×I to the input to each layer. Thus, B is obtained by B = (A ◦ ξ) ∗ W, where ◦ indicates the Hadamard product. In Gaussian dropout, the noise ξ is endowed with a Gaussian distribution N(1, δ) with its variance δ = p/(1 - p) determined by the dropout rate p [Srivastava et al., 2014]. Multiplying weight wij by a Gaussian noise ξ ∼ N(1, δ) is equivalent to sampling in a new Gaussian distribution N (wij , w2 ij δ). Note that each weight in each layer can be assigned an individual dropout rate δij, rather than a global δ. Using the additive reparameterization [Molchanov et al., 2017], we can substitute θij = wijδij with a lowvariance form θij = wij(1 + √δijϵij) with ϵij ∼ N(0, 1). This additive reparameterization can remove samples from the calculation graph to allow gradient back propagation.

2.3 Semi-implicit Distribution

Semi-implicit distribution p(x) is defined based on an implicit distribution p(z) and a conditional distribution with explicit PDF p(x|z).

Here, p(z) does not have explicit distribution, but its expectation can be estimated based on sampling. Given N samples from p(z), the semi-implicit distribution p(x) is estimated by

where zn represents the n-th sample extracted from p(z).

3 Methodology

In this section, we elaborate the proposed SI-VDNAS, in which a hierarchical structure for NAS is developed based on the semi-implicit variational dropout with an annealed ARD prior. Semi-implicit variational dropout is able to obtain structureswith varying numbers of operations on one edge. In addition to operation selection, the hierarchical structure enables separate selection of varying numbers of edges to improve the interpretability and stability of network search.

3.1 Semi-implicit Variational Dropout for NAS

SI-VDNAS leverages semi-implicit variational dropout to achieve NAS with an individual dropout rate for each operation, rather than a shared hyperparameter for all operations. The dropout noise is supposed to obey the semi-implicit distribution q(ξ) = ∫ q(ξ|ψ)q(ψ)dψ, where q(ξ|ψ) = N(ξ|ψ, ψδ) and q(ψ) is a Bernoulli distribution parameterized with δ, i.e., P(ψ = 1) = 1/(1 + δ) and P(ψ = 0) = δ/(1 + δ). Note that δ is also introduced in the variance of q(ξ|ψ). Thus, the semi-implicit distribution qδ(ξ) is formulated with regard to δ.

In one-shot NAS, the dropout noise ξ ∼ qδ(ξ) is assigned to each candidate operation o(x) on x. Let us define τ o i,j = γ o i,j · ξ o i,j for the weight γ o i,j of the operation o in the edge (i, j). For simplicity, we omit the superscripts and subscripts in the rest of this subsection. According to Eq. (3), τ ∼ qδ(τ ) = ∫ qδ(τ |ψ)qδ(ψ)dψ with qδ(τ |ψ) = N (τ |γψ, γ 2ψδ) for each τ = γ · ξ. Due to the discrete Bernoulli distribution qδ(ψ), the integral for the semi-implicit distribution qδ(τ) is intractable. Thus, we leverage the Monte Carlo sampling to calculate qδ(τ) fromK samples ψ1 , · · · , ψK from qδ(ψ).

Here, we set K = 1 to make stochastic update based on semi-implicit variational dropout. The Evidence Lower BOund (ELBO) L(δ) is developed for semi-implicit variational dropout by taking qδ(ψ) as the posterior distribution.

In Eq. (5), LD(δ) is empirically obtained based on the training set that consists of N pairs of observations and labels {xn, yn}, n = 1, · · · ,N. Consequently, KL divergence DKL(qδ(τ )‖p(τ )) is minimized for operation selection under the properly designed prior p(τ), as shown in Section. 3.2.

3.2 Annealed Semi-implicit ARD Prior

We further design a semi-implicit prior pλ,δ(τ) for semiimplicit variational dropout.

where pδ(φ) is also a Bernoulli distribution parameterized with δ and pλ(τ|φ) = N(τ|λγφ, η-1φ) is a Gaussian PDF

associated with the variance η-1 of ARD prior and the annealed temperature λ. From Eq. (4) and (6), we formulate the KL divergence DKL (qδ(τ )‖p(τ )) for minimization.

When we set K = N = 1 for stochastic update, we have

Here, we use Φ to represent ψ∗ and φ∗ , as they are taken from the same Bernoulli distribution. Thus, we have DKL(qδ(τ )‖pλ,δ(τ )) ≈ Eqδ(τ |Φ) log(qδ(τ |Φ)/pδ(τ |Φ)).

Proposition 1. The proposed annealed semi-implicit ARD prior pλ,δ(τ) generalizes the ARD prior N(τ|0, η-1 ) to induce the KL divergence DKL(qδ(τ)‖pλ,δ(τ)) that are independent of the weight parameters γ.

Proof. Given Φ sampled from the Bernoulli distribution qδ(φ), we have qδ(τ) ≈ qδ(τ|Φ) = N(τ|γΦ, γ 2Φδ) and pλ,δ(τ) ≈ pλ(τ|Φ) = N(τ|λγΦ, η-1Φ). When Φ = 0, qδ(τ) and pδ(τ) degenerate to the constant 0, which means their KL divergence equals to 0. For Φ = 1, the KL divergence between qδ(τ) and pλ,δ(τ) is approximated by

The optimal value η∗ = argminη DKL(qδ(τ|Φ)||pλ(τ|Φ)). From the gradient of DKL(qδ(τ |Φ)||pλ(τ |Φ)), we have

Thus, we can approximate DKL(qδ(τ )‖pλ,δ(τ )) by

This KL divergence does not depend on weight parameters γ. Note that pλ,δ(τ) degenerates to the vanilla ARD prior N(τ|0, η-1 ), when λ = 0 and Φ = 1. The KL divergence degenerates to 0.5 · log(1 + δ-1 ), which is adopted for the hierarchical structure in [Liu et al., 2019b].

3.3 Hierarchical Structure

Edge selection is implicit in previous one-shot approaches, in which only operation weight parameters are used, resulting ambiguous edge selection and poor evaluation performance. Edge normalization introduces the edge weight parameters to assist operation weight parameters in searching the final structure [Xu et al., 2019]. It can restrict the search of candidate edges between nodes, but cannot separately achieve the selection of edges and operations. In SI-VDNAS, we assign individual dropout rates to the edge weight parameters as well as operation weight parameters, separating the edge selection from the operation selection and achieving architecture with variable edges.

Let us denote ˜ fi,j(x) the output of mixed operations in edge (i, j) and ˜hj(x) the mixed edge on the node j. The edge weight parameters ϕ and the operation weight parameters w are referred as the architecture parameters γ. Considering vanilla variational dropout, for the operation weight parameters, we assign the dropout rate ξ o i,j ∼ N(ξ o i,j |1, δo i,j ) to the input opoi,j (x) and obtain:

˜ fi,j(x) can be calculated using samples from N( ˜ fi,j (x)|wo i,j opoi,j , (wo i,j opoi,j )2 δo i,j ).

For the edge weight parameters, we assign the dropout rate ζi,j ∼ N(ζi,j|1, σi,j) to ˜ fi,j(x). ˜hj(x) can be similarly sampled fromN(˜hj(x)|ϕi,j ˜ fi,j(x), (ϕi,j ˜ fi,j (x))2σi,j ).

Without loss of generality, we consider arbitrary one edge (i, j). Denote µo = wo · ξ o ∼ N(µo |wo , (wo )2 δo ) and υ = ϕ · ζ ∼ N(υ|ϕ, ϕ2σ). The KL divergence can be developed for the approximate posterior distribution q(υ, µ) = q(υ|µ) ∏o∈O q(µo ) and the hierarchical prior distribution p(υ, µ) = p(υ|µ) ∏o∈O p(µo ) under independent µ and υ.

Consequently, we introduce semi-implicit variational dropout. The posterior distribution q(υ|µ) is approximated by

where Φe is the sample from the Bernoulli distribution q(ψe). The posterior q(µo ) is similarly approximated with Φoo sampled from the Bernoulli distribution q(ψo o ).

Subsequently, we adopt the annealed semi-implicit ARD prior to formulate the KL divergence for training.

Algorithm 1 Semi-implicit Variational Dropout NAS

Input: Data {xn, yn}1:N, network parameter p, architecture parameter (w, ϕ, δ, σ): operation weight parameter w, edge weight parameter ϕ, dropout rate δ for operation and σ for edge, operation Bernoulli variable ψo and edge Bernoulli variable ψe, dropout noise ξ ∼ N(ξ|ψo, ψoδ) for w, the variational posterior distribution N(µ|wψo,w2ψoδ), dropout noise ζ ∼ N(ζ|ψe, ψeσ) for ϕ, the variational posterior distribution N(υ|ϕψe, ϕ2ψeσ), prior distribution pλo,δ(µ) with the annealed temperature λo for operation and pλe,σ(υ|µ) with the annealed temperature λe for edge.

Output: The individual dropout rate δ for operation and σ for edge, the operation weight parameter w and the edge weight parameter ϕ, the sample of ψo and ψe.

• 1: Initialize w, ϕ, δ and σ.
• 2: while not converged do
• 3: Sample Φo and Φe fromψo and ψe respectively.
• 4: Approximate the semi-implicit posterior distribution as q(µ) ≈ q(µ|Φo), q(υ|µ) ≈ q(υ|µ, Φe). Approximate the semi-implicit prior as p(µ) ≈ p(µ|Φo), p(υ|µ) ≈ p(υ|µ, Φe).
• 5: Calculate DKL(q(υ|µ,Φe)‖p(υ|µ,Φe)) using Eq. (17) and DKL(q(µ|Φo)‖p(µ|Φo)) using Eq. (18).
• 6: Update w, ϕ, δ, and σ by gradient descent.
• 7: Update network parameters p by gradient descent .
• 8: (1-λnew o )2 = 0.95 · (1-λold o )2 .
• 9: (1-λnew e )2 = 0.95 · (1-λold e )2 .
• 10: end while
• 11: return w, ϕ, δ, and σ and final samples Φo, Φe from ψo and ψe respectively.

From Eq. (17) and (18), we obtain the objective KL divergence in Eq. (14). Algorithm 1 elaborates the implementation of SI-VDNAS.

4 Experiments

Following the pipeline of DARTS, our experiments can be divided into 3 stages. In the first stage, we search the architecture on CIFAR-10 with a simple network. The optimized architecture in search is stored and stacked to generate a complex network for new training processes from scratch. The new training processes consist of two stages, i.e., evaluation on CIFAR-10/100 and ImageNet.

4.1 Datasets

CIFAR-10/100 [Krizhevsky and Hinton, 2009] is a popular dataset consisting of 60K images, 50K training images and 10K test images. All these images share the same spatial resolution of 32×32, and are categorized into 10/100 classes.

ImageNet [Deng et al., 2009] is a large-scale benchmark for image classification. It contains 1.3M training images and

50K test images that are equally distributed into 1000 classes. As in [Zoph et al., 2018; Xie et al., 2019; Liu et al., 2019a], we adopt the mobile setting with a spatial resolution of 224× 224 for input images to limit FLOPS by 600M during testing.

4.2 Architecture Search

Implementation Details

Following DARTS, the super-net is formed by stacking two kinds of basic cells. Specifically, 6 normal cells and 2 reduction cells are stacked to form the super-net. Each cell contains seven nodes, including two input nodes, four intermediate nodes, and one output node. The output of the four intermediate nodes are concatenated as the input for the output node. Each cell has 14 candidate edges, where a hybrid operation consisting of 7 candidate operations is assigned to each candidate edge.

For fair comparison, we adopt the same candidate operations as existing one-shot NAS, except for 'none' operation. We assign dropout rate to operation weight parameters w and edge weight parameters φ, respectively. So, we have two stages of our object function. In each stage, we use the semi-implicit distribution qδ(ξ) = ∫ qδ(ξ|ψ)qδ(ψ)dψ as the approximate posterior. The implicit distribution qδ(ψ) is chosen to be the Bernoulli distribution parameterized by the dropout rate δ. Thus, we can sample the super-net by sampling from qδ(ψ). The final structure is also determined by the samples of qδ(ψ). The first part of semi-implicit distribution is the Gaussian distribution which is the same as that of VDNAS when ψ = 1. The prior distribution is chosen to be the annealed semi-implicit ARD distribution, which is the generalized version of ARD prior used in previous work [Liu et al., 2019b]. For our training, we set the annealed temperature as (1-λnew )2 = 0.95 · (1-λold )2 , and the second part of prior shares the same sample of qδ(ψ).

We utilize the bi-level update algorithm to update the architectural parameters and conventional network weights, respectively. For the training, we split the training images into two subsets with the same size. One subset is used for training network parameters, the other is used for architectural parameters. We can not only train the network for 50 epochs following DARTS to get the optimal structure, but also train 150 epochs for convergence(the number of epochs does not have to be 150, but it can also be 100 or 300. We can search to get a convergent result without the need to carefully design the search process for preventing degeneration), with the initial number of channels being 16. Following [Chen et al., 2019], we freeze architectural parameters and only update the network parameters in first 15 epochs. The batch-size is set to 64 to enable the searching process on single GPU.

Search Results

The search process requires 8 GPU-hours for optimal structure within 50 epochs and 20 GPU-hours for a convergent result within 150 epochs on a single NVIDIA GTX 1080Ti GPU. The search time can be reduced by about 50% on a single Tesla V100 GPU. Figure 1 shows the optimal structure within 50 epochs and the convergent structure within 150 epochs. Due to space limitations, we only show the searched normal cells, as the number of reduction cells is small.

(a) Base normal cell

(b) Convergent normal cell

Figure 1: Searched normal cells on CIFAR-10.

4.3 Architecture Evaluation

Evaluation on CIFAR-10/100

Evaluation on CIFAR-10/100 follows that of DARTS and PDARTS [Chen et al., 2019]. The network is composed of 20 cells (18 normal cells and 2 reduction cells). Each cell has the same architecture learned in our searching stage. The initial number of channels is set to 36. The network weights are trained from scratch using all the 50K training images with a batch size of 96. The network is trained for 600 epochs. We use the SGD optimizer with an initial learning rate of 0.025 (annealed down to zero following a cosine schedule without restart), a momentum of 0.9, a weight decay of 3 × 10-4 /5 × 10-4 and a norm gradient clipping at 5. We apply the drop-path trick with the probability of 0.3. Cutout is also used in our evaluation. Table 1 summarizes the results obtained by the state-of-the-arts and SI-VDNAS on CIFAR10/100. SI-VDNAS outperforms all the benchmarks in terms of test accuracy. It can achieve 2.50% top-1 error rate with 3.6M parameters in 0.3 GPU-days. The convergent architecture of SI-VDNAS also achieves 2.60% top-1 error rate and 2.7M parameters in 0.8 GPU-days.

Evaluation on ImageNet

We further evaluate the searched architectures from CIFAR10 on ImageNet to validate its generalizability. The evaluation process is similar to DARTS. The input images share the resolution of 224×224 for Mobile Setting. Three convolution layers of stride 2 are first employed to reduce the sizes of input images from 224×224 to 28×28. The network consists of 12 normal cells and 2 reduction cells. The initial number of channels is set to 48. The network is trained from scratch for 250 epochs using a batch size of 1024. We use the SGD optimizer with a momentum of 0.9, an initial learning rate of 0.5 (decayed down to zero linearly) and a weight decay of 3 × 10-5 . We also adopt label smoothing and an auxiliary loss tower during training. Learning rate warm-up is applied for the first 5 epochs.

Table 2 provides the evaluation results. The effectiveness of transferring SI-VDNAS is demonstrated. The basic architecture searched on CIFAR-10 can achieve 25.3% top-1 error rate and 8.0% top-5 error rate with 5.0 MB parameters and 577M FLOPS. The convergent architecture searched on CIFAR-10 can achieve 25.6% top-1 error rate and 8.2% top5 error rate with only 4.1 MB parameters and 462M flops, which outperforms DARTS (26.7% top-1 error rate and 8.7% top-5 error rate). SI-VDNAS also reduces 1%-2% top-1 and top-5 error rate in comparison to SNAS. Note that AmendedDARTS and P-DARTS adopt deeper search space (more than 8 layers stacked as the super-net) than DARTS. Furthermore, PC-DARTS and DARTS+ directly performs NAS on ImageNet. Structures achieved by these approaches are usually deeper and naturally perform better on large dataset consisting of high-resolution images. However, these approaches actually change the search space and are time consuming.

4.4 Ablation Study

To verify the design of SI-VDNAS, we further evaluate the strategies excluding the semi-implicit distribution from SIVDNAS. Two search strategies, namely H-VDNAS and VVDNAS, are adopted, where V-VDNAS utilizes only vanilla variational dropout and H-VDNAS considers hierarchical structure over V-VDNAS. Table 3 shows the top-1 error rates on CIFAR-10/100 obtained by SI-VDNAS, V-VDNAS, HVDNAS and DARTS, respectively. These results imply that both variational dropout and hierarchical structure contribute to the performance gain by SI-VDNAS. We also depict the searched normal cells by V-VDNAS and H-VDNAS in Figure 2. Moreover, H-VDNAS and V-VDNAS are also affected by the over-selection of skip-connect operations with a large number of epochs, as semi-implicit variational dropout is not adopted to preserve variable operations. Note that it is not desirable to directly separate the hierarchical structure from SI-VDNAS, as multiple operations are preserved in one edge.

5 RelatedWork

In this section, we briefly introduce previous works paying attention to NAS. Evolutionary algorithms were adopted in [Elsken et al., 2019; Miikkulainen et al., 2019; Real et al., 2017; Real et al., 2019] to evolve one single network or a family of networks towards better performance. Reinforcement learning (RL) based methods [Zoph and Le, 2017; Zoph et al., 2018; Bender et al., 2018; Pham et al., 2018] utilized a meta-controller to guide the search process in the huge space of architecture by optimizing the reward function for the inference accuracy of the selected network. To narrow the search space, ENAS [Pham et al., 2018] stacked repeated cells to form the final structure, which was also adopted in gradient-based methods such as DARTS.

Despite the good performance, a large amount of search time and consumption of computing resources limit the development of evolution based methods and reinforcement based methods. Recently, Gradient based methods, especially oneshot models, have gradually become the mainstream method of NAS. Among them, DARTS [Liu et al., 2019a] proposed a differentiable framework by introducing architectural parameters to measure the importance of candidate operations. Many works are then based on DARTS. SNAS [Xie et al., 2019] introduced the concrete distribution to NAS and replaced softmax function in DARTS with Gumbel-softmax function. PC-DARTS proposed partial channel connection to reduce the memory cost in NAS and accelerate the search process. P-DARTS gradually improved the width and depth of search space to bridge the gap between search and evaluation. Recently, DARTS is discovered to degenerate when searched until convergence. DARTS+ [Liang et al., 2019] introduced early stopping into NAS, Fair-DARTS [Chu et al., 2019] replaced the softmax function with sigmoid function, Amended-DARTS added the hyperparameter to th gradient of DARTS. Too many artifacts and lack of interpretability are common to these methods, our proposed SI-VDNAS leverages variational dropout based on an annealed semi-implicit ARD prior, solving the degeneration and also improving the evaluation performance by enlarging the evaluation space.

6 Conclusion

In this paper, we proposed a probabilistic NAS approach, named Semi-Implicit Variational Dropout Neural Architecture Search (SI-VDNAS). The core idea of SI-VDNAS is to use the semi-implicit variational dropout and annealed semiimplicit ARD prior replace the vanilla variational dropout and ARD prior, hierarchical structure is also used to separate the edge selection and operation selection. SI-VDNAS can solve the degeneration occurred in previous one-shot NAS sharing the search space of DARTS. SI-VDNAS can also improve the evaluation performance by preserving variable operations and variable edge. In terms of performance, SI-VDNAS can approach the state-of-art result in CIFAR-10/100 and outperforms the benchmark when transformed to ImageNet.

Acknowledgements

This work was supported in part by the National Natural Science Foundation of China under Grants 61971285, 61720106001, 61932022, 61931023, 61972256, 61871267 and 91838303, and in part by the Program of Shanghai Academic Research Leader under Grant 17XD1401900.