Source-linked AI summary
BayesNAS: A Bayesian Approach for Neural Architecture Search
Hongpeng Zhou, Minghao Yang, Jun Wang, Wei Pan
TL;DR
One-shot NAS reduces search time but existing methods can mishandle operation dependencies and prune using questionable magnitude criteria. BayesNAS addresses these issues with HARD priors, finds architectures after one training epoch, and achieves competitive CIFAR-10 and transferable ImageNet performance while supporting sparse networks.
Problem
Existing one-shot NAS methods may disregard dependencies around zero operations and rely on architecture-parameter magnitude for pruning, despite magnitude depending on weight scale and not necessarily identifying optimal edges.
Method
BayesNAS models architecture parameters with hierarchical automatic relevance determination priors and uses Bayesian uncertainty for pruning, with Laplace approximation for inference.
Results
BayesNAS finds a CIFAR-10 architecture in 0.2 GPU days using one GPU and one training epoch, achieves competitive test error with fewer parameters, and transfers competitively to ImageNet.
Takeaways & Limitations
The approach directly learns sparse neural architectures and can compress convolutional neural networks by enforcing structural sparsity without accuracy deterioration.
Takeaways & Limitations
Exact Bayesian posterior and predictive computations are generally intractable, so the method relies on approximate inference and assumes data likelihoods belong to the exponential family.
Abstract
from arXiv · showhide
One-Shot Neural Architecture Search (NAS) is a promising method to significantly reduce search time without any separate training. It can be treated as a Network Compression problem on the architecture parameters from an over-parameterized network. However, there are two issues associated with most one-shot NAS methods. First, dependencies between a node and its predecessors and successors are often disregarded which result in improper treatment over zero operations. Second, architecture parameters pruning based on their magnitude is questionable. In this paper, we employ the classic Bayesian learning approach to alleviate these two issues by modeling architecture parameters using hierarchical automatic relevance determination (HARD) priors. Unlike other NAS methods, we train the over-parameterized network for only one epoch then update the architecture. Impressively, this enabled us to find the architecture on CIFAR-10 within only 0.2 GPU days using a single GPU. Competitive performance can be also achieved by transferring to ImageNet. As a byproduct, our approach can be applied directly to compress convolutional neural networks by enforcing structural sparsity which achieves extremely sparse networks without accuracy deterioration.
1. Introduction
One-shot NAS reduces search cost by sharing weights in an over-parameterized supergraph, but existing methods can mishandle dependencies, zero operations, and magnitude-based pruning. BayesNAS addresses these issues with Bayesian sparse priors and a variance-aware pruning procedure.
- Motivation: One-shot NAS treats candidate architectures as subgraphs of a weight-sharing supergraph, reducing search time without separate training.This contrasts with reinforcement-learning and neuro-evolutionary NAS methods, which require enormous computational power.
- Limitations: Existing one-shot methods can disregard node dependencies, producing isolated or dysfunctional graph structures when zero operations dominate edges.The paper notes that post-processing disconnected graphs can require extensive extra computation.
- Limitations: Magnitude-based pruning is questionable because architecture-parameter magnitude depends on weight scale and does not necessarily identify the optimal edge.The paper connects this issue to sensitivity to architectural choices and the need for pre-training in network compression.
- Approach: BayesNAS models architecture parameters with hierarchical automatic relevance determination priors that promote sparsity and encode dependencies between nodes.The resulting uncertainty provides a pruning metric beyond parameter magnitude.
- Approach: BayesNAS uses posterior variance in an iteratively re-weighted ℓ1-type algorithm and trains the model for only one epoch before updating the architecture.The posterior-variance update uses Laplace approximation and Hessian computation.
- Implications: BayesNAS also extends directly to network compression by enforcing structural sparsity, producing extremely sparse models with minimal or no accuracy loss across tested architectures.The sparse-kernel extension can be integrated with architecture search for resource-limited hardware.
2. Related Work
Related work spans magnitude-based pruning for network compression and one-shot NAS, alongside Bayesian approaches that use sparsity-inducing priors for neural-network modeling.
- Network Compression: Network-compression methods commonly prune weights by magnitude, often requiring repeated pruning and learning steps.Penalty-based sparsity is a conventional route to removing weights below a threshold.
- One-shot NAS: DARTS ranks operations using a softmax over architecture-parameter magnitudes, while ProxylessNAS binarizes parameters and selects the highest-weight edge.These methods use magnitude-derived criteria to determine retained operations or edges.
- BayesNAS Context: Figure 1 contrasts disconnected graphs caused by dependency disregard with dependency-aware switches and prioritized zero operations.The illustration also presents a multi-input-multi-output motif as a graph building block.
- One-shot NAS: DSO-NAS relaxes the ℓ0 objective with an ℓ1 penalty and gradually prunes edges using a threshold during training.The regularization strength is tied to the learning rate and a predefined parameter.
- Bayesian Learning: Bayesian neural-network methods include Laplace approximation, Hamiltonian Monte Carlo, expectation propagation, and variational inference.BayesNAS adopts Laplace approximation despite the computational challenge of inverse-Hessian calculation for large networks.
3. Search Space Design
BayesNAS represents the search space as a directed acyclic graph whose edges carry operations and scaling parameters, then adds switches to encode dependency-aware edge activation.
- DAG Representation: The NAS search space is represented as a directed acyclic graph with edge e_ij denoting a connection from node i to node j.The search space determines which architectures can be discovered in principle.
- Architecture Parameters: Scaling scalars w_ij are assigned over operations on edges to control information flow and act as architecture parameters.The mixed-operation outputs and node representations are built from these edge-associated quantities.
- Optimization Objective: The search objective is to learn a sparse subgraph that maintains or improves the accuracy of the over-parameterized DAG under a desired nonzero-edge budget.The formulation separates network parameters from edge parameters and constrains sparsity with an ℓ0 norm.
- Dependency-Aware Switches: Each architecture parameter receives a switch whose on-off state is not determined solely by magnitude and whose predecessor controls successor connectivity.The switch is computed directly rather than updated by gradient descent and supports proxy and proxyless settings.
- Motif Design: A motif serves as a general multi-input-multi-output building block in which successors may have multiple predecessors and edges may carry primitive operations.Convolution and pooling are examples of edge-associated operations.
4. Dependency Based One-Shot Performance Estimation Strategy
BayesNAS encodes valid information flow probabilistically so successor edges depend on predecessor activity, then uses switches and sparse Bayesian learning to identify redundant connections.
- Dependency Proposition: Information flows from node j to k when the successor operation is nonzero and at least one predecessor operation supplying node j is nonzero.The equivalent zero-flow condition requires all predecessor operations or the successor weight to be zero.
- Dependency Proposition: This dependency criterion explains why nonzero edges can become dysfunctional and why an unconstrained sparsity formulation is incomplete.The paper illustrates the issue with edges from node 2 to nodes 3 and 4.
- Probabilistic Encoding: BayesNAS represents dependency logic through probability distributions over architecture parameters and switches that encode whether edges are active or redundant.The paper notes that multiple expressions can encode the same proposition.
- Switch Construction: Continuous switches use infinity for ON and zero for OFF, with addition and multiplication representing set union and intersection; switches encode uncertainty rather than weight magnitude.Gaussian distributions provide a direct way to encode this logic.
- Sparse Bayesian Learning: Automatic relevance determination and sparse Bayesian learning supply sparsity-inducing priors because Gaussian distributions alone generally do not promote sparsity.The method uses these Bayesian techniques to obtain sparse probabilistic solutions.
- Zero-Operation Handling: To prioritize zero operations, BayesNAS inserts an auxiliary identity node between connected nodes so switching it off disables downstream operations.The inserted edge weight and switch are trainable and initialized to 1.
5. Bayesian Learning Search Strategy
BayesNAS formulates architecture and network parameters probabilistically, using hierarchical priors and approximate Bayesian inference to promote sparsity and support architecture search and compression.
- Bayesian Neural Network: The likelihood models network weights W and noise precision σ^-2 given data D = (X, Y).The model assumes the data likelihood belongs to the exponential family.
- Bayesian Neural Network: Gaussian priors are assigned to entries of the network weight matrices, with hyperparameters controlling their variances.The formulation introduces a Gaussian prior for each entry in each weight matrix.
- Bayesian Neural Network: Gamma hyperpriors make the latent-variable marginal prior Student's t and promote sparsity.The analysis uses uniform scale priors by fixing the Gamma parameters a and b to zero.
- Bayesian Neural Network: The posterior over parameters is obtained with Bayes' rule, while predictions use a predictive distribution for new inputs.Exact posterior and predictive computations are generally intractable, requiring approximate inference methods.
- Search Algorithm: The search algorithm jointly trains weights and architecture parameters, computes architecture-parameter Hessians, updates variables, prunes parameters, and then trains the pruned network.The framework also supports proxy tasks and structural sparsity for network compression.
6. Experiments
BayesNAS is evaluated in proxy and proxyless NAS settings, using Bayesian architecture evaluation and pruning to learn compact convolutional cells and paths. On CIFAR-10, it achieves competitive accuracy with fewer parameters, while proxyless search substantially reduces model size despite a slight test-error increase.
- Proxy Search: BayesNAS evaluates proxy NAS by searching convolutional cells on CIFAR-10 and retraining a network assembled from the learned cells.The search follows DARTS and SNAS, then trains a larger stacked-cell network for comparison.
- Proxyless Search: In proxyless NAS, BayesNAS searches the optimal paths within PyramidNet tree-cells and reconstructs the network by retaining those paths.The tree-cell has nine possible paths, and each bottleneck block replaces its 3 × 3 convolution with this structure.
- Proxy Search: γ replaces validation accuracy as the performance criterion, enabling BayesNAS to update architecture parameters in a one-shot manner.This distinguishes the search procedure from DARTS and SNAS, which use validation accuracy during or after search.
- Proxy Search: The search space stacks seven-node convolutional cells, with reduction cells at one-third and two-thirds of network depth and no zero operations.Cell inputs connect to outputs from the preceding two cells, while outputs concatenate intermediate nodes.
- Proxy Search: 3 hours on a single GPU: the CIFAR-10 search trains an eight-cell network with cached feature maps and batch size 18.The experiments used an NVIDIA TITAN V GPU.
- Proxy Search: BayesNAS finds CIFAR-10 cells with competitive test error and fewer parameters than DARTS and SNAS.The learned normal and reduction cells are shown in Figure 3, and the retrained large network uses 20 cells.
- Proxyless Search: Proxyless search yields chain-like cells with one path or inception cells with divergence and convergence, while some cells are dispensable.Although test error increases slightly relative to ProxylessNAS, the number of learned model parameters drops significantly.
7. Transferability to ImageNet
BayesNAS transfers cells learned on CIFAR-10 to ImageNet in the mobile setting. The transferred cells achieve competitive ImageNet performance under the reported training configuration.
- Transferability to ImageNet: Cells learned on CIFAR-10 transfer to ImageNet and achieve competitive performance in the mobile setting.The ImageNet network uses 14 cells, 224×224 inputs, and 250 training epochs.
8. Conclusion and Future Work
BayesNAS learns sparse neural network architectures while reducing search time to one epoch, with current implementation limitations noted for Hessian computation and scaling.
- BayesNAS directly learns a sparse neural network architecture.
- One epoch is used to obtain the candidate architecture, significantly reducing search time.
- Caching all feature maps for Hessian computation currently makes the implementation inefficient.
- Hessian computation could instead be performed with backpropagation to reduce search time and scale to larger search spaces.
A.1. Algorithm Derivation
The derivation formulates Bayesian architecture-parameter learning as iterative optimization using sparse priors, Laplace-based uncertainty estimates, and dependency-aware updates.
- The likelihood and sparse hierarchical prior are combined into an optimization problem over architecture parameters and hyperparameters.
- Laplace approximation uses a Taylor expansion with gradient and Hessian terms to approximate the posterior distribution.
- The optimization alternates between estimating architecture parameters and updating their associated hyperparameters.
- Dependency-aware Gaussian priors incorporate relationships between edges when computing architecture-parameter uncertainty.
- The parameter update becomes an ℓ1-regularized regression-type optimization, summarized in Algorithm 1.
- The broader algorithm jointly trains weights and architecture parameters, computes Hessians, updates variables, prunes parameters, and standardly trains the pruned network.
B.2. Experiments
The experiments evaluate structured sparsity on MNIST and CIFAR-10 models, while accompanying procedures define Hessian computation and compression settings.
- Experiments: LeNet-300-100 on MNIST learns a 465 −37 −90 structure with 1.54% test error and 0.04 FLOPS.
- Experiments: Structured sparsity examples for convolutional filters are illustrated in Figure S5.
- Experiments: ResNet-18 on CIFAR-10 applies shape-wise and filter-wise regularization to convolutional layers and row-wise and column-wise regularization to the fully connected layer.
- Hessian computation: The approximate Hessian method reduces computation from 107.97 MMACs to 0.04 MMACs for n = 100 and m = 100.
C.2. Compute the Hessian of Conv Layer
The convolutional-layer Hessian is computed recursively by converting convolutions to equivalent fully connected operations, then approximated using diagonal information to reduce cost.
- Explicit Hessian computation is computationally intensive, especially for convolutional layers, motivating a recursive equivalent fully connected formulation.
- Convolutional inputs, outputs, and weights are reshaped into fully connected representations for recursive Hessian computation.
- The recursive procedure propagates pre-activation Hessians between layers and converts results back to convolutional dimensions.
- The architecture-parameter Hessian is computed from input features and downstream pre-activation Hessians, using their independence for efficiency.
D.1. Architecture Search on CIFAR-10
The CIFAR-10 architecture search uses standard image augmentation, a defined operation set, and SGD-based training settings, with BayesNAS producing normal and reduction cells.
- Data Pre-processing and Augmentation Techniques: Training images are centrally padded, randomly cropped, horizontally flipped, and channel-normalized.Cutout is also used as an additional enhancement.
- Implementation Details of Operations: The search space includes separable and dilated convolutions, max and average pooling, and skip connections.Operations generally use stride one, while operations adjacent to input nodes in reduction cells use stride two.
- Detailed Training Settings: Network parameters are optimized with momentum SGD using learning rate 0.1, momentum 0.9, weight decay 1 × 10−4, batch size 16, and 16 initial channels.
- Detailed Training Settings: A separate training configuration uses batch size 128, SGD, weight decay 3 × 10−4, momentum 0.9, and cosine-annealed learning rate 0.1.
- Architecture Search Results: BayesNAS finds both a normal cell and a reduction cell for the CIFAR-10 architecture.