Differentiable architecture search

Overview

Differentiable architecture search (DAS) treats the choice of architecture as a set of continuous parameters that can be trained by gradient descent. The search space is a cell: a small directed acyclic graph of ordered nodes whose edges each carry one operation picked from a candidate set [arxiv:1908.00709 §4.1]. In the reference form, each edge is replaced by a softmax-weighted sum of all candidate operations, with one architecture parameter \alpha per candidate 1908.00709 eq. 4 1908.00709 §4.2.3.

The reference loop fits the architecture parameters \alpha on validation data and the network weights \theta on training data, in alternation 1908.00709 §4.2.3 1908.00709 eq. 5. When the search ends, each edge keeps only its highest-weight operation 1908.00709 §4.2.3 1908.00709 fig. 16. The survey describes DARTS as the representative method of this family 1908.00709 §4.2.3. The variants below depart from the reference in different ways. BayesNAS trains on the training data and prunes by a posterior variance 1905.04919 alg. 1. ProxylessNAS keeps only one sampled path per edge in memory and adds a latency estimate to the loss 1812.00332 §3.2 1812.00332 §3.3.1. LocalMamba keeps the softmax relaxation but applies it to scan orders in a vision model, and keeps the top four candidates per layer instead of one 2403.09338 §4.2.

The method sits in the gradient-descent branch of the survey's optimizer taxonomy, next to evolutionary, reinforcement-learning and surrogate methods 1908.00709 §4.2. This page covers gradient-based NAS only. Gradient methods for hyperparameter tuning are a separate branch [arxiv:1908.00709 §4.3.3].

This page draws on four cards. AutoML: A Survey of the State-of-the-Art runs no experiments of its own. Its DARTS description and table rows come from the original papers, and the primary DARTS paper has no card here, so the DARTS specifics below are the survey's summary, not verified against the source. Check the primary paper before implementing. BayesNAS is a method paper with its own experiments. Its numbers are author-reported, and this wiki holds no independent replication. ProxylessNAS is also a method paper with its own experiments, on CIFAR-10 and ImageNet. Its numbers are likewise author-reported. LocalMamba is a method paper whose search picks scan directions for a vision state-space model. Its ImageNet, COCO and ADE20K results are author-reported too.

Variants

The variants differ in how they relax the discrete choice, how they connect the two optimisation problems, what criterion decides what to keep, and how they limit memory or search bias.

Bilevel DARTS

The reference form optimises \alpha on validation loss with \theta^* defined as the minimiser of the training loss for that \alpha 1908.00709 eq. 5. The survey gives two rows for it on CIFAR-10: first-order at 97.00% accuracy and 1.5 GPU-days, and second-order at 97.23% and 4 GPU-days, both with Cutout 1908.00709 tab. 3. The GPU-day figure is the survey's copy of each paper's report; GPU model and count differ across rows 1908.00709 eq. 9. Card: AutoML: A Survey of the State-of-the-Art.

Single-level and mixed-level objectives

Single-level optimisation minimises the training loss over \theta and \alpha together 1908.00709 eq. 6. The survey warns that \alpha then tends to overfit the training set 1908.00709 §4.2.3. Mixed-level optimisation (MiLeNAS) adds a weighted validation term to the training loss 1908.00709 eq. 7. The survey reports MiLeNAS at 97.66% on CIFAR-10 with 0.3 GPU-days 1908.00709 tab. 3. The survey gives no detail on how the weight of the validation term is set.

Sampled mixtures (Gumbel-softmax)

GDAS, SNAS and FBNet sample the operation choice with the Gumbel-softmax relaxation instead of using all operations at once 1908.00709 eq. 8 1908.00709 §4.2.3. GDAS runs only the argmax operation in the forward pass and backpropagates through the relaxation 1908.00709 §4.2.3. The survey reports GDAS at 97.18% on CIFAR-10 with 0.17 GPU-days on one GPU, and SNAS at 97.15% with 1.5 GPU-days 1908.00709 tab. 3. Because only sampled operations run, memory does not grow with the number of candidates the way it does for DARTS 1908.00709 §4.2.3.

Single-path binarised gates with hardware cost (ProxylessNAS)

ProxylessNAS gives each candidate path its own architecture weight, samples a binary gate vector from their softmax, and activates exactly one path 1812.00332 §3.2 1812.00332 eq. 2. Only the active path's activations are stored, which is the basis of its memory argument 1812.00332 §3.2. The architecture update is a gradient with respect to the sampled gate, pushed through the softmax, in place of the path probability. The paper calls this an approximation and gives an intuitive reason for it 1812.00332 §3.2.1 1812.00332 eq. 4. Computing that gradient for every path needs all path outputs, which brings back the memory cost. The paper's appendix describes a backward pass that avoids storing the unused outputs 1812.00332 app. D. The authors' repository also has a mode that keeps every path's output, which uses about N times the memory. The ProxylessNAS card lists the mode differences.

