Neural Architecture Search

A neural network architecture is a discrete structural object: a choice of layers, operation types, connectivity pattern, and hyperparameters like kernel size or channel width. Traditionally, architects fix this object by hand and only optimize the continuous weights inside it via gradient descent. Neural Architecture Search (NAS) instead treats the architecture as a variable to be optimized, typically to maximize validation accuracy (or a multi-objective combination of accuracy, latency, and parameter count) subject to a training budget. Formally, let $\mathcal{A}$ denote a search space of candidate architectures. For an architecture $a \in \mathcal{A}$, let $w^*(a) = \arg\min_w \mathcal{L}_{\text{train}}(a, w)$ be its optimally trained weights, and $\text{Acc}_{\text{val}}(a, w^*(a))$ its validation performance. NAS solves the bilevel optimization problem

The outer maximization over $a$ is discrete and non-differentiable in general; the inner minimization over $w$ is the ordinary training loop. Solving this exactly requires training every candidate to convergence, which is computationally infeasible for spaces with $10^{10}$–$10^{20}$ candidate architectures. The entire field of NAS is organized around three design axes that make this problem tractable: 1. Search space — what structures are representable. A macro (whole-network) space enumerates entire topologies; a cell-based space searches only a small repeatable building block (a "cell"), stacked $N$ times to form the network, dramatically shrinking the space while retaining expressiveness (used in NASNet, AmoebaNet, DARTS). 2. Search strategy — the algorithm exploring $\mathcal{A}$: reinforcement learning (an RNN controller sampling architectures, rewarded by validation accuracy), evolutionary search (mutate and select a population of architectures), Bayesian optimization, or continuous relaxation for gradient-based search. 3. Performance estimation strategy — how a candidate's quality is judged without full training: proxy tasks (fewer epochs, smaller datasets), learning-curve extrapolation, low-fidelity estimates, and — most importantly for modern NAS — weight sharing, where one large supernetwork is trained once and all candidate sub-architectures reuse its weights. These three axes interact: a small search space can tolerate slow strategies; a slow performance estimator forces coarse search strategies. The history of NAS is largely the story of compressing the performance-estimation cost by orders of magnitude, which is what turned NAS from a $\sim$2,000 GPU-day research curiosity (2016) into an overnight, single-GPU procedure (2018-2019).

Component 1: The search space

The simplest search space is a chain-structured space: a sequence of $L$ layers, each with a type (conv, pool, fc) and hyperparameters (filter count, kernel size, stride). This space is easy to reason about but scales poorly and rarely captures useful motifs like skip connections. The dominant modern choice is a cell-based space. A cell is a small directed acyclic graph (DAG) with $B$ internal nodes; each node's input is a learned combination of previous nodes' outputs, each edge carries an operation chosen from a fixed operation set (e.g., $3{\times}3$ separable conv, $5{\times}5$ dilated conv, max pool, identity, zero). Two cell types are searched: a normal cell (preserves spatial resolution) and a reduction cell (halves resolution, doubles channels). The final network stacks copies of these two cells in a fixed macro pattern. Because the cell is small (typically $B=4$–$7$ nodes), the combinatorial space collapses from astronomically large (whole-network) to a space of size roughly $|\text{ops}|^{2B} \cdot (\text{DAG topologies})$ — still large ($\sim 10^{10}$–$10^{18}$) but far more tractable, and critically, cells discovered on a small proxy task (e.g., CIFAR-10) transfer reasonably well to larger tasks (e.g., ImageNet) because the motif, not the whole network, is what's searched.

Component 2: Search strategies

Reinforcement learning (Zoph & Le, 2016). An RNN controller with parameters $\theta_c$ generates a sequence of tokens describing an architecture (layer types, filter sizes, connections), one token at a time, autoregressively — the architecture description is treated like a sentence. Each sampled architecture $a$ is fully trained, its validation accuracy $R(a)$ is used as a reward, and the controller is updated by the REINFORCE policy-gradient estimator:

