PaperlayerView exact revision on arXiv

Source-linked AI summary

DrNAS: Dirichlet Neural Architecture Search

Xiangning Chen, Ruochen Wang, Minhao Cheng, Xiaocheng Tang, Cho-Jui Hsieh

arXiv:2006.10355v4cs.LGcs.CVstat.ML

TL;DR

Differentiable NAS can overfit validation data and often relies on smaller proxy tasks because of memory demands. DrNAS learns a Dirichlet distribution over operation weights and adds progressive learning for large-scale tasks. It reports strong benchmark results, including a 2.46% CIFAR-10 average error rate and 23.7%/7.1% ImageNet top-1/5 test errors.

  • Problem

    Differentiable NAS can overfit validation data and produce distorted architectures, while memory demands often require proxy tasks that differ from evaluation.

    Source: §1 · p. 1, §2.2 · p. 3

  • Method

    DrNAS learns a Dirichlet distribution over operation mixing weights and uses progressive learning that increases channel fractions while pruning operations based on the learned distribution.

    Source: §2.2 · p. 3, §2.4 · p. 5

  • Results

    DrNAS reports strong performance across image-classification tasks, including 2.46% average error on CIFAR-10 and 23.7%/7.1% top-1/5 test error on ImageNet's mobile setting.

    Source: §1 · p. 2, §4.2 · p. 8, §5 · p. 9

  • Takeaways & Limitations

    The lower bound links minimizing expected validation loss to control of the validation-loss Hessian trace norm.

    Source: §2.3 · p. 4

    Using the learned distribution to find architectures under resource restrictions in a post-search phase is left for future work.

    Source: §2.2 · p. 4

Abstract

from arXiv · show

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. Specifically, we obtain a test error of 2.46% for CIFAR-10, 23.7% for ImageNet under the mobile setting. On NAS-Bench-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.

1 INTRODUCTION

Differentiable NAS reduces search costs, but validation overfitting and memory-driven proxy tasks limit stability and create a search–evaluation gap. DrNAS learns a distribution over operation weights and adds progressive learning, reporting strong results across several benchmarks.

1 INTRODUCTION

  • Earlier NAS approaches used reinforcement learning, evolutionary algorithms, or Bayesian optimization; DARTS made search more efficient by learning continuous operation mixing weights with gradient optimization.The paper frames DrNAS as addressing shortcomings in differentiable NAS.Source: §1 · p. 1
  • Differentiable NAS methods can overfit validation data and produce distorted architectures, while memory demands often force proxy searches with smaller datasets or narrower networks.These proxy setups create disparities between search and evaluation phases.Source: §1 · p. 1
  • DrNAS treats operation mixing weights as random variables sampled from a learnable Dirichlet distribution, using pathwise derivatives for end-to-end optimization.A distance regularizer on the concentration parameter balances exploration and exploitation; a theoretical bound links the constrained objective to stability and generalization.Source: §1 · pp. 1–2
  • Progressive learning maintains a task-specific super-network at evaluation depth and width, then increases the channel fraction in stages and prunes operations using the learned distribution.The scheme is designed to preserve memory efficiency while enabling direct search on large-scale tasks.Source: §1 · p. 2
  • DrNAS reports a 2.46% average error rate on CIFAR-10 and ImageNet top-1/5 errors of 23.7%/7.1%, versus 24.0%/7.3% for the previous state of the art in the mobile setting.It also reports state-of-the-art results on all three NAS-Bench-201 datasets.Source: §1 · p. 2

2 THE PROPOSED APPROACH

DrNAS frames differentiable architecture search as learning a distribution over operation-mixing weights, and derives an implicit Hessian regularization result. It also proposes progressive learning to reduce memory costs while searching at evaluation-scale network dimensions.

2.1 PRELIMINARIES: DIFFERENTIABLE ARCHITECTURE SEARCH

  • Continuous relaxation represents each edge as a weighted sum of candidate operations, with the highest mixing weight selecting the operation.The mixing weights lie on the probability simplex.Source: §2.1 · p. 2

2.2 DIFFERENTIABLE ARCHITECTURE SEARCH AS DISTRIBUTION LEARNING

  • DrNAS treats operation-mixing weights as random variables from a learnable Dirichlet distribution rather than directly optimized point estimates.The method uses pathwise derivative estimators to optimize the distribution parameters and regularizes their distance from the symmetric Dirichlet anchor.Source: §2.2 · p. 3
  • Pathwise derivative estimators enable gradient-based updates through Dirichlet samples, with model weights and concentration parameters updated alternately.Each forward pass draws a mixing-weight sample from Dir(β).Source: §2.2 · pp. 3–4

2.3 THE IMPLICIT REGULARIZATION ON HESSIAN

  • Minimizing expected validation loss under the derived bound controls the Hessian trace norm; empirically, DrNAS maintains a low dominant Hessian eigenvalue.The bound is derived using a Laplacian approximation to the Dirichlet distribution.Source: §2.3 · p. 4

