Neural Architecture Search

Neural Architecture Search (NAS) automates the design of neural network architectures. Instead of hand-crafting architectures, NAS optimizes the architecture and weights jointly or sequentially. This module covers search spaces, search strategies, and efficient NAS methods with mathematical foundations.

1. Search Space Design

1.1 Cell-Based Search Space

Most modern NAS methods search for a cell that is repeated to form the full architecture:

Normal cell: Preserves spatial resolution Reduction cell: Halves spatial resolution, doubles channels

A cell is a directed acyclic graph (DAG) with nodes. Each intermediate node is computed from all predecessor nodes:

where is the mixed operation:

1.2 Operation Search Space

The set of candidate operations typically includes:

1.3 Hierarchical Search Space

The full architecture is parameterized as:

where is the number of channels, is the resolution, and is the number of layers in stage .

2. Search Strategies

2.1 Reinforcement Learning (NASNet)

Zoph & Le (2016) model architecture search as a sequential decision process:

Controller RNN outputs architecture parameters:

Reward: Validation accuracy of the child architecture.

Policy gradient (REINFORCE):

where is a baseline (exponential moving average of past rewards).

2.2 Evolutionary Search (AmoebaNet)

Real et al. (2019) use tournament selection:

1. Initialize population with random architectures
2. Evaluate each architecture
3. Select parents via tournament selection
4. Mutate parent to produce offspring
5. Replace weakest with offspring
6. Repeat

Mutation operations: Add/remove nodes, change operations, modify skip connections.

Age-based selection: Prioritize younger architectures to prevent premature convergence:

2.3 Bayesian Optimization

Model the architecture-performance mapping as a Gaussian process:

Acquisition function (Expected Improvement):

3. Differentiable Architecture Search (DARTS)

3.1 Continuous Relaxation

DARTS (Liu et al., 2019) relaxes the discrete architecture space to continuous:

3.2 bilevel Optimization

Inner loop (architecture weights fixed, optimize network weights ):

Outer loop (fix , optimize architecture ):

The gradient:

Computing requires differentiating through the optimization, which is expensive. DARTS approximates using:

where is the learning rate for .

3.3 DARTS Discretization

After continuous search, the architecture is discretized:

Keep top- predecessors per node (typically ).

3.4 DARTS Variants

PC-DARTS (Pang et al., 2020): Partial channel connections to reduce memory.

Fair DARTS (Sai et al., 2020): Avoids operation overlap by using independent Gumbel distributions.

DARTS- (Zela et al., 2020): Progressive narrowing of the search space.

GDAS (Dong & Yang, 2019): Gumbel-Softmax sampling for discrete architecture selection.

4. One-Shot NAS

4.1 Weight Sharing

Train a supernet containing all candidate architectures:

Each sub-architecture shares weights with the supernet.

4.2 Path Sampling

Sample paths through the supernet:

where is a path (sequence of operations) and are the shared edge weights.

4.3 Hardware-Aware NAS

Incorporate hardware constraints:

Or use Lagrangian relaxation:

4.4 Neural Architecture Search with Weight Sharing (NASWS)

ENAS (Pham et al., 2018): Share weights across all subgraphs of a single directed graph.

SPOS (Guo et al., 2020): Single Path One-Shot—sample one path per iteration.

5. Efficient NAS Methods

5.1 Proxy Tasks

Search on a smaller proxy task (fewer epochs, smaller dataset):

Zero-cost proxies: Estimate architecture quality without full training:

where is random initialization.

5.2 Progressive NAS

Start with a small search space and progressively expand:

5.3 Differentiable NAS with Regularization

6. Implementation

import torch
import torch.nn as nn
import torch.nn.functional as F

class MixedOp(nn.Module):
    def __init__(self, C, stride, ops_dict):
        super().__init__()
        self.ops = nn.ModuleList([
            ops_dict[name](C, C, stride) for name in ops_dict
        ])
        self.num_ops = len(ops_dict)

    def forward(self, x, weights):
        return sum(w * op(x) for w, op in zip(weights, self.ops))

