NAS I: Searching for Efficient Networks

Chapter 1: What is NAS?

Neural Architecture Search (NAS) is the process of automatically finding the best neural network architecture from a candidate set, optimizing some objective (accuracy, latency, energy, memory). Every NAS algorithm has three components working together: a search space (the universe of architectures that can be expressed), a search strategy (how we navigate that universe), and a performance estimation strategy (how we quickly evaluate any candidate without full training).

Think of it like this: the search space is the map. The search strategy is the navigation algorithm (random walk vs GPS vs intuition). The performance estimator is a compass that gives you a noisy reading of "how good is this point?" without requiring you to walk the whole path to find out.

The field spent 2016–2022 pushing each of these three components. Early NAS used random search for the strategy and full training for estimation — 450 GPU-days per search run (the original NASNet paper). Modern NAS uses gradient-based search (DARTS, 4 GPU-days) with weight sharing (one supernet, zero retraining per candidate). That's a 100× reduction in search cost.

Chapter 2: Search Space

The search space defines what architectures the NAS algorithm can even consider. A poorly designed search space — one that excludes the right building blocks — means the best possible find is mediocre. There are two main paradigms: cell-based search spaces and network-level (macro) search spaces.

Cell-Based Search Space

Instead of searching the entire network end-to-end, you search for a small repeating module called a cell. The cell is then stacked N times to form the full network. This dramatically shrinks the search space: instead of choosing operations for 50+ layers individually, you choose operations for a 5-block cell and the stacking rules handle the rest.

NASNet (Zoph et al., CVPR 2018) introduced two cell types: a Normal Cell that preserves feature map resolution and a Reduction Cell that halves the spatial dimensions. You stack Normal Cells in groups, with Reduction Cells between groups. The RNN controller only needs to design these two cells — everything else (how to stack them, how many stages) is fixed by the human designer.

Each cell contains B blocks. Each block takes two inputs (hidden states from earlier in the cell), applies an operation to each, and combines them. The search choices per block are: which two inputs to use, which operation to apply to each, and how to combine the results (add or concatenate). With B=5 blocks and |O| candidate operations, the combinatorial count grows as (2×2×M×M×N)^B — which we'll count precisely in the next chapter.

Network-Level Search Space

Macro NAS searches the whole network topology: depth (how many layers per stage), width (how many channels), kernel sizes, resolution, and even connectivity patterns (which layers receive skip connections from which earlier layers). ProxylessNAS and EfficientNet use this approach. Instead of a repeated cell, each layer independently chooses from a menu of operations (MBConv 3×3, MBConv 5×5, MBConv 7×7 with expansion ratios 3 or 6, identity, etc.).

For TinyML specifically, the search space must also include memory constraints. A GPU has 32 GB. An MCU like the STM32F746 has 320 KB of SRAM. Memory is not just a soft constraint — it is a hard wall. MCUNet addresses this by jointly searching the space (which width/depth combinations allow models to fit in SRAM) before searching within the space (which specific architecture has best accuracy).

Chapter 3: Space Size & Cell Counting

Let's count exactly how many architectures live in NASNet's cell-based search space. This is not an academic exercise — it motivates why we need smart search strategies and why exhaustive evaluation is hopeless.

NASNet Cell Structure

Recall the NASNet cell: it contains B = 5 blocks. Each block performs these 5 choices:

1. Which hidden state to feed as the first input (choose from 2 candidates at the start, growing as blocks complete)
2. Which hidden state to feed as the second input
3. Which operation to apply to the first input (choose from M = 5 operations)
4. Which operation to apply to the second input (choose from M = 5 operations)
5. Which combination method to merge the two transformed states (N = 2: add or concatenate)

Assuming 2 candidate inputs for both input choices, M = 5 operations, N = 2 combination methods, per block the number of choices is: 2 × 2 × 5 × 5 × 2 = 200. Over B = 5 blocks: 200⁵? No — that overcounts because blocks are ordered and each new block can also use outputs of prior blocks. The simplified bound from the lecture is:

More precisely, the NASNet paper (with M = 13 operations including sep-conv 3×3, sep-conv 5×5, sep-conv 7×7, dilated sep-conv, max-pool 3×3, avg-pool 3×3, identity, etc.) counts the space as approximately 1.3 × 10¹⁰ unique cells (for the Normal Cell alone). Both the Normal and Reduction Cells must be discovered — so the joint space is even larger.

Worked Example: M=5, N=2, B=5

Let's plug in concrete numbers from the slide: M = 5 candidate operations, N = 2 combination methods, B = 5 blocks.

Training each candidate for just 1 GPU-hour gives 3.2 × 10¹¹ GPU-hours = 36.5 million GPU-years. Even with 800 GPUs running in parallel: 45,600 years. This is why random search (which can sample densely from a good region) and RL/gradient methods (which can move toward better regions) are necessary.

