Source-linked AI summary

SNAS: Stochastic Neural Architecture Search

Sirui Xie, Hehui Zheng, Chunxiao Liu, Liang Lin

arXiv:1812.09926v3cs.LGcs.AIstat.ML

TL;DR

SNAS addresses the efficiency and differentiability trade-off in neural architecture search by jointly learning operation and architecture-distribution parameters through back-propagation. It reports strong CIFAR-10 results with fewer parameters or resources, and the searched result transfers to ImageNet.

  • Problem

    Neural architecture search must trade off search efficiency, result optimality, and feasible resource use while preserving a complete differentiable pipeline.

  • Method

    SNAS reformulates cell architecture search as optimization of a factorized joint distribution, learning architecture parameters and operation parameters together with a differentiable stochastic relaxation and resource regularization.

  • Results

    SNAS achieves strong CIFAR-10 performance with fewer parameters or computation resources than several NAS baselines, while the searched result transfers to ImageNet.

  • Takeaways & Limitations

    SNAS provides an economical, less-biased candidate for efficient NAS on large datasets by using more efficient credit assignment and locally decomposable resource constraints.

  • Takeaways & Limitations

    SNAS assumes the architecture distribution is fully factorizable, rather than requiring the Markov Decision Process assumption used in ENAS.

Abstract

from arXiv · show

We propose Stochastic Neural Architecture Search (SNAS), an economical end-to-end solution to Neural Architecture Search (NAS) that trains neural operation parameters and architecture distribution parameters in same round of back-propagation, while maintaining the completeness and differentiability of the NAS pipeline. In this work, NAS is reformulated as an optimization problem on parameters of a joint distribution for the search space in a cell. To leverage the gradient information in generic differentiable loss for architecture search, a novel search gradient is proposed. We prove that this search gradient optimizes the same objective as reinforcement-learning-based NAS, but assigns credits to structural decisions more efficiently. This credit assignment is further augmented with locally decomposable reward to enforce a resource-efficient constraint. In experiments on CIFAR-10, SNAS takes less epochs to find a cell architecture with state-of-the-art accuracy than non-differentiable evolution-based and reinforcement-learning-based NAS, which is also transferable to ImageNet. It is also shown that child networks of SNAS can maintain the validation accuracy in searching, with which attention-based NAS requires parameter retraining to compete, exhibiting potentials to stride towards efficient NAS on big datasets. We have released our implementation at https://github.com/SNAS-Series/SNAS-Series.

1 INTRODUCTION

SNAS addresses the computational cost and limited gradient use of NAS by jointly learning operation parameters and architecture distributions through differentiable search. It reports strong CIFAR-10 and ImageNet results with reduced search resources and competitive model size.

  • SNAS jointly trains neural operation parameters and architecture distribution parameters in one back-propagation round while preserving a complete, differentiable NAS pipeline.
  • SNAS replaces constant reinforcement-learning rewards with gradient feedback from generic differentiable loss and reformulates architecture search using a stochastic model.
  • 2.85±0.02% CIFAR-10 test error with 2.8M parameters outperforms 1st-order DARTS at 3.00±0.14%-3.3M and ENAS at 2.89%-4.6M.
  • 88% validation accuracy during search exceeds ENAS’s around 70% in fewer epochs, and SNAS maintains search validation accuracy without child-network finetuning.
  • 27.3% top-1 error after transfer to ImageNet is comparable to 26.9% from 2nd-order DARTS.

2 METHODOLOGY

SNAS models cell architectures as sampled graph structures and makes their optimization differentiable through a relaxed joint distribution. Its search gradient assigns local structural credit and incorporates resource cost into the objective.

  • 2.1 SEARCH SPACE AND ARCHITECTURE SAMPLING: A cell is represented as a directed acyclic parent graph whose edges select operations between ordered latent-representation nodes.
  • 2.1 SEARCH SPACE AND ARCHITECTURE SAMPLING: SNAS jointly searches operations and topology by adding a zero operation that can remove an edge, while cells use two previous-cell inputs and concatenate intermediate nodes.
  • 2.1 SEARCH SPACE AND ARCHITECTURE SAMPLING: The architecture distribution p(Z) is fully factorizable and parameterized by α, which is learned alongside operation parameters θ without requiring an MDP assumption.
  • 2.1 SEARCH SPACE AND ARCHITECTURE SAMPLING: SNAS uses training/testing loss directly as reward, optimizing expected performance over architectures sampled from p(Z) rather than DARTS’s analytical edge-wise expectation.
  • 2.2 PARAMETER LEARNING FOR OPERATIONS AND ARCHITECTURES: Concrete-distribution relaxation makes discrete architecture sampling continuous and differentiable through reparameterization; as temperature approaches zero, the relaxation becomes unbiased once converged.
  • 2.3 CREDIT ASSIGNMENT: SNAS’s search gradient is equivalent to a policy gradient while assigning each structural decision credit through differentiable loss contributions rather than delayed rewards.
  • 2.3 CREDIT ASSIGNMENT: At each edge, continuous masks initially share credit among operations; annealing toward one-hot masks progressively concentrates credit on selected operations and updates their sampling probabilities.
  • 2.4 RESOURCE CONSTRAINT: The objective adds ηC(Z) to expected loss, where C(Z) approximates child-network time using parameter size, FLOPs, or memory access cost.

