BayesNAS: A Bayesian Approach for Neural Architecture Search

In plain words

Finding a good neural network design by trying many designs is slow. One shortcut builds one large network that contains every candidate design and lets training decide which parts to keep. Existing versions of this shortcut often judge a part as unimportant because its weight is small, and they can cut away a piece that the rest of the network still needs.

BayesNAS keeps the same shortcut but uses a statistical view. Each candidate connection gets an uncertainty value, and a connection is kept or dropped depending on that uncertainty, not on its raw weight. The uncertainty of a connection also depends on the connections that feed into the same point, so a dead feeder can switch off the connections it would have powered. Training runs for only a short time before the design is chosen, which makes the search quick.

The same idea can also shrink an ordinary trained network by removing whole groups of weights. The card's results are the authors' own numbers and have not been checked by anyone else.

TL;DR

BayesNAS is a one-shot architecture search. It places a hierarchical sparse (ARD-style) prior on each architecture weight, tied to the weights of predecessor edges, and uses a Laplace approximation with a cheap Hessian to update the priors. After one epoch of training, edges whose prior variance is at most 1/(2πe) are removed §5.3. Use it if you need a fast, gradient-based search with a principled prune rule; the paper reports CIFAR-10 search in about 0.2 GPU days and no independent replication is in the wiki.

Problem

Reinforcement-learning and evolutionary NAS methods cost very large amounts of compute §1. One-shot NAS methods are cheaper, but the paper identifies two problems with them §1:

• Dependencies are ignored. A node's outgoing edges can stay non-zero while the node has no non-zero input, so the sampled graph is disconnected. Zero operations therefore receive improper treatment §1 §4.
• Pruning by architecture-weight magnitude is questionable. Magnitude depends on weight scale, which needs pre-training and is sensitive to the architecture choice §1 §2.

The paper also notes that the probability of sampling a disconnected graph is large in one cited example, which motivates the dependency handling §1.

Method

BayesNAS treats architecture search as sparse Bayesian learning over a DAG-shaped over-parameterised network §3. Its steps:

1. Mixed operations. Each edge i \to j carries a scalar weight w_{ij}^o per candidate operation, and the node output is a weighted sum of operation outputs eq. 1 §3. The search space excludes zero operations in the experiments, so zero is not a candidate §6.1.
2. Dependency-aware prior. A switch s_{jk}^{o'} is attached to each architecture weight. The prior on w_{jk}^{o'} is a zero-mean Gaussian whose variance \gamma_{jk}^{o'} combines the switches of all predecessor edges i < j with the switch of the edge itself eq. 6 §4.1. Equation 7 is the simplified form the paper uses in the prior eq. 7 §4.1. To represent "no operation" without a zero operation, a single identity node i' is inserted before j, with its own switch eq. 8 §4.2.
3. Hierarchical ARD prior. The variances get Gamma hyper-priors, which gives a Student-t marginal and promotes sparsity. The hyper-prior parameters are fixed to zero, meaning flat priors §5.1.
4. Laplace approximation. The data term is approximated by a second-order expansion around the current weights, with Hessian \mathbf{H} of the energy with respect to each architecture weight §5.2 app. S1.1. The paper's Proposition and Corollary derive the cost that is minimised app. S1.1.
5. One search iteration (Algorithm 1 in the paper, reproduced below). Train the network and architecture weights on a re-weighted \ell_1 loss; compute the Hessian; update the posterior variance; derive the re-weighting coefficient; derive the switch; update the variance; prune edges whose variance is at most 1/(2\pi e); then fix the architecture weights to 1 and train the pruned network in the standard way alg. 1 §5.3.
6. Pruning rule. The entropy of a zero-mean Gaussian with variance \gamma is \tfrac{1}{2}\ln(2\pi e\gamma). An edge is removed when this entropy is non-positive, which is \gamma \le 1/(2\pi e) \approx 0.0585 §5.3.

Key equations (numbers as in the compiled paper):

o_j(z_i) = \sum_{o \in \mathcal{O}} w_{ij}^o \, o_{ij}(z_i), \qquad z_j = \sum_{i<j} o_j(z_i) eq. 1 §3

\gamma_{jk}^{o'} = \left( \frac{1}{\sum_{i<j} \sum_{o} s_{ij}^{o}} + \frac{1}{s_{jk}^{o'}} \right)^{-1} eq. 6 §4.1

