Neural Architecture Search: Theory & Methods

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. The NAS Problem
5. Search Spaces
6. RL-Based Search (NASNet)
7. Evolutionary Search (AmoebaNet)
8. DARTS: Differentiable NAS
9. One-Shot NAS & Weight Sharing
10. Hardware-Aware NAS
11. Common Pitfalls
12. Summary
13. Exercises

Learning Objectives

1. Formulate NAS as a bilevel optimization problem.
2. Compare RL, evolutionary, and gradient-based search strategies.
3. Derive DARTS continuous relaxation of the discrete search space.
4. Explain weight sharing and why it reduces search cost.
5. Design hardware-aware search with latency constraints.

Notation

• $\mathcal{A}$ — search space (set of candidate architectures)
• $\alpha$ — architecture parameters
• $w$ — network weights
• $\mathcal{L}_{\text{val}}, \mathcal{L}_{\text{train}}$ — validation and training loss

Core Intuition

Designing neural network architectures requires expertise and extensive experimentation. NAS automates this: define a SEARCH SPACE of possible architectures, then use SEARCH STRATEGIES to find the best one. Early methods (NASNet) used RL — training thousands of networks to evaluate each candidate. Modern methods (DARTS, one-shot) make the search differentiable or use weight sharing — finding architectures in GPU-hours instead of GPU-months.

The NAS Problem

Bilevel optimization:

Three components:

1. Search space $\mathcal{A}$: What architectures are possible?
2. Search strategy: How to explore $\mathcal{A}$ efficiently?
3. Performance estimation: How to evaluate a candidate cheaply?

Challenge: $|\mathcal{A}|$ is combinatorially large (e.g., $10^{18}$ possible architectures). Can't evaluate all.

Search Spaces

Cell-based (NASNet): Search for two cell types (normal + reduction), then stack them:

• Operations: 3x3 conv, 5x5 conv, pooling, identity, dilated conv.
• Connections: which operations feed into which.
• Cell repeated $N$ times → full architecture.

Network-level: Search for depth, width, resolution per stage.

• EfficientNet: compound scaling of depth × width × resolution.

Operation-level: For each edge in a DAG, choose one operation from a set.

Macro search: Search overall topology (skip connections, branching). Micro search: Fix topology, search operation choices.

RL-Based Search (NASNet)

Zoph & Le (2017): Controller RNN generates architecture descriptions; trained with REINFORCE:

where $R(a)$ = validation accuracy of architecture $a$.

Process:

1. Controller samples an architecture (sequence of choices).
2. Train the architecture from scratch (full training).
3. Evaluate on validation set → reward.
4. Update controller with REINFORCE.
5. Repeat for 20,000+ architectures.

Cost: 2000 GPU-days for original NASNet. Prohibitively expensive.

Result: Found architectures outperforming human-designed ones (NASNet-A on ImageNet).

Evolutionary Search (AmoebaNet)

Real et al. (2019): Tournament selection + mutation:

1. Maintain population of architectures.
2. Select two random candidates; keep the better one (tournament).
3. Mutate winner (change one operation, add/remove connection).
4. Train mutant; add to population.
5. Remove oldest member.

Advantages over RL:

• Simpler (no controller training).
• More diverse exploration (multiple candidates simultaneously).
• Comparable results to RL-based search.

Cost: Similar to RL (thousands of full training runs).

DARTS: Differentiable NAS

Liu et al. (2019): Make architecture search differentiable:

Continuous relaxation: Instead of choosing ONE operation per edge, use a weighted sum:

where $\alpha_o$ are learnable architecture parameters.

Bilevel optimization with gradient descent:

• Update $w$ on training loss: $w \leftarrow w - \eta_w \nabla_w \mathcal{L}_{\text{train}}(w, \alpha)$.
• Update $\alpha$ on validation loss: $\alpha \leftarrow \alpha - \eta_\alpha \nabla_\alpha \mathcal{L}_{\text{val}}(w, \alpha)$.

After search: Discretize by keeping the top-1 operation per edge ($\arg\max_o \alpha_o$).

Cost: 1-4 GPU-days (1000x cheaper than RL-based NAS).

Limitations: Skip connection collapse (DARTS often converges to mostly skip connections). Fix: early stopping of architecture search; regularization of $\alpha$.

One-Shot NAS & Weight Sharing

Key insight: Train a SUPERNET containing all possible architectures as subnetworks. Evaluate any architecture by extracting its weights from the supernet.

Process:

1. Build a supernet (all operations on all edges active simultaneously).
2. Train supernet with uniform path sampling (each step, randomly activate one subnetwork).
3. After training: evaluate candidate architectures by using their paths' weights from the supernet.
4. Select the architecture with best supernet-inherited performance.

Cost: One training run of the supernet (same cost as training ONE architecture).

Accuracy: Weight sharing introduces noise (shared weights aren't optimal for any single architecture). But ranking correlation with independent training is high enough for search.

Hardware-Aware NAS

Optimize for latency, not just accuracy:

or equivalently: $\min_\alpha \mathcal{L}_{\text{val}}(\alpha)$ subject to $\text{Latency}(\alpha) \leq T$.

Latency prediction:

• Lookup table: measure latency of each operation on target hardware.
• Latency model: learned predictor from architecture → latency.
• On-device measurement: directly measure (slow but accurate).

MNASNet (Tan et al., 2019): Found architectures Pareto-optimal for accuracy vs mobile latency.

EfficientNet: NAS-found base architecture (B0), then compound-scaled (B1-B7) by jointly increasing depth/width/resolution.

Common Pitfalls

Pitfall 1. DARTS skip-connection collapse. Without regularization, DARTS converges to architectures dominated by skip connections (identity operations that reduce training loss but limit capacity). Use auxiliary loss on operation diversity.

Pitfall 2. Supernet weight coupling. In one-shot NAS, weights are shared but operations compete. A "bad" operation sharing weights with a "good" one may appear better than it is (free-riding).

Pitfall 3. Searching on proxy tasks that don't transfer. Architectures found on CIFAR-10 (small images, 10 classes) may not be optimal for ImageNet (large images, 1000 classes). Use search tasks similar to the deployment target.

Summary

• NAS: Automate architecture design as bilevel optimization.
• RL/Evolutionary: Train thousands of architectures; expensive but thorough.
• DARTS: Differentiable search via continuous relaxation; 1000x cheaper.
• One-shot: Weight-sharing supernet; cost of ONE training run.
• Hardware-aware: Optimize accuracy + latency jointly.
• Results: NAS-found architectures (EfficientNet, MobileNet) dominate human designs.

Exercises

Exercise 1. Define the search space for a 6-layer transformer NAS: specify what choices exist at each layer (attention heads, FFN size, activation).

Exercise 2. Derive the gradient $\nabla_\alpha \mathcal{L}_{\text{val}}$ in DARTS using the chain rule through the softmax relaxation.

Exercise 3. For one-shot NAS with 8 operations per edge and 14 edges: compute the total number of architectures in the search space.

Exercise 4. Build a latency lookup table for: 3x3 conv, 5x5 conv, 3x3 depthwise conv, skip connection on a target hardware. Use these to predict full network latency.

Exercise 5. Compare DARTS vs random search on a small search space (100 architectures). When does DARTS outperform random?