where $b$ is a baseline (e.g., an exponential moving average of past rewards) that reduces gradient variance without introducing bias, since $\mathbb{E}[\nabla_{\theta_c}\log \pi_{\theta_c}(a)] = 0$. This is exactly REINFORCE applied with the architecture as the "action" and final validation accuracy as the (delayed, sparse) reward.

Evolutionary search. Maintain a population of architectures. At each generation, sample two candidates, mutate the weaker (or a randomly chosen) one — e.g., change an operation on one edge, add/remove a skip connection — evaluate the child, and replace the oldest (or worst) population member (tournament / aging evolution, as in AmoebaNet). No gradient or controller network is needed; the operator is a discrete mutation plus a selection rule.

Gradient-based / differentiable search (DARTS). Replace the discrete choice of operation on each edge with a continuous relaxation: the output of an edge $(i,j)$ becomes a softmax-weighted mixture of all candidate operations,

where $\alpha^{(i,j)} \in \mathbb{R}^{|\mathcal{O}|}$ are continuous architecture parameters. Now the search problem is a bilevel optimization over continuous variables — architecture weights $\alpha$ and network weights $w$ — solved by alternating gradient steps (approximating the bilevel objective with one-step unrolling):

After training, the final discrete architecture is recovered by keeping, per edge, the operation with the largest $\alpha$ weight ("derive" step). This makes NAS end-to-end differentiable and reduces search cost to roughly the cost of training a single (larger, weight-sharing) network — often under a GPU-day.

Component 3: Performance estimation via weight sharing

The single biggest cost driver in early NAS was that every sampled architecture was trained from scratch (600+ epochs) purely to obtain one reward signal. ENAS (Efficient NAS) and DARTS solve this with weight sharing: build one large "supernetwork" containing every possible operation on every edge, train it once, and let all candidate sub-architectures reuse (share) the same pool of weights, rather than owning independent weights. A sampled sub-architecture's accuracy is estimated by masking out unused operations and running a forward pass with shared weights — no retraining. This is the mechanism that took NAS from ~2,000 GPU-days (original RL-NAS) to ~0.5–4 GPU-days (ENAS, DARTS).

Pseudocode: RL-based NAS controller loop

```

# RL-based Neural Architecture Search (Zoph & Le style) Input: search space encoding (token vocabulary), reward function R(a) controller RNN with parameters theta_c number of search iterations T, architectures per batch K Initialize theta_c randomly Initialize baseline b = 0, decay rate beta = 0.99 for iteration = 1 to T: batch_rewards = [] batch_log_probs = [] for k = 1 to K:

# 1. Sample an architecture description autoregressively tokens = [] hidden = controller.init_hidden() for step = 1 to num_decisions: logits, hidden = controller.forward(tokens, hidden) token = sample_categorical(softmax(logits)) tokens.append(token) architecture a_k = decode(tokens) # e.g. layer types, kernel sizes, skips

# 2. Estimate performance (full training OR weight-shared forward pass) w_k = train_or_reuse_weights(a_k) # bottleneck step reward_k = R(a_k, w_k) # e.g. validation accuracy batch_rewards.append(reward_k) batch_log_probs.append(log_prob(tokens | theta_c))

# 3. Update baseline and compute advantage for reward_k in batch_rewards: b = beta b + (1 - beta) reward_k

# 4. REINFORCE policy-gradient update grad = 0 for k = 1 to K: advantage_k = batch_rewards[k] - b grad += advantage_k grad_of(batch_log_probs[k], theta_c) theta_c = theta_c + learning_rate (grad / K) return best architecture found across all iterations ``` ```

# DARTS: differentiable architecture search (bilevel gradient descent) Input: supernetwork with mixed operations per edge, operation set O train set D_train, validation set D_val learning rates: eta_w (weights), eta_alpha (architecture params) Initialize weights w, architecture params alpha (one vector per edge) repeat until convergence:

# Step 1: update architecture parameters (outer objective)

# approximate w(alpha) by a single training step (first-order approx) w_virtual = w - eta_w grad_w( L_train(w, alpha) ) alpha = alpha - eta_alpha * grad_alpha( L_val(w_virtual, alpha) )

# Step 2: update network weights (inner objective) on real data w = w - eta_w * grad_w( L_train(w, alpha) )

# Derive final discrete architecture for each edge (i, j): keep the operation o = argmax_o alpha_o^(i,j) prune all other operations -> discrete architecture a

# Retrain a from scratch with fresh weights (standard practice) w_final = train(a, D_train) return a*, w_final ``` The DARTS loop's Step 1 is the crux: because computing the true $w^*(\alpha)$ requires a full inner training loop, DARTS approximates it with one virtual gradient step — a first-order (or, in the paper's second-order variant, one-step unrolled) approximation to the bilevel gradient, trading exactness for tractability.