2.4 PROGRESSIVE ARCHITECTURE LEARNING

  • Differentiable NAS memory grows with the operation candidate space, so progressive learning starts with a task-specific super-network matching evaluation depth and channel count.It increases channel fractions across stages to preserve more information while retaining a consistent memory budget through operation pruning.Source: §2.4 · pp. 4–5
  • With only 1/8 of channels sent to mixed operations, partial-channel distribution learning reduces test accuracy by over 3% on CIFAR-10 and 18% on CIFAR-100.The proposed scheme gradually increases the forwarded channel fraction and prunes less important operations using the learned distribution, while widening the super-network between stages.Source: §2.4 · p. 5

3 DISCUSSIONS AND RELATIONSHIP TO PRIOR WORK

Prior NAS methods trade computational efficiency against stability, generalization, or fidelity between search and evaluation. The discussion contrasts these approaches with DrNAS’s distribution learning and progressive search, and reports ImageNet accuracy gains over several methods.

Memory overhead

Differentiable NAS

  • Table 1 reports NAS-Bench-201 test accuracy for derived architectures searched with different partial-channel fractions, where 1/K channels enter the mixed operation.The caption specifies the comparison scope but supplies no accuracy values.Source: Table 1 · p. 5
  • DARTS optimizes one simplex point per search epoch, with no guarantee of generalization after discretization; its stability and generalization have been challenged.The passage cites Li & Talwalkar (2019), Zela et al. (2020a), Chen & Hsieh (2020), and Wang et al. (2021).Source: §3 · p. 5
  • SNAS and GDAS use Gumbel-softmax reparameterization, while DrNAS learns an architecture distribution without a handcrafted temperature schedule.The passage characterizes Gumbel-softmax as a softmax approximation and notes that SNAS and GDAS require temperature tuning.Source: §3 · p. 6

Memory overhead

  • On ImageNet, DrNAS reports accuracy gains of +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.The passage presents these as advantages on challenging datasets such as ImageNet.Source: §3 · p. 6
  • Discrete path sampling and partial-channel connections reduce memory use, but the passage links discretization to premature convergence and instability, and partial connections to information loss and a search–evaluation gap.PC-DARTS searches a shallower network with fewer channels; the proposed approach progressively prunes operations while widening the network.Source: §3 · p. 6

4 EXPERIMENTS

DrNAS is evaluated on the DARTS search space for CIFAR-10 and ImageNet, on NAS-Bench-201’s three datasets, and through an exploration–exploitation study. It achieves strong benchmark results, while sampled-architecture accuracy ranges narrow during search and the Dirichlet-mean architecture matches the best sampled performance.

4 EXPERIMENTS

  • Experiments cover DARTS-space CIFAR-10 and ImageNet, all three NAS-Bench-201 datasets, and search-process exploration and exploitation.The CIFAR-10 and ImageNet evaluations use DARTS space; NAS-Bench-201 includes CIFAR-10, CIFAR-100, and ImageNet-16-120.Source: §4 · p. 6

4.1 RESULTS ON CIFAR-10

  • 2.46% average test error on CIFAR-10 ranks DrNAS among the top recent NAS results; ProxylessNAS is lower but uses a different space, longer search, and larger model.On the conventional proxy task, DrNAS reaches 2.54% test error, and progressive learning further reduces test error by eliminating the proxy–target gap.Source: Table 2 · p. 7, §4.1 · p. 7

4.2 RESULTS ON IMAGENET

  • 23.7%/7.1% top-1/5 test error on ImageNet mobile outperforms all compared baselines and achieves state-of-the-art performance.With the 8-cell, 16-channel proxy setup, DrNAS obtains 24.2% top-1 accuracy; progressive learning further increases top-1/5 accuracy by 0.5%/0.2%.Source: Table 3 · p. 8, §4.2 · p. 8

4.3 RESULTS ON NAS-BENCH-201

  • DrNAS achieves the best accuracy on all three NAS-Bench-201 datasets and reaches the global optimum on CIFAR-100.Experiments use task-specific search and evaluation; each method is run four times with different random seeds.Source: §4.3 · p. 8, Table 4 · p. 9

4.4 EMPIRICAL STUDY ON EXPLORATION V.S. EXPLOITATION

  • The sampled architectures’ accuracy range narrows over search, while the Dirichlet-mean architecture consistently matches the best sampled performance.The study samples 100 architectures after each search epoch on NAS-Bench-201.Source: Figure 1 · p. 9, §4.4 · p. 9

5 CONCLUSION

DrNAS formulates differentiable NAS as constraint distribution learning, explicitly modeling stochasticity in architecture mixing weights to balance exploration and exploitation. It also introduces progressive learning to eliminate the search–evaluation gap and reports strong performance across image-classification tasks.