\gamma_{jk}^{o'} = \left( \frac{1}{s_{ii'}} + \frac{1}{\sum_{i'<j} \sum_{o} s_{i'j}^{o}} + \frac{1}{s_{jk}^{o'}} \right)^{-1} eq. 8 §4.2

\mathbf{C}_{jk}^{o'}(t) = \left( \frac{1}{\gamma_{jk}^{o'}(t-1)} + \mathbf{H}_{jk}^{o'}(t) \right)^{-1} alg. 1

\omega_{jk}^{o'}(t) = \frac{\sqrt{\gamma_{jk}^{o'}(t-1) - \mathbf{C}_{jk}^{o'}(t)}}{\gamma_{jk}^{o'}(t-1)}, \qquad s_{jk}^{o'}(t) = \left| \frac{w_{jk}^{o'}(t)}{\omega_{jk}^{o'}(t)} \right| alg. 1 §5.3

The compiled loss for the architecture search is a maximum-likelihood term plus a re-weighted \ell_1 penalty on the architecture weights and an \ell_2 penalty on the network weights, \mathcal{L}_D = E_D + \lambda_w \sum \|\omega \, w\|_1 + \lambda \|\mathcal{W}\|_2^2 §5.3. The proxy-task variant replaces the \ell_1 term with a group \ell_2 term over repeated cells app. S1.3.

# BayesNAS search on one over-parameterised network  [alg. 1] [§5.3]
# Input:  training data; sparsity intensity lambda_w; lambda = 0.01
# Output: boolean mask of edge operations to keep

def bayesnas_search(data, lambda_w, num_iters):
    # Initialise architecture weights, variances and l1 weights to 1  [alg. 1]
    arch = ones(num_edge_ops)
    gamma = ones(num_edge_ops)
    omega = ones(num_edge_ops)
    for t in range(1, num_iters + 1):
        # Train network and architecture weights on the re-weighted loss;
        # the inner schedule is not stated  [§5.3] [alg. 1]
        net, arch = sgd_step(data, net, arch, lambda_w * omega)
        # Hessian of the energy for each architecture weight, from cached
        # feature maps  [alg. 1] [app. S3]
        hessian = architecture_hessian(net, arch, data)
        # Posterior variance, using the previous variance as the prior  [alg. 1]
        posterior = 1 / (1 / gamma + hessian)
        # Re-weighting coefficient for the next l1 step  [alg. 1]
        omega = math.sqrt(gamma - posterior) / gamma
        # Switch value of each architecture weight  [alg. 1]
        switch = abs(arch / omega)
        # Sum of switches over predecessor edges of each edge  [§4.1]
        pred_sum = sum_predecessor_switches(switch)
        # Variance of each edge prior, Eq. 6 or Eq. 8  [§5.3]
        gamma = 1 / (1 / pred_sum + 1 / switch)
    # Remove edges whose maximum entropy is non-positive  [§5.3]
    keep = gamma > 1 / (2 * math.pi * math.e)
    # Fix architecture weights to 1 and retrain the pruned net  [alg. 1]
    return keep

Notes on the sketch. The loop variable t is an iteration index, and the paper does not state how many iterations a search runs in the CIFAR-10 experiments §6.1. The gamma used in the Hessian step is the previous iteration's value, as in alg. 1.

Proxy and proxyless variants. For proxy search (cells shared across the network), the switches and omegas are tied within a group of identical cells, and the group variance uses the combined norm of the group app. S1.3. For proxyless search, each cell is searched separately §6.2.

Structural compression. The same machinery prunes groups of conv or FC weights (shape-wise, row-wise, column-wise, filter-wise, etc.) by a reweighted group Lasso, with the Hessian computed layer by layer app. S2 app. S2.1.

Hessian approximations. The exact Hessian of an FC layer is replaced by its diagonal in the recursive formula app. S3.1. For conv layers the approximation replaces inputs and pre-activation Hessians by their means over the feature map app. S3.2. For each architecture weight the Hessian is the squared mean of the absolute input times the pre-activation Hessian of the target node app. S3.3.

Hyperparameters

Results

Author-reported. Setup: CIFAR-10 test error with cutout for the search and retraining protocols in the hyperparameter table above; ImageNet in the mobile setting with the cell searched on CIFAR-10 §7. The paper does not say whether baseline numbers were re-run; they are the values cited in each table row.

CIFAR-10, subset of tab. 1:

ImageNet (mobile setting), from tab. 2:

Structural compression app. S2.2, baseline and proposed as reported:

The paper's own summary is that the cell searched on CIFAR-10 transfers to ImageNet with competitive error and fewer parameters than DARTS and SNAS §6.1 §7. The paper reports no run-to-run variance for the ImageNet rows.

Paper vs code

No official code. Searched: (1) the paper's TeX source has no URL or project page, and the arXiv metadata has no comments field; (2) no project page is linked; (3) GitHub repository search for "BayesNAS" returned BayesNAS/BayesNAS (Apache-2.0; README says it is "code for ICML paper BayesNAS") and Contantine/BayesNA-display (a UI, no license). BayesNAS/BayesNAS belongs to a user account named BayesNAS with no name, company or bio, and the README names no author, so it cannot be verified as the authors' repository. It is not used here and was not cloned. Nothing in this card is checked against code.

Pitfalls

• Search-cost inconsistency. The abstract and tab. 1 say 0.2 GPU days. The experiments section says the search "takes about 3 hours on a single GPU" §6.1. Three hours is about 0.125 GPU days. Budget for the larger figure until the authors clarify.
• Batch size and learning rate conflict. The search batch is 18 in the main text §6.1 and 16 in the appendix app. S4.1. The search learning rate is "fixed 0.1" in the main text and "initial 0.1" in the appendix, so a schedule is possible but not stated.
• Iteration count is not stated. The introduction says the network is trained for "only one epoch" before the architecture update §1. Algorithm 1 loops over t with "several epochs" per iteration §5.3, and the CIFAR-10 search has no stated T_{\max}.
• Variance, not standard deviation. The prior is written \mathcal{N}(w \mid 0, s) and the derivation uses \log|s| as a log-determinant app. S1.1. That only works if s and \gamma are variances. The paper does not say so in words.
• Two forms of the switch update. Algorithm 1 uses s = |w/\omega| alg. 1; the derivation gives s = w/\omega app. S1.2.2 and then uses the absolute value app. S1.2.2. Use the absolute value, which keeps the variance positive.
• Hessian-based posterior uses the previous variance. \mathbf{C}(t) is built from \gamma(t-1) and the Hessian at step t alg. 1. Using \gamma(t) instead changes the update.
• Penalty notation. The loss writes \lambda_w §5.3 and the tables write \lambda_w^o. The paper does not say whether they are the same parameter, or whether \lambda_w is one value per operation or per edge. The three CIFAR-10 settings differ only in this value tab. 1.
• Zero operations. The main contribution in §4 is a dependency-aware zero treatment. The experiments exclude zero operations §6.1. The paper does not report an experiment that isolates the effect of §4's dependency prior.
• Which Hessian was used. The appendix gives an exact recursive Hessian and two approximations (diagonal, and mean over feature maps) app. S3.1 app. S3.2. The paper does not say which one the experiments use.
• Garbled formula. The architecture-weight Hessian in the approximation \mathbf{H}_{ij}^o = (\mathbb{E}(|B_i|))^2 H_j is missing a closing parenthesis in the source app. S3.3. The implied form is the squared mean absolute input times H_j; check against the derivation before implementing.
• Operation count claim. The paper says the exact FC-layer Hessian cost is n(2m^2+2n^2+4mn+3m-1) multiply-accumulates and the approximate cost is n(2+4m). With n = m = 100 the formulas give about 8.03 M and 0.04 M respectively, while the text states 107.97 M for the exact case app. S3.1. The approximate figure matches; the exact figure does not.
• Memory. The current implementation caches all feature maps in memory to compute the Hessian, and the authors call this inefficient §8. Memory, not compute, sets the batch size (18) in the search §6.1.
• Retraining is separate. The pruned network is retrained from scratch with its own schedule §6.1 §6.2 §7. Search-time accuracy is not the reported accuracy.
• Compression results have no variance. The MNIST and CIFAR-10 structural results report no spread across runs app. S2.2. The ResNet-18 schedule and per-layer \lambda are not stated.

Lineage

• Builds on: one-shot and gradient-based NAS (DARTS, SNAS, ProxylessNAS, DSO-NAS cited in §2), and classic Bayesian learning with automatic relevance determination and sparse Bayesian learning §2 §4.1.
• Compared against: NASNet, AmoebaNet, ENAS, PNAS, DARTS, SNAS, DSO-NAS, ProxylessNAS, as in tab. 1 and tab. 2.
• Differs from DARTS-style methods by pruning on a posterior variance rather than on architecture-weight magnitude §1 §5.3.