class DARTSCell(nn.Module):
    def __init__(self, steps, C, ops_dict):
        super().__init__()
        self.steps = steps
        self.ops = nn.ModuleList()
        self._indices = []
        
        for i in range(steps):
            for j in range(2 + i):
                stride = 1
                self.ops.append(MixedOp(C, stride, ops_dict))
                self._indices.append((i, j))

    def forward(self, s0, s1, alphas):
        states = [s0, s1]
        offset = 0
        
        for i in range(self.steps):
            s = 0
            for j, h in enumerate(states):
                weights = F.softmax(alphas[offset], dim=-1)
                s += self.ops[offset](h, weights)
                offset += 1
            states.append(s)
        
        return torch.cat(states[-self.steps:], dim=1)

class DARTS(nn.Module):
    def __init__(self, C=16, steps=4, num_classes=10, layers=8):
        super().__init__()
        self.num_nodes = steps
        num_edges = sum(range(2, steps + 2))
        
        self.alphas_normal = nn.Parameter(1e-3 * torch.randn(num_edges, len(OPS)))
        self.alphas_reduce = nn.Parameter(1e-3 * torch.randn(num_edges, len(OPS)))
        
        self.cells = nn.ModuleList()
        for i in range(layers):
            reduction = (i in [layers // 3, 2 * layers // 3])
            C_curr = C if not reduction else C * 2
            cell = DARTSCell(steps, C_curr, OPS)
            self.cells.append(cell)

    def forward(self, x):
        s0 = s1 = self.stem(x)
        
        for i, cell in enumerate(self.cells):
            alphas = self.alphas_reduce if cell.reduction else self.alphas_normal
            s0, s1 = s1, cell(s0, s1, alphas)
        
        return self.classifier(s1)

def compute_architecture_loss(model, input_train, target_train, 
                               input_val, target_val, w_optim, 
                               arch_lr=0.0003):
    # Update architecture params
    loss_val = F.cross_entropy(model(input_val), target_val)
    grads = torch.autograd.grad(loss_val, model.arch_parameters())
    
    alpha = model.alphas_normal.clone().detach()
    for p, g in zip(model.arch_parameters(), grads):
        alpha -= arch_lr * g
    
    # Update architecture
    with torch.no_grad():
        model.alphas_normal.copy_(alpha)
    
    # Compute training loss for weight update
    loss_train = F.cross_entropy(model(input_train), target_train)
    return loss_train

7. SVG: NAS Search Space Illustration

8. SVG: DARTS Architecture Optimization

9. Search Space Analysis

9.1 Over-parameterization

The search space is heavily over-parameterized: with nodes and operations, there are possible cells.

Reduction techniques:

• Restrict to predecessors per node
• Limit operations to
• Use hierarchical search spaces with progressively larger cells

9.2 Transferability

Architectures found on CIFAR-10 transfer well to ImageNet (same family of architectures). However:

• Channel scaling: Scale channels by when transferring
• Depth scaling: Adjust number of cells
• Resolution: Match target resolution

10. Comparison of NAS Methods

11. Open Problems

• Search space design: How to define search spaces that enable discovering truly novel architectures?
• Evaluation protocol: Proxy tasks may not accurately predict full training performance
• Multi-objective: Balancing accuracy, latency, memory, and energy
• Hardware-aware: Co-optimizing architecture and hardware
• NAS for new domains: 3D, graphs, sparse attention, mixture-of-experts
• Interpretability: Understanding why certain architectures work better

References

1. Zoph, B., & Le, Q. V. (2016). Neural Architecture Search with Reinforcement Learning. ICLR.
2. Liu, H., Simonyan, K., & Yang, Y. (2019). DARTS: Differentiable Architecture Search. ICLR.
3. Real, E., Aggarwal, A., Huang, Y., & Le, Q. V. (2019). Regularized Evolution for Image Classifier Architecture Search. AAAI.
4. Pham, H., Guan, M. Y., Zoph, B., Le, Q. V., & Dean, J. (2018). Efficient Neural Architecture Search via Parameter Sharing. ICML.
5. Elsken, T., Metzen, J. H., & Hutter, F. (2019). Neural Architecture Search: A Survey. JMLR.
6. White, C., et al. (2023). A Neural Architecture Search (NAS) Algorithm. arXiv.