The paper adds hardware cost to the search. The expected latency of a block is the probability-weighted sum of its candidates' latencies, and the expected network latency is the sum over blocks 1812.00332 eq. 5 1812.00332 eq. 6. The loss adds this term with a scale λ2 1812.00332 eq. 7. The paper does not state λ2, so it must be tuned 1812.00332 §3.3.1. Per-candidate latencies come from a predictor described in the appendix 1812.00332 app. B. For objectives that cannot be differentiated, the paper offers a REINFORCE estimator for the architecture parameters 1812.00332 §3.3.2 1812.00332 eq. 8.

The reported CIFAR-10 result is 2.08% test error with 5.7M parameters 1812.00332 tab. 1. On ImageNet under an 80 ms mobile latency target, the paper reports 74.6% top-1 at 78 ms on a Pixel 1, with about 200 GPU hours of search and the GPU type and count not stated 1812.00332 tab. 2. Its specialisation table shows that a model searched for one device is slower on the other two 1812.00332 tab. 4. Card: ProxylessNAS: Direct Neural Architecture Search on Target Task and Hardware.

Posterior-variance pruning (BayesNAS)

BayesNAS keeps a one-shot over-parameterised network but decides what to keep from a posterior variance, not from the size of each architecture weight. Each architecture weight has a zero-mean Gaussian prior. Its variance combines the switches of the edges that feed the same node with the switch of the edge itself, and the variances get hierarchical Gamma hyper-priors [arxiv:1905.04919 §4.1] 1905.04919 §5.1. A Laplace approximation, with a Hessian for each architecture weight, updates the variances 1905.04919 §5.2 [arxiv:1905.04919 alg. 1]. An edge is removed when its variance is at most 1/(2\pi e), the value at which the entropy of the Gaussian reaches zero [arxiv:1905.04919 §5.3]. The name says Bayesian, but this is sparse Bayesian learning over the weights, not Bayesian optimisation of hyperparameters.

The paper's motivation is that magnitude pruning depends on the weight scale, which needs pre-training, and that one-shot methods ignore dependencies between edges 1905.04919 §1 1905.04919 §2. Its dependency-aware prior is the main design claim 1905.04919 §4. The experiments exclude zero operations, and the paper reports no run that isolates the prior's effect 1905.04919 §6.1. The search caches all feature maps for the Hessian, which the authors call inefficient, and memory sets the search batch size 1905.04919 §8 1905.04919 §6.1.

The reported CIFAR-10 result is 2.81% test error with cutout at \lambda_w = 0.005, for about 0.2 GPU-days. A PyramidNet and TreeCell variant reports 2.41% at 0.1 GPU-days 1905.04919 tab. 1. Both are author-reported. Card: BayesNAS: A Bayesian Approach for Neural Architecture Search.

Scan-order search for vision state-space models (LocalMamba)

Here the candidates are not operations on a cell edge but scan orders for a sequence model. Each layer has one learnable vector over eight candidate directions: horizontal and vertical scans, each with its flip, and two local window scans (window 2 and window 7), each with its flip 2403.09338 §4.2. The layer output is the softmax-weighted sum of the eight branch outputs 2403.09338 eq. 5, so the relaxation is the reference one. The supernet trains all eight branches, and afterwards the four most probable directions are kept for each layer, not one 2403.09338 §4.2. Card: LocalMamba: Visual State Space Model with Windowed Selective Scan.

The paper's ablation isolates the search. LocalVim-T reaches 76.2% top-1 on ImageNet-1K with the search, and 75.8% with fixed horizontal and 2×2 local scans and no search 2403.09338 tab. 2. The gain is 0.4 points. The card records no seeds or variance for either number. The paper reports that searched layers favour local scans at the start and end of the network and global scans in the middle, but it gives no quantitative analysis of that pattern 2403.09338 §5.5 2403.09338 fig. 4.

When to use which

The choice depends on the constraints. The rows below follow the papers' own statements, and the costs they report are not on a common hardware basis.

Implementation sketch

The survey gives the alternation, not the update rule. The learning rates and the first- or second-order approximation are not stated in the card.

# Architecture parameters: one weight per edge and candidate op [§4.2.3]
alpha = init_architecture_params(edges, ops)
theta = init_network_weights()

# Alternate the two problems (eq. 5). Step counts and step sizes not stated.
for step in range(T):
    # Outer problem: update alpha on validation data
    loss_val = L_val(theta, alpha, next(val_batches))
    alpha = alpha - eta_alpha * grad(loss_val, alpha)

    # Inner problem: update theta on training data
    loss_tr = L_train(theta, alpha, next(train_batches))
    theta = theta - eta_theta * grad(loss_tr, theta)

# Final cell: keep the highest-weight op on each edge [§4.2.3] [fig. 16]
cell = {e: argmax_op(alpha[e]) for e in edges}

Each edge is a softmax over alpha. For sampled variants, replace the softmax with the Gumbel-softmax of eq. 8 and sample a hard choice per step. Where the two sides disagree on the temperature, see Contradictions below.