Search space size. For a cell with $B$ intermediate nodes, each node chooses 2 predecessor nodes (from $2+b-1$ available at node $b$) and an operation (from $|\mathcal{O}|$) per incoming edge. The number of distinct cell topologies grows combinatorially:

which for $B=4$, $|\mathcal{O}|=8$ gives roughly $10^{10}$ distinct normal cells (and similarly for reduction cells, multiplying the two). Exhaustive search is infeasible; this is why a strategy is required at all.

Cost of naive (RL, full-training) NAS. If each candidate requires $T_{\text{train}}$ GPU-hours to train fully, and the controller needs $K$ samples per iteration over $I$ iterations, total cost is $O(K \cdot I \cdot T_{\text{train}})$. The original Zoph & Le (2016) NAS used $K \approx 800$ architectures each trained for hours to days, totaling roughly 22,400 GPU-days (≈ 800 GPUs for a month) — the number that motivated every subsequent efficiency improvement.

Cost with weight sharing. ENAS/DARTS instead train one supernetwork containing all operations, at roughly the cost of training a single large network, $O(T_{\text{train}}^{\text{super}})$, independent of $|\mathcal{A}|$ (up to a constant factor for the extra mixture-operation compute per edge, since a mixture evaluates all $|\mathcal{O}|$ operations rather than one). This is the qualitative jump from exponential-in-space-size cost to constant cost in the number of architectures, at the price of an approximation: a shared-weight sub-network's accuracy is only a proxy for its accuracy if retrained independently — the shared weights are jointly optimized to be good for all sub-architectures simultaneously, not perfectly specialized to any one.

Why weight sharing is a defensible approximation (and where it breaks). The core assumption is that a supernetwork's weights, co-adapted across many sub-architectures, still rank sub-architectures in roughly the same order as if each were retrained independently — a hypothesis validated empirically in ENAS/DARTS but not exact: operations that are structurally advantaged under shared training (e.g., skip-connections, which are easy to co-optimize and receive gradient from many paths) tend to be systematically over-selected relative to their true standalone performance. This is the documented DARTS failure mode where search collapses to cells dominated by skip-connections, which perform poorly once retrained in isolation — a real correctness gap between the proxy objective and the true objective, mitigated in follow-up work (e.g., early-stopping the search, regularizing operation strength, or restricting skip-connection counts).

DARTS gradient step complexity. Each search iteration performs two backward passes (one through $\mathcal{L}_{\text{val}}$ w.r.t. $\alpha$ via the virtual weight update, one through $\mathcal{L}_{\text{train}}$ w.r.t. $w$), each costing $O(|\text{supernet params}|)$ — comparable to standard SGD on a network roughly $|\mathcal{O}|\times$ larger than the final derived network (since every edge evaluates all candidate operations). Total search cost is therefore $O(\text{epochs} \times |\mathcal{O}| \times \text{per-epoch training cost})$, empirically ~1–4 GPU-days on CIFAR-10-scale search spaces — several thousand times cheaper than the original RL-NAS.

Correctness caveat — proxy-task transfer. Architectures are almost always searched on a small proxy (CIFAR-10, a few epochs, reduced channel width) and then evaluated on the true target task (ImageNet, full training). The implicit assumption — that relative rankings of architectures are preserved under this rescaling — is empirically reasonable for cell-based motifs but is not a theorem; it can and does fail for search spaces or tasks with large domain gaps, which is why cross-space and cross-task NAS remains an active empirical concern rather than a solved problem.

Trace a miniature DARTS-style search on a toy cell with 2 input nodes, 1 intermediate node, and an operation set $\mathcal{O} = \{\text{conv3x3}, \text{skip}, \text{zero}\}$, over 3 search epochs, to see architecture parameters evolve. Setup: intermediate node $n_1$ receives a mixed edge from each input, $c_{k-2}$ and $c_{k-1}$. Each edge has its own $\alpha$ vector of length 3 (one weight per operation), initialized uniformly ($\alpha = [0,0,0]$, softmax gives $[\tfrac13,\tfrac13,\tfrac13]$ for each op).

At each epoch, Step 1 of the DARTS loop nudges $\alpha$ via the validation-loss gradient (through the one-step virtual weight update), and Step 2 updates the shared weights $w$ on the training set as usual. Notice: on the $(c_{k-2}\!\to\!n_1)$ edge, `conv3x3`'s weight rises from 0.33 to 0.59 — the gradient signal consistently favors it since a learnable convolution reduces loss more than an identity or zero map on this edge. On the $(c_{k-1}\!\to\!n_1)$ edge, `skip` wins (0.51) — plausible if $c_{k-1}$ is already a well-formed representation and adding a convolution there mostly adds noise/overfitting risk early in training.

Derive step: after search, take $\arg\max$ per edge: edge 1 → `conv3x3`, edge 2 → `skip`. The discrete architecture $a^*$ is a cell where $n_1 = \text{conv3x3}(c_{k-2}) + \text{skip}(c_{k-1})$. This $a^*$ is then retrained from scratch with fresh weights (the shared supernet weights are discarded), typically for the full training schedule, to get its true reported accuracy — this retraining step is essential and is often the part omitted when NAS search cost is quoted, inflating apparent efficiency.

Architecture search as an automated discipline crystallized around 2016–2018, building on decades of prior work in neuroevolution and hyperparameter optimization, and quickly connected to core reinforcement-learning and meta-learning ideas.

NAS-derived design principles and directly searched architectures now appear across production systems. MobileNetV3 used a combination of NAS (platform-aware search, optimizing for real mobile-device latency alongside accuracy) and manual refinement (NetAdapt) to hit its efficiency targets, and is deployed broadly in on-device vision pipelines. EfficientNet's compound scaling was informed by NAS-discovered baseline architectures (EfficientNet-B0 was found via a small, multi-objective (accuracy + FLOPs) NAS search), then scaled up by a principled depth/width/resolution rule rather than re-searching at each scale. Google's AutoML platform commercialized NAS-style automated model design for customers without deep-learning expertise, and cloud ML services (e.g., automated vision/tabular model builders) apply related search-and-select pipelines under the hood. NAS techniques also underpin hardware-aware model design, where the reward signal incorporates measured on-device latency or energy rather than accuracy alone — critical for deploying models on phones, wearables, and edge accelerators where the "best" architecture depends on the target chip. This connects tightly to Model Compression and Knowledge Distillation, since many NAS pipelines search for a compact student architecture directly, rather than compressing a fixed large model after the fact — the two techniques are often composed: search for an efficient topology, then further compress or distill the found network. NAS is also a direct generalization of Hyperparameter Optimization, extending the search variable from continuous knobs (learning rate, weight decay) to the discrete structure of the computational graph itself, and it shares its core exploration/exploitation machinery with Reinforcement Learning (the REINFORCE controller) and Evolutionary Algorithms (population-based search operators).

Ask the tutor: How does the choice between first-order and second-order gradient approximation in DARTS affect search stability and cost?