3 EXPERIMENTS

SNAS searches convolutional cells on CIFAR-10, evaluates their child networks, and transfers the learned cells to ImageNet. Across these experiments, it combines competitive accuracy with lower model complexity and search resources, while resource constraints produce smaller architectures.

  • Architecture Search on CIFAR-10: SNAS searches seven-node convolutional cells with three resource-constraint levels for 150 epochs on a single GPU.The search network uses eight stacked cells and single-level optimization of operation and architecture parameters.
  • Architecture Search on CIFAR-10: SNAS takes fewer epochs than ENAS to converge to higher validation accuracy, while DARTS converges faster but its search accuracy is inconsistent with the child network.SNAS maintains performance after architecture derivation, whereas DARTS exhibits a large gap.
  • Architecture Search on CIFAR-10: Higher architecture-distribution entropy in DARTS makes operation and edge derivation substantially alter the continuous network output.The paper attributes the gap to removing weak operations and manually selecting ambiguous edges after search.
  • Architecture Search on CIFAR-10: Increasing the resource constraint sparsifies SNAS cells, with aggressive constraints producing smaller reduction-cell structures that ENAS and DARTS cannot discover.Under mild constraints, each node uses two input edges; stronger constraints further sparsify reduction cells.
  • Architecture Evaluation on CIFAR-10: 2.85±0.02% test error with 2.8M parameters makes SNAS better than 1st-order DARTS and ENAS, while using three orders of magnitude fewer computation resources than RL- and evolution-based NAS.SNAS is also comparable to 2nd-order DARTS with fewer parameters, and aggressive constraints retain performance comparable to 1st-order DARTS.
  • Architecture Transferability Evaluation on ImageNet: SNAS cells learned on CIFAR-10 transfer successfully to ImageNet under a mobile setting restricting multiply-add operations to less than 600M.The ImageNet evaluation uses 224 × 224 inputs and reports competitive performance with substantially fewer search resources than RL-based NAS.

4 RELATED WORKS

SNAS targets more efficient NAS by combining gradient-based architecture learning with an end-to-end pipeline. It draws on ideas from parameter sharing and continuous structural decisions while using generic differentiable loss for architecture updates.

  • SNAS follows ENAS in using parameter sharing among all possible child graphs to accelerate NAS.
  • SNAS uses a continuous distribution for structural decisions at each edge and optimizes it with an l0 complexity regularizer.
  • SNAS leverages generic differentiable loss gradients to update architecture distributions while preserving the complete NAS pipeline.
  • The paper combines insights from generative modeling and reinforcement learning to discuss efficiency improvements for NAS.

5 CONCLUSION

SNAS is presented as an economical end-to-end NAS framework that improves efficiency through gradient-based architecture search and complexity regularization. Experiments report strong CIFAR-10 results, transferability to ImageNet, and a possible path toward NAS on larger datasets.

  • SNAS uses gradient information from generic differentiable loss without sacrificing the completeness of the NAS pipeline.
  • Complexity regularization lets SNAS trade off testing error and forwarding time.
  • SNAS searches well on CIFAR-10, and its result can be transferred to ImageNet.
  • The authors position SNAS as a possible candidate for full-fledged NAS on large datasets because it is more efficient and less biased.

B DIFFERENCE BETWEEN SNAS AND DARTS