python
# Search space size for cell-based NAS (NASNet-style)
def cell_space_size(M: int, N: int = 2, B: int = 5) -> float:
    """
    M: number of candidate operations per input
    N: combination methods (add=1, concat=2)
    B: number of blocks per cell
    Each block: 2 input choices × 2 input choices × M ops × M ops × N combine
    But input choices grow as blocks complete (for blocks 1..B, choices are 2+b-1)
    Simplified (lecture formula): (2*2*M*M*N)^B
    """
    return (4 * M**2 * N) ** B

# Lecture values
size = cell_space_size(M=5, N=2, B=5)
print(f"M=5, N=2, B=5: {size:.2e} candidates")  # 3.20e+11

size_nasnet = cell_space_size(M=13, N=2, B=5)
print(f"NASNet M=13: {size_nasnet:.2e} candidates")  # ~1.3e+13

# At 1 GPU-hour per candidate with 800 GPUs:
gpu_years = size / (800 * 8760)
print(f"GPU-years with 800 GPUs: {gpu_years:.0f}")  # 45,662 years

Chapter 4: RL Controller

The first successful NAS algorithm (Zoph and Le, ICLR 2017) used reinforcement learning to search the space. The idea is beautiful: treat architecture design as a sequential decision-making problem. An RNN controller samples an architecture token by token — "first layer: 3×3 conv, second layer: 5×5 depthwise, third layer: identity..." — like a language model generating a sentence. The architecture is the "sentence." The validation accuracy of the trained network is the "reward."

The RL Controller Loop

The controller is a recurrent neural network. At each step, it emits a categorical distribution over the possible choices at that position (operation type, kernel size, skip connection target, etc.) via a softmax. It samples from this distribution to generate one discrete architectural token. After all tokens are generated, the sampled architecture is trained for a few epochs and evaluated on a validation set.

The REINFORCE Update — Derived

The controller has parameters θ. It defines a policy π_θ(a) — the probability of sampling architecture a. We want to maximize the expected reward J(θ) = E_a~πθ[R(a)].

We can't compute this expectation exactly (the space has 10¹¹+ architectures). But the REINFORCE gradient trick lets us estimate the gradient unbiasedly from samples:

In words: if architecture a got reward R(a), push the controller's parameters in the direction that would have made a more likely — scaled by R(a). Architectures with high reward get positively reinforced. Architectures with low reward get negatively reinforced.

Worked Step: One REINFORCE Update

Suppose the controller samples architecture a = [3×3 conv, 5×5 dwconv, skip]. The RNN's log-probability of generating this sequence is log π_θ(a) = log(0.3) + log(0.4) + log(0.5) = −4.12. The architecture achieves R = 74.2% validation accuracy on a held-out set. The reward baseline (exponential moving average of past rewards) is b = 72.0%.

We subtract the baseline to reduce variance:

The positive (R − b) = 2.2 means: this architecture was better than average. The gradient update pushes θ to make this architecture's tokens more likely in the future. If the reward were below baseline, the signal would be negative — suppressing those choices.

python
# Sketch: RL controller sampling + REINFORCE update
import torch
import torch.nn as nn

class NASController(nn.Module):
    def __init__(self, num_ops=5, num_layers=10, hidden=64):
        super().__init__()
        self.rnn = nn.LSTMCell(num_ops, hidden)
        self.heads = nn.ModuleList([
            nn.Linear(hidden, num_ops) for _ in range(num_layers)
        ])
        self.baseline = 0.0

    def sample_architecture(self):
        # returns (architecture_tokens, log_prob_sum)
        tokens, log_probs = [], []
        h = c = torch.zeros(1, self.rnn.hidden_size)
        x = torch.zeros(1, self.rnn.input_size)
        for head in self.heads:
            h, c = self.rnn(x, (h, c))
            logits = head(h)                      # (1, num_ops)
            dist = torch.distributions.Categorical(logits=logits)
            tok = dist.sample()                  # discrete choice
            tokens.append(tok.item())
            log_probs.append(dist.log_prob(tok))
            x = nn.functional.one_hot(tok, self.rnn.input_size).float()
        return tokens, torch.stack(log_probs).sum()

    def reinforce_update(self, log_prob_sum, reward, lr=1e-3, ema=0.9):
        self.baseline = ema * self.baseline + (1 - ema) * reward
        loss = -log_prob_sum * (reward - self.baseline)
        # maximize reward ⟺ minimize -reward * log π(a)
        loss.backward()

Chapter 5: Evolutionary Search

The second major NAS search strategy mimics biological evolution. You maintain a population of candidate architectures. You evaluate each on a fitness function (usually validation accuracy, sometimes combined with latency). You keep the fittest individuals, mutate some of them (randomly change one operation or one layer depth), and crossover others (recombine choices from two parents). The population evolves over generations toward better architectures.