5 CONCLUSION

  • DrNAS models stochasticity in architecture mixing weights through constraint distribution learning, balancing exploration and exploitation.The method is optimized efficiently with a gradient-based algorithm and has a theoretical benefit for generalization ability.Source: §5 · p. 9
  • Progressive learning is proposed to eliminate the gap between architecture search and evaluation.Source: §5 · p. 9
  • DrNAS consistently achieves strong performance across several image-classification tasks.The paper describes this as revealing potential for future end-to-end deep-learning platforms.Source: §5 · p. 9

A.1 PROOF OF PROPOSITION 1

The proof uses a Laplace approximation that represents Dirichlet samples through a softmax-transformed Gaussian. It relates the approximation’s parameters to the Dirichlet concentration and establishes a positive-semidefinite property for the objective.

A.1 PROOF OF PROPOSITION 1

  • Sampling from the approximation draws h from a multivariate normal distribution and applies softmax to obtain θ; its mean and diagonal covariance depend on β.The softmax-transformed mean equals the normalized Dirichlet concentration parameters.Source: §A.1 · p. 14
  • The proof applies a Laplace approximation to the Dirichlet distribution to rewrite the unconstrained upper-level objective.The approximation uses a softmax-transformed multivariate normal variable.Source: §A.1 · p. 14
  • The full objective constrains the distance between the learned concentration and the fixed prior, and the proof establishes a positive-semidefinite property with empirical justification in Appendix A.4.The stated constraint is ||β − 1||2 ≤ δ.Source: §A.1 · p. 14

A.2 SEARCHED ARCHITECTURES

The paper visualizes normal and reduction cells discovered by DrNAS on CIFAR-10 and ImageNet.

A.2 SEARCHED ARCHITECTURES

A.3 ABLATION STUDY ON ANCHOR REGULARIZER PARAMETER λ

The anchor-regularizer ablation reports stable performance across a wide range of λ values, with accuracy boosted by using the regularizer.

A.3 ABLATION STUDY ON ANCHOR REGULARIZER PARAMETER λ

  • DrNAS performs stably across a wide range of λ values, and using the anchor regularizer can boost accuracy.Table 5 reports test accuracy for searched architectures on NAS-Bench-201 CIFAR-10; λ = 1e−3 is the setting used in all experiments.Source: §A.3 · p. 15, Table 5 · p. 15

A.4 EMPIRICAL STUDY ON THE HESSIAN REGULARIZATION EFFECT

On NAS-Bench-201, DrNAS maintains a low Hessian norm during search, while DARTS’ largest eigenvalue expands about tenfold after 100 epochs. DrNAS also has lower reported endpoint dominant eigenvalues than DARTS and R-DARTS across four simplified spaces.

A.4 EMPIRICAL STUDY ON THE HESSIAN REGULARIZATION EFFECT

  • After 100 epochs on NAS-Bench-201, DARTS’ largest eigenvalue expands about 10 times, while DrNAS maintains the Hessian norm at a low level.The Hessian trajectory is averaged over four independent runs.Source: §A.4 · p. 15
  • DrNAS endpoint dominant eigenvalues are 0.0392, 0.0390, 0.0286, and 0.0389 across four simplified spaces, and it consistently outperforms DARTS and R-DARTS.The four spaces use different sets of candidate operations.Source: §A.4 · p. 15

A.5 CONNECTION TO VARIATIONAL INFERENCE

The section connects DrNAS to variational inference by treating operation-mixing weights as a Dirichlet variational approximation to the true posterior. It relates the resulting objective to the paper’s search objective and reports that simple L2 regularization works across datasets and search spaces.

A.5 CONNECTION TO VARIATIONAL INFERENCE

  • Figures 4 and 5 track Hessian-norm trajectories on NAS-Bench-201 while searching with CIFAR-10, with Figure 5 varying λ.The captions specify the plotted scope but do not report trajectory outcomes.Source: Figure 4 · p. 16, Figure 5 · p. 16
  • Table 6 reports CIFAR-10 test error across four simplified search spaces.The caption does not provide the table’s values or comparisons.Source: Table 6 · p. 16
  • The variational objective minimizes KL divergence between the Dirichlet approximation q(θ|β) and the true posterior, equivalently maximizing an evidence lower bound.The prior p(θ) is assumed to be symmetric Dirichlet with concentration 1.Source: §A.5 · pp. 16–17
  • The paper identifies the variational-inference upper-level objective as part of its bilevel optimization framework.The supplied passage introduces the objective but does not state its formula.Source: §A.5 · p. 17
  • The variational objective resembles the main search objective when negative log likelihood is the loss and the distance term is replaced by KL divergence.In practice, simple L2 distance regularization works across datasets and search spaces.Source: §A.5 · p. 17
Loading 2006.10355v4…