SNAS optimizes the expected loss of sampled architectures with a factorized stochastic distribution and differentiable relaxation. Unlike DARTS’s analytical expectation, this preserves the sampled-architecture objective and avoids the inconsistency caused by nonlinear operations.

  • The example search space contains three intermediate nodes, with nodes arranged in series and parallel to illustrate SNAS–DARTS differences.
  • The NAS objective is the expected reward E_Z∼pα(Z)[R(Z)], where pα(Z) is the architecture distribution.
  • SNAS uses the differentiable loss R(Z) = Lθ(Z) as reward, allowing architecture and operation parameters to be trained together.
  • SNAS factorizes p(Z) without relying on the MDP assumption, while preserving the stated objective without introducing bias under its independent-edge assumption.
  • DARTS replaces sampling with analytical expectations over incoming-edge operations and optimizes a relaxed loss with deterministic gradients.
  • Because the operations are nonlinear ReLU-Conv-BN stacks, DARTS’s transformation introduces unbounded bias and can require parameter retraining after architecture derivation.
  • SNAS represents each edge’s operation choice with one-hot stochastic variables and relaxes sampling through a concrete distribution and reparameterization.
  • The SNAS search gradient is derived for architecture parameters at each edge, with successor-node gradients stopped under an enforced independence assumption.

E TAYLOR DECOMPOSITION FOR CONTRIBUTION ANALYSIS

The section develops contribution analysis for SNAS by decomposing network outputs into node-level credits and defining resource measures for candidate operations. It concludes that parameter size, FLOPs, and MAC are jointly needed to distinguish all operation types.

  • Contribution analysis: Taylor decomposition distributes network output credit among nodes across effective DAG layers, including integrated credits through skip connections.Nodes involved in multiple layers accumulate credits from those layers.
  • Resource criteria: For convolutional layers, H, W, f, and k denote spatial and filter dimensions, while I, O, and g denote input channels, output channels, and groups.These quantities support parameter-size, FLOPs, and memory-access calculations.
  • Resource criteria: MAC is simplified as memory access for input/output feature maps and kernel weights under an assumption of sufficient cache.The approximation follows the stated cache assumption.
  • Resource criteria: Parameter size alone suffices for convolution-only constraints because FLOPs are directly proportional to parameter size and MAC is positively correlated with it.This equivalence does not extend to parameter-free operations.
  • Resource criteria: Pooling and skip connections are parameter free, so parameter size cannot distinguish them from convolutional operations.Their resource costs require additional criteria.
  • Resource criteria: FLOPs distinguish pooling from skip connections, MAC distinguishes skip connections from none, and all three criteria distinguish the four operation types.The combined constraint uses parameter size, FLOPs, and MAC as locally decomposable criteria.

G.1 ARCHITECTURE SEARCH ON CIFAR-10

The CIFAR-10 search uses standard image preprocessing, a defined operation set, and separate optimization settings for neural-operation and architecture-distribution parameters.

  • Data preprocessing: Training images are padded to 40×40, randomly cropped to 32 × 32, horizontally flipped, and channel-normalized.The same normalization procedure is applied to training and validation images.
  • Operations: The operation set includes separable and dilated separable convolutions, pooling, skip connection, and zero operation.Operations use stride one except operations adjacent to reduction-cell inputs, which use stride two.
  • Operations: Convolved feature maps are padded to preserve spatial resolution, and convolutions use ReLU-Conv-BN ordering.
  • Optimization: Neural operation parameters use momentum SGD with initial learning rate ηθ = 0.025, while architecture distribution parameters use Adam with initial learning rate ηα = 3 × 10−4.The two parameter groups therefore use different optimizers and learning rates.
  • Optimization: The batch size is 64, with momentum 0.9 and weight decay 3 × 10−4 for neural operation parameters.

G.2 ARCHITECTURE EVALUATION ON CIFAR-10

The evaluation setup adds regularization and auxiliary components commonly used in prior NAS systems.

  • Enhancements: Cutout, path dropout, and auxiliary towers are used as additional enhancement techniques.
  • Enhancements: Path dropout is applied with probability 0.2.The setting matches DARTS in the authors’ publicly released code.
  • Enhancements: Auxiliary towers use weight 0.4.

G.3 ARCHITECTURE TRANSFERABILITY EVALUATION ON CIFAR-10

The transferability evaluation trains a CIFAR-10 network with specified optimization settings and examines SNAS-derived normal and reduction cells.

  • Training settings: The network is trained with batch size 128, weight decay 3 × 10−5, and an initial SGD learning rate of 0.1.The learning rate is decayed by a factor of 0.97 after each epoch.
  • Training settings: Auxiliary towers are adopted with weight 0.4 as an additional enhancement.
  • Cell architectures: Figure 7 presents SNAS cells found on CIFAR-10 under a moderate constraint: one normal cell and one reduction cell.
Loading 1812.09926v3…