Search spaces, random search, evolution and reinforcement-learning controllers
On this page
Designing architectures by hand is slow and mostly trial and error: a layer made wider here, a skip connection (a path that carries a layer's input past other layers) added there.
Neural architecture search (NAS) hands this job to an automated search.
NAS runs once, before the final model is trained; the search is its own job (Section 3 has its cost), and its output architecture is then trained from scratch. During the search the strategy keeps a memory: a population for evolution, a controller network for reinforcement learning, or architecture weights for differentiable search (One-Shot NAS: Weight Sharing and DARTS).
This page builds on convolutional networks: kernels and feature maps from CNN Basics: Convolution, Stride and Pooling, and the skip connections of ResNet from Classic CNN Architectures: AlexNet to ResNet. Every NAS method has the same three pieces (Figure 1).
Get these wrong and thousands of GPU hours rediscover ResNet, a hand-designed network.
Say each of L layers picks one op out of k candidates. Take an illustrative space of eight layers and five ops, the one Listing A searches: a 3x3 convolution, a 5x5 convolution, a 3x3 separable convolution, max pooling and a skip that passes its input through. One layer gives 5 options, two layers give 25 pairs, and every further layer multiplies by 5 again, so eight layers give 58 = 390,625 architectures. The card states that count N for any k and L.
At 390,625 architectures for only eight layers, brute force is not a plan. So what matters is the search strategy, not the size of the space.
Random search and evolution both fit in a few lines on a toy space of 200 architectures with fixed accuracies; score is a lookup standing in for training, and mutate moves up to 15 places.
import random
N, BUDGET = 200, 100
acc = [random.random() for _ in range(N)]
# stands in for training
def score(a):
return acc[a]
def mutate(a):
step = random.randint(-15, 15)
# stay inside the space
return min(N - 1, max(0, a + step))
samples = random.choices(range(N), k=BUDGET)
best_random = max(samples, key=score)score looks up a stored accuracy instead of training.mutate moves to a nearby architecture, at most 15 places away (this page's choice, not a published setting). A larger step explores more but can jump off a good peak; a smaller one refines a peak but can stay stuck on a small one.BUDGET = 100 random picks and keep the best. A bigger budget gets every strategy closer to the best architecture at proportionally more GPU time.Evolution spends the same BUDGET of architectures to train and score on a population of 8, replacing the worst with a mutated copy of the best. Hill climbing is a population of one that keeps a mutation only if it scores higher.
pop = random.choices(range(N), k=8)
for _ in range(BUDGET - 8):
# worst first
pop.sort(key=score)
pop[0] = mutate(pop[-1])
best_evo = max(pop, key=score)NAS can only find what the search space contains.
Figure 3. Three searches race; hill climbing can stall.
Toy simulation with a fixed seed, not real NAS results.
Reinforcement learning is the NASNet approach. A controller, usually a small RNN (recurrent neural network, which emits its choices one at a time; taught in Recurrent Neural Networks: RNN, GRU, and LSTM), samples architectures; each is trained, and its validation accuracy (on held-out data) comes back as a reward.
The controller is a policy: for each choice it gives every option a probability, written π(a) for the option a it actually picked, and the reward R is the child's validation accuracy. The first NAS updates it with REINFORCE against a baseline (step 4 of Figure 5); NASNet uses PPO (Proximal Policy Optimization, which limits how far one update can move the policy). REINFORCE, the log-probability it raises and the baseline it subtracts are taught, with PPO, in Policy-Based and Actor-Critic Methods: REINFORCE, PPO and SAC.
Figure 5. One loop round, illustrative numbers; reward lifts each pick.
loss = -logp * reward.The listing runs that update without a baseline on a one-layer toy: logits choose among kernels 3, 5 and 7, and a lookup table perf stands in for training.
import torch
from torch.distributions import Categorical
# kernels 3, 5, 7
perf = torch.tensor([0.70, 0.90, 0.60])
logits = torch.nn.Parameter(torch.zeros(3))
opt = torch.optim.Adam([logits], lr=0.1)
for step in range(50):
dist = Categorical(logits=logits)
action = dist.sample()
reward = perf[action]
logp = dist.log_prob(action)
loss = -logp * reward
opt.zero_grad()
loss.backward()
opt.step()The probability mass drifts toward kernel 5; without a baseline some runs wander first. With real training in place of the lookup, this is the first NAS controller's REINFORCE update in miniature, minus its baseline and its RNN. NASNet trained the same kind of RNN controller with PPO instead.
NASNet (Zoph et al. 2018) cut the cost of the first NAS by searching for small cells instead of whole networks (Figure 6).
NASNet searches for two cells on CIFAR-10 (60,000 tiny labeled images), a cheap stand-in for ImageNet (the large standard image-classification dataset): a normal cell keeps the feature map (a layer's output grid) the same size and a reduction cell halves it. The same cell design scales by construction.
EfficientNet-B0, the best-known result of RL search, reaches 77.1% top-1 on ImageNet (first guess right) with 5.3M weights. The controller searched the MnasNet mobile space with a reward trading accuracy against FLOPs (the count of arithmetic operations). The authors then scaled B0 up with compound scaling, one coefficient growing depth, width and resolution (the card below).
EfficientNet's own ratios (Tan and Le 2019) are 1.2 for depth, 1.1 for width and 1.15 for resolution, applied once per step of the knob φ (B0 is φ = 0); the card names them α, β and γ.
| φ | depth | width | resolution | FLOPs |
|---|---|---|---|---|
| 1 | ×1.2 | ×1.1 | ×1.15 | ×1.92 |
| 2 | ×1.44 | ×1.21 | ×1.32 | ×3.7 |
The strategy is an outer loop; each child is an ordinary short training job on a proxy, and its validation score goes back to the strategy. The evolution loop of Section 2 cannot be reused: it rescores the whole population every round (744 calls to score for a budget of 100), and its mutate moves an index along a line, not a choice inside an architecture.
Listing A trains one child: an architecture is a list of L choices from the five ops of Section 2.1, and score trains it briefly on proxy_loader (a training subset) and returns accuracy on the held-out val_loader.
import torch
import torch.nn as nn
import torch.nn.functional as F
C = 16
MAKE = {
"conv3": lambda: nn.Sequential(nn.Conv2d(C, C, 3, padding=1), nn.ReLU()),
"conv5": lambda: nn.Sequential(nn.Conv2d(C, C, 5, padding=2), nn.ReLU()),
"sep3": lambda: nn.Sequential(nn.Conv2d(C, C, 3, padding=1, groups=C),
nn.Conv2d(C, C, 1), nn.ReLU()),
"pool": lambda: nn.MaxPool2d(3, stride=1, padding=1),
"skip": lambda: nn.Identity(),
}
def build(arch):
layers = [MAKE[op]() for op in arch]
return nn.Sequential(nn.Conv2d(3, C, 3, padding=1), *layers,
nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(C, 10))
def score(arch, epochs=5):
net = build(arch)
opt = torch.optim.SGD(net.parameters(), lr=0.05, momentum=0.9)
net.train()
for _ in range(epochs):
for x, y in proxy_loader:
loss = F.cross_entropy(net(x), y)
opt.zero_grad()
loss.backward()
opt.step()
net.eval()
with torch.no_grad():
hits = sum((net(x).argmax(1) == y).sum().item() for x, y in val_loader)
return hits / len(val_loader.dataset)build, which stacks one op per layer between a first convolution and a classifier. Convolution and max pooling are taught in CNN Basics: Convolution, Stride and Pooling; sep3 is a depthwise-separable convolution (groups=C), taught in Efficient CNNs: MobileNet and EfficientNet.Listing B is aging evolution (Algorithm 1 of Real et al. 2019, the AmoebaNet search). Members keep their scores, so each child trains once; each cycle mutates the best of S sampled members and drops the oldest.
import random, collections
OPS = ["conv3", "conv5", "sep3", "pool", "skip"]
L, P, S, BUDGET = 8, 10, 3, 100
def mutate(arch):
child = list(arch)
i = random.randrange(L)
others = [op for op in OPS if op != arch[i]]
child[i] = random.choice(others)
return child
# each member is (arch, acc): trained once, kept
pop = collections.deque()
history = []
while len(pop) < P:
a = [random.choice(OPS) for _ in range(L)]
pop.append((a, score(a)))
history.append(pop[-1])
while len(history) < BUDGET:
sample = random.choices(list(pop), k=S)
parent = max(sample, key=lambda m: m[1])[0]
child = mutate(parent)
pop.append((child, score(child)))
history.append(pop[-1])
# drop the oldest
pop.popleft()
best = max(history, key=lambda m: m[1])[0]mutate changes one layer to a different op.score runs once per child, 100 times in all.The paper used P = 100 and S = 25, 25 epochs per child and 20,000 models on 450 K40 GPUs, about 7 days; NNI's RegularizedEvolution defaults to the same P and S. A larger S picks parents more greedily (S = 1 is random search); a larger P remembers more but adapts more slowly. Repeat architectures are retrained on purpose, so only those that retrain well survive.
The per-child budget sets the bill. The first NAS (Zoph and Le 2017) trained every child for 50 epochs (Figure 6 has the totals). MnasNet trained each of about 8K models for 5 epochs on ImageNet, 4.5 days on 64 TPUv2 devices, with a reward that also charged for latency, the time one prediction takes on the target hardware.
Implementations include MnasNet's code (tensorflow/tpu) and NNI's ready-made random, evolution and RL (PolicyBasedRL) strategies.
Each method here trains every candidate separately, thousands of GPU-days per search. Can the choice itself be trained by gradient descent, inside one network? One-Shot NAS: Weight Sharing and DARTS does that.