ProxylessNAS keeps the same alternation but samples one path per edge from the softmax before each weight step 1812.00332 §3.2.1. It then adds the expected latency of eq. 5 to the training loss 1812.00332 eq. 7. The paper does not give how often α is updated, so an implementation must choose 1812.00332 §3.2.1.

BayesNAS does not follow this loop. Its search trains on the training loss with a re-weighted \ell_1 penalty, updates a per-edge variance from a Hessian, and removes edges below the threshold 1905.04919 alg. 1. The card's sketch of that algorithm is on BayesNAS.

For the LocalMamba variant, the search loop is the same, but the branch output is the softmax-weighted sum of eight scanned branches and the selection is top-k rather than argmax 2403.09338 §4.2. The LocalMamba code normalises each branch output before the weighted sum, a step that eq. 4 does not include multi_mamba.py. Its softmax logits start from a small random initialisation, which the paper does not state multi_mamba.py.

Contradictions

• Temperature direction in the Gumbel relaxation. The survey's text says the relaxation approaches one-hot as \tau \to \infty [arxiv:1908.00709 §4.2.3]. Read from eq. 8, dividing the logits by a large \tau pushes all weights toward uniform, so the one-hot limit is \tau \to 0 [arxiv:1908.00709 eq. 8]. The card's pitfall note takes the equation as the reference. Check the primary paper before you fix an annealing direction.
• DARTS CIFAR-10 error. The text gives 2.83% for DARTS [arxiv:1908.00709 §4.1.2]. The table rows give 97.00% and 97.23% accuracy for first- and second-order DARTS 1908.00709 tab. 3, which is 3.00% and 2.77% error. Neither row matches 2.83%. BayesNAS's table gives 2.76% for the same second-order row 1905.04919 tab. 1. ProxylessNAS's CIFAR-10 table also lists DARTS with cutout at 2.83% 1812.00332 tab. 1, which agrees with the survey's text and not with its table rows.
• DARTS search cost. The survey lists second-order DARTS at 4 GPU-days on four 1080Ti GPUs 1908.00709 tab. 3. The BayesNAS table lists the same row at 1 GPU-day 1905.04919 tab. 1. The two papers do not agree, and the wiki does not resolve it. Check the DARTS paper.
• BayesNAS search cost and row. The survey gives BayesNAS 97.59% accuracy at 0.1 GPU-days 1908.00709 tab. 3. That matches BayesNAS's PyramidNet variant (2.41% error, 0.1 GPU-days), not its main cell (2.81%, 0.2 GPU-days) 1905.04919 tab. 1. The BayesNAS card also states about 3 hours on one GPU for the search 1905.04919 §6.1, which is about 0.125 GPU-days and does not match its 0.2 GPU-day figure. Budget for the larger number until the authors clarify.
• DARTS ImageNet top-5. The survey lists 81.3% top-5 for DARTS searched on CIFAR-10, about 10 points below every other method with similar top-1 1908.00709 tab. 4. ProxylessNAS's ImageNet table lists DARTS at 73.1% top-1 and 91.0% top-5 1812.00332 tab. 3. The two papers disagree on this row. The survey card's own pitfalls note flags its entry as suspicious and says to verify it against the primary paper. Not resolved here.
• Random search versus weight-sharing methods. One cited finding is that random search with weight sharing beats ENAS and DARTS [arxiv:1908.00709 §6.1]. Another cited study finds weight sharing degrades individual architectures and the true ranking 1908.00709 §6.3. A third finds state-of-the-art NAS not significantly better than random search 1908.00709 §4.2.5. These are cited, not reconciled.
• Eq. 7 condition. The survey says \lambda = 0 gives single-level optimisation and then calls the bilevel form "in contrast", without naming the \lambda for that case 1908.00709 §4.2.3. Take the exact update from the MiLeNAS paper.
• Costs. The GPU-days column combines different GPU models and counts 1908.00709 tab. 3 1908.00709 eq. 9. The survey's own table reading is that these costs are not directly comparable [arxiv:1908.00709 §6.1.2]. ProxylessNAS reports its ImageNet cost in GPU hours on hardware it does not name 1812.00332 tab. 2, so it is not on the same basis as the survey's GPU-day rows.

Papers

• AutoML: A Survey of the State-of-the-Art: survey that places DARTS and its variants in the gradient-based branch of NAS, with equations for the bilevel, single-level, mixed-level and sampled forms.
• BayesNAS: A Bayesian Approach for Neural Architecture Search: one-shot search that prunes by posterior variance under a hierarchical sparse prior, with a dependency-aware prior and author-reported CIFAR-10 results.
• ProxylessNAS: Direct Neural Architecture Search on Target Task and Hardware: searches on ImageNet and CIFAR-10 directly with one sampled path per edge in memory, and adds an expected-latency term to the loss for a chosen device.
• LocalMamba: Visual State Space Model with Windowed Selective Scan: applies a softmax relaxation over eight scan orders per layer, keeps the top four, and reports author-reported ImageNet, COCO and ADE20K results.