Why Evolution Works for Architecture Search

Evolution is effective when the fitness landscape has structure — when similar architectures tend to have similar performance. Architecture space does have this property: an architecture that uses 5×5 depthwise conv in layer 3 is likely to perform similarly to one that uses 3×3 depthwise conv in layer 3, because they share the same blueprint everywhere else. Small mutations produce small performance changes, not catastrophic ones. This smooth landscape means evolved populations can navigate toward good regions.

Mutation

Mutation on depth: Change Stage 1 from 3 layers to 2 layers, keeping all operations intact. This tests whether the network can achieve the same accuracy with fewer parameters at that stage.

Mutation on operator: Keep all layer counts fixed but change one operation from MB6 3×3 to MB6 5×5 in Stage 2 Layer 1. This tests whether a larger receptive field at that specific position improves accuracy.

In Once-for-All (Cai et al., ICLR 2020), a supernet is trained once and mutations are applied to sub-network architectures sampled from the supernet. Mutation cost = zero retraining, just re-evaluate the mutated sub-network using the shared weights.

Crossover

Crossover takes two parent architectures and randomly combines their layer choices. Parent 1 uses MB6 3×3 in layers {1,3,5}; Parent 2 uses MB6 5×5 in layers {2,4,6}. For each layer, randomly pick from Parent 1 or Parent 2 to form the child. The child inherits the best genes (hopefully) from both parents.

Comparison: RL vs Evolution

python
# Evolutionary NAS sketch — mutation + tournament selection
import random

def mutate(arch, ops_per_layer, p_mutate=0.1):
    """Randomly change each layer's op with probability p_mutate."""
    return [
        random.choice(ops_per_layer) if random.random() < p_mutate else op
        for op in arch
    ]

def crossover(parent1, parent2):
    """For each layer, randomly pick from parent1 or parent2."""
    return [
        p1 if random.random() < 0.5 else p2
        for p1, p2 in zip(parent1, parent2)
    ]

def evolutionary_nas(search_space, fitness_fn, pop_size=50, generations=30):
    ops = search_space['ops']
    # Initialize random population
    population = [[random.choice(ops) for _ in range(search_space['layers'])]
                  for _ in range(pop_size)]
    for gen in range(generations):
        scores = [fitness_fn(a) for a in population]
        # Keep top 50%, generate children via mutation + crossover
        top_k = sorted(zip(scores, population), reverse=True)[:pop_size//2]
        survivors = [arch for _, arch in top_k]
        children = [mutate(crossover(*random.sample(survivors, 2)), ops)
                    for _ in range(pop_size - len(survivors))]
        population = survivors + children
    return max(population, key=fitness_fn)

Chapter 6: DARTS: Gradient-Based NAS

RL and evolutionary search are both black-box methods: they treat the search space as discrete, treat performance evaluation as an oracle, and can't compute gradients through architecture choices. DARTS (Differentiable Architecture Search, Liu et al., ICLR 2019) fixes this by making the architecture search itself differentiable.

The Continuous Relaxation

In a standard cell, each edge (i→j) carries exactly one operation (e.g., 3×3 conv). DARTS replaces that single choice with a mixture of all operations, weighted by learned scalars. The key equation:

Here O is the set of candidate operations. α^(i,j) is a vector of real-valued scalars (one per operation) — the architecture parameters. The softmax converts them to weights that sum to 1. The "mixed operation" ō simultaneously runs ALL operations on x and produces a weighted average of their outputs.

Worked Example: 3 Operations

Suppose O = {3×3 conv, 5×5 conv, skip}. The architecture parameters for edge (0,1) are α = [2.1, 0.3, −0.8]. Let's compute the DARTS mixed operation output for input x:

The gradient of the validation loss with respect to α flows through this weighted sum. If conv3×3 outputs lead to lower loss (better validation accuracy), its α value increases, increasing its weight in future mixed operations. The α values converge toward the best operation on each edge.

Discretization

After the search phase, each edge is discretized: pick the operation with the highest α weight. In the example above, that's conv3×3 (weight 0.820). The supernet collapses to a single discrete architecture. This is then retrained from scratch to get the final accuracy.

Bilevel Optimization

DARTS has two sets of parameters: network weights w (the convolution filters) and architecture parameters α. They're optimized jointly via bilevel optimization:

In practice this is approximated by alternating: one gradient step on w using training data, one gradient step on α using validation data. This avoids the expensive inner optimization loop.

python
# DARTS mixed operation — the core idea
import torch
import torch.nn as nn
import torch.nn.functional as F

class MixedOp(nn.Module):
    """DARTS mixed operation: weighted sum over all candidate ops."""
    def __init__(self, ops: list[nn.Module]):
        super().__init__()
        self.ops = nn.ModuleList(ops)
        # Architecture parameters α — initialized near uniform
        self.alphas = nn.Parameter(torch.zeros(len(ops)))

    def forward(self, x):
        weights = F.softmax(self.alphas, dim=0)  # shape: (|O|,)
        return sum(w * op(x) for w, op in zip(weights, self.ops))

    def discretize(self) -> nn.Module:
        """Return the operation with the highest α weight."""
        best_idx = self.alphas.argmax().item()
        return self.ops[best_idx]

# During DARTS search (alternating bilevel optimization):
# Step 1: one step on w using train_batch, α fixed
# Step 2: one step on α using val_batch, w fixed
for train_batch, val_batch in zip(train_loader, val_loader):
    w_optim.zero_grad()
    loss_train = model(*train_batch)
    loss_train.backward()
    w_optim.step()  # update weights w, α unchanged

    alpha_optim.zero_grad()
    loss_val = model(*val_batch)
    loss_val.backward()
    alpha_optim.step()  # update α, w unchanged

Chapter 7: Performance Estimation

The bottleneck in any NAS algorithm isn't the search strategy — it's how you evaluate each candidate. Training a network to convergence on ImageNet takes 90 epochs × 1.2M images × 4 GPU-hours = days per candidate. With 10¹¹ candidates in the search space, you need much faster estimation.

Strategy 1: Lower-Fidelity Evaluation

Train each candidate for fewer epochs on a smaller dataset (e.g., CIFAR-10 as a proxy for ImageNet), or at lower resolution, or with fewer channels. The relative ranking of architectures is usually preserved even under these degraded conditions — an architecture that wins at 5 epochs often wins at 90 epochs. This reduces cost by 10–20×.

Strategy 2: Weight Sharing (One-Shot / Supernet)

The most impactful innovation in NAS. Instead of training each candidate from scratch, train one giant supernet whose weights are shared across all sub-network architectures. A sub-network (a specific architecture) is just a choice of which edges/operations to activate in the supernet. All sub-networks sharing edges also share the weights on those edges.

SMASH (Brock et al.) and ENAS (Pham et al., 2018) were early weight-sharing methods. Once-for-All (OFA, Cai et al., ICLR 2020) extends this: train one supernet once for 300 epochs, then search over sub-networks for free by just evaluating them using the shared weights — no retraining at all.

Strategy 3: Accuracy Predictors

Train a regression model (usually a small MLP or GNN) to predict a new architecture's accuracy from its encoding. The predictor is trained on ~500 architecture-accuracy pairs from early in the NAS run. Once trained, it scores millions of new candidates in milliseconds. This is the approach used by NAT, BRP-NAS, and Semi-Supervised NAS.

The Cost Comparison

Chapter 8: Showcase: NAS Strategy Explorer

This showcase puts all three NAS search strategies side-by-side: Random Search, Evolutionary, and DARTS (simulated as gradient ascent in accuracy space). A 2D accuracy-vs-MACs landscape represents the search space. The goal: find the Pareto-optimal front of high-accuracy, low-MACs architectures. Watch how each strategy navigates — and how quickly each converges.

Chapter 9: Connections & Cheat Sheet

NAS is not one algorithm — it's a framework with three independently-improving components. The best modern NAS systems (OFA, EfficientNet, ProxylessNAS) mix and match these components to achieve 100× faster search than the 2017 baseline.

NAS Component Cheat Sheet

RL vs Evolution vs DARTS — When to Use Which

Bridge to NAS II: Hardware-Aware and Once-for-All

This lecture covered NAS fundamentals. The next frontier is hardware-aware NAS: making the search space and objective directly aware of target-hardware latency, not just FLOPs. ProxylessNAS models latency as a differentiable function of architecture parameters. MnasNet adds a latency term to the RL reward. OFA trains a single supernet and searches for specialized sub-networks for each hardware target — a phone, a Raspberry Pi, an MCU — without retraining.

NAS II (MIT 6.5940 Lecture 8) covers: hardware-aware NAS, Once-for-All progressive shrinking, zero-shot NAS (scoring architectures without ANY training using gradient statistics), and neural-hardware co-design (searching architecture and accelerator simultaneously).

Related Gleams

• ← TinyML L2: NN Building Blocks & Compute — the ops NAS picks from (MBConv, depthwise-sep, bottleneck)
• ← TinyML L4: Pruning II — AMC uses RL (DDPG) to search sparsity ratios per layer, closely related to NAS RL
• ← TinyML L5: Quantization I — NAS and quantization are often co-searched (mixed-precision NAS)
• Policy Gradients — REINFORCE and the policy gradient theorem, the theoretical foundation for RL-based NAS