Neural Architecture Search on ImageNet in Four GPU Hours: A Theoretically Inspired Perspective

In plain words

Finding a good neural network design usually means training many candidate designs and seeing which one works best. That is slow and expensive, and shortcuts such as training each candidate for only a little while can mislead the search.

This paper asks whether a candidate can be judged before any training. It uses two properties of a network at its random starting point. One says how easily the network can be trained by gradient descent. The other says how many different input regions the network can separate. Neither needs labels.

The search starts from a large network that contains every candidate operation, then removes weak operations in rounds until one path is left. Each round ranks the operations by how much removing them changes the two properties, and the paper argues this trade-off gives a good design.

TL;DR

A 2021 method for neural architecture search with no training of candidates. It scores operations at initialisation with two indicators: the condition number of the NTK, a trainability measure, and a count of linear regions, an expressivity measure §3.1. It removes operations by a rank sum over both indicators, one per edge per round §3.2 alg. 1. Author-reported: it matches or beats the compared NAS methods on NAS-Bench-201 and DARTS while using far less search time §4 tab. 1 tab. 2 tab. 3. Its search cost figures are not comparable with training-based methods on identical hardware, so check the units in the tables.

Problem

Neural architecture search needs to know which candidate architecture is good. Existing methods either train a super-network or train many candidates, which costs a lot of GPU time, and truncated training can bias the ranking §1. The paper asks whether trainability and expressivity, both measurable at initialisation without labels, can rank architectures well enough to replace that evaluation §1 §3.1. It targets cell-based search spaces where the number of candidate cells is |O|^E for E edges and |O| operators, so sampling every candidate is not feasible §3.2.

Method

Indicator 1: trainability by the NTK condition number. The finite-width NTK is J(x) J(x')^T, where J is the Jacobian of the outputs with respect to the weights §3.1.1. Under gradient descent the training dynamics can be written in terms of the NTK's eigenvalues eq. 1 eq. 2. The paper cites the hypothesis of Lee et al. (2019) that the largest feasible learning rate scales as 2 divided by the largest eigenvalue. Under that assumption the smallest eigenvalue converges at a rate set by the condition number §3.1.1. The paper therefore uses the condition number as a trainability score, where lower is better eq. 3. It is computed without gradient descent or labels §3.1.1.

Indicator 2: expressivity by the number of linear regions. For a ReLU network, a linear region is the set of inputs that share one activation pattern eq. 4. The paper averages the region count over sampled weight initialisations, where higher is better eq. 5 §3.1.2. In practice it counts distinct activation patterns on sampled inputs app. A.

Combining the two. The two values have very different scales, so the paper ranks operators instead of using raw values. An equal-weight sum of the two ranks is the default. Weighted sums and other combinations gave lower accuracy in the paper's ablations §3.2 app. D.

Pruning-by-importance search. The search starts from a super-network that contains every operation on every edge alg. 1. Each round does the following.

1. For every active operation, remove it alone and measure the change in the condition number (Δκ) and in the region count (ΔR) alg. 1.
2. Rank the candidates. Largest Δκ gets rank 0, since removing it helps trainability most. Smallest ΔR gets rank 0, since removing it costs expressivity least. Both rankings say "prune first" for low index §3.2 alg. 1.
3. Add the two ranks. On each edge, remove the candidate with the lowest rank sum alg. 1 §3.2.
4. Stop when the network is a single path alg. 1.

Each round removes one operator per edge, so the number of networks scored falls from |O|^E to about |O|·E §3.2. The paper also shows how the trainability and expressivity values change during the search §3.2.1 fig. 5.

# TE-NAS pruning loop  [alg. 1] [§3.2]
# Input: supernet with every candidate operator on every edge

def te_nas(supernet):
    # Repeat rounds until each edge keeps a single operator  [alg. 1]
    while not is_single_path(supernet):
        d_kappa = {}  # (edge, op) -> change in NTK condition number
        d_regions = {}  # (edge, op) -> change in linear-region count
        # Remove each active op on its own and measure the change  [alg. 1]
        for edge, op in active_ops(supernet):
            # Larger drop in kappa means removing the op helps trainability
            # [§3.1.1]
            d_kappa[edge, op] = kappa(supernet) - kappa(supernet - op)
            # Smaller drop in regions means removing the op costs little
            # expressivity [§3.1.2]
            d_regions[edge, op] = regions(supernet) - regions(supernet - op)
        # Rank 0 is the most prunable for each indicator  [§3.2] [alg. 1]
        s_kappa = rank_descending(d_kappa)
        s_regions = rank_ascending(d_regions)
        # Equal-weight sum of the two ranks  [§3.2]
        importance = {c: s_kappa[c] + s_regions[c] for c in s_kappa}
        # On each edge, the lowest rank sum is removed first  [alg. 1]
        for edge in edges(supernet):
            candidates = candidates_on(edge, importance)
            victim = min(candidates, key=lambda c: importance[c])
            supernet = supernet - victim
    return supernet

Trainability score, computed on one mini-batch §3.1.1 app. A:

# Condition number of the NTK on one training batch  [§3.1.1] [eq. 3]

def ntk_condition(network, batch):
    # One gradient row per sample, over weight parameters only
    # [code: lib/procedures/ntk.py#L44-L50]
    grads = [sample_gradient(network, x) for x in batch]  # (N, P)
    # Gram matrix of the gradients, the NTK over the batch
    # [code: lib/procedures/ntk.py#L55]
    ntk = grads @ grads.T  # (N, N)
    eigenvalues = sorted(eigvalsh(ntk))  # ascending
    # Largest over smallest eigenvalue  [eq. 3]
    return eigenvalues[-1] / eigenvalues[0]

Expressivity score, on sampled inputs §3.1.2 app. A:

# Distinct ReLU activation patterns over sampled inputs  [§3.1.2] [eq. 5]

def linear_regions(network, samples):
    # Sign of every ReLU output, flattened per sample into one 0/1 row
    # [code: lib/procedures/linear_region_counter.py#L126]
    patterns = activation_signs(network, samples)  # (S, D), 0 or 1
    # Number of samples with the same pattern as each sample, itself included
    # [code: lib/procedures/linear_region_counter.py#L129-L135]
    same = count_identical_rows(patterns)  # (S,)
    # Each group of identical patterns adds exactly one region in total
    # [code: lib/procedures/linear_region_counter.py#L136]
    return sum(1 / same)

Hyperparameters

Values stated in the paper and the values in the official code differ in several places. Both are listed. "Not stated" means neither the paper nor the code sets it.

Results

Author-reported. The paper runs TE-NAS four times with different seeds for NAS-Bench-201 and CIFAR-10 §4.1 §4.2; the ImageNet run count is not stated §4.3. Search costs are as printed; units differ between tables.

NAS-Bench-201, test accuracy (%), search on each dataset tab. 1

CIFAR-10, DARTS space, test error (%) tab. 2

The paper marks the ProxylessNAS row as a different space with a PyramidNet backbone, and the PNAS row as trained without cutout augmentation tab. 2.

ImageNet, DARTS space, mobile setting tab. 3

The paper's ImageNet claim is the searched-on-ImageNet row: 24.5% top-1 and 7.5% top-5 error in about four GPU hours on one 1080Ti §4.3.

Ablations, NAS-Bench-201, CIFAR-100 test accuracy (%)

Only trainability, only expressivity, and both tab. 4:

Other ways to combine the two indicators tab. 5:

Paper vs code

Official repository: VITA-Group/TENAS, linked from the abstract. The README clone command names chenwydj/TENAS; this card follows the paper's link. The code was cloned at commit 9df78ffd98573035375b12e19b9007578cc4155d. The repository licence is MIT LICENSE. Only the search code is included; retraining and evaluation of DARTS architectures live in another repository named in the README. NAS-Bench-201 accuracies come from the benchmark database README.md.

• Weight initialisation is not applied. The paper says the weights are drawn from Kaiming normal initialisation for each score app. A §3.1.2. The code's init_model only handles the strings kaiming_norm_fanin and kaiming_norm_fanout prune_tenas.py. The launcher passes kaiming_normal prune_launch.py, and the repeat loop appends _fanout or _fanin to it prune_tenas.py prune_tenas.py. Neither string matches, so the call does nothing. The networks keep the weights from their constructor, and the repository has no custom initialisation models. Within one pruning round the three repeats therefore share weights and differ only in the batches drawn.
• The none operator is removed before search. The paper's operator spaces include none §4.1 §4.2. The code sets the first operator of every edge to -INF before search and during each scoring call prune_tenas.py prune_tenas.py. Index 0 is none in both spaces cell_operations.py.
• Region counting uses a different network and inputs. The paper forwards 5000 images through "the network" app. A. The code scores regions with a thin network: one channel, no stem, depth 1 for NAS-Bench-201 and depth 2 for DARTS prune_tenas.py prune_tenas.py, with 3000 patches of size 3x3 prune_tenas.py linear_region_counter.py. The paper does not state a thin network or a patch size. The patches come from random 3x3 crops with one randomly chosen colour channel linear_region_counter.py. For imagenet-1k the image is first resized to 32x32 linear_region_counter.py.
• Batch sizes differ. The paper uses 32 app. A. The launcher uses 72, 14 and 24 prune_launch.py prune_launch.py prune_launch.py.
• Ranks use relative changes, averaged over repeats. The paper defines Δκ and ΔR as absolute differences alg. 1. The code divides each change by the unpruned value, then averages the three repeats prune_tenas.py prune_tenas.py. Ties share a rank prune_tenas.py.
• Ranking is global; selection is per edge. All candidates from both cell types are ranked together, then each edge keeps the lowest-sum candidates up to prune_number prune_tenas.py.
• Operators per round differ from the paper's text. The paper says one operator per edge per round §3.2. The DARTS launcher removes three (CIFAR-10) or two (ImageNet) prune_launch.py prune_launch.py.
• A second DARTS stage is not in the paper. After the single-path loop, DARTS runs a group-wise pruning that keeps two operators in each of four edge groups prune_tenas.py prune_tenas.py. The paper's footnote says the DARTS result keeps two operators per edge §3.2 fn. 2. The loop's stop test requires one operator per edge prune_tenas.py. The paper does not describe the group stage, and I did not resolve how the two stages fit together.
• The NTK kernel sums over outputs. The paper writes the NTK with a Jacobian over output units §3.1.1. The code backpropagates a vector of ones through the logits, so each sample's gradient is the gradient of the sum of its outputs ntk.py. The paper does not say how output units are combined.
• Only weight gradients enter the NTK. Parameters whose names contain weight are used; biases are excluded. This includes batch-norm scale parameters ntk.py. The paper does not state which parameters are used.
• Networks run in training mode. Scores are computed with batch statistics in batch-norm layers ntk.py prune_tenas.py. The paper does not say which mode is used.
• Algorithm 1 wording conflicts with the text. The algorithm's comment on the final argmin says "greatest importance", while §3.2 says the lowest importance is pruned alg. 1 §3.2. The code removes the lowest rank sum prune_tenas.py.

Pitfalls

• Search cost units. Table 1 gives GPU seconds, while Tables 2 and 3 give GPU days. The abstract's "0.5 GPU hours" for NAS-Bench-201 on CIFAR-10 matches Table 1's 1558 GPU seconds (about 0.43 hours), not the DARTS CIFAR-10 figure of 0.05 GPU days tab. 1 tab. 2. Do not mix the two.
• Baseline rows in Table 1 look copied. ENAS, DARTS (1st) and DARTS (2nd) all show 54.30, 15.61 and 16.32 with zero deviation tab. 1. This looks like a table error, so do not use those rows as baselines.
• Search evaluates many candidates. Every active operator is scored in each round, and each score repeats the NTK and region counts three times with freshly built networks prune_tenas.py. The NTK needs one backward pass per sample ntk.py. The run needs a CUDA device; the script asserts it prune_tenas.py.
• Following the paper alone gives a different search. Batch size, the thin scoring network and the none removal all change the scores, so a reimplementation from the paper will not reproduce the released choices.
• The condition number depends on batch size. It is computed on the Gram matrix of one batch, so a batch of 32 (paper) and batches of 72, 14 or 24 (launcher) give different values. The paper does not say how batches are drawn app. A.
• Retraining is not in this repository. The appendix's retraining settings app. A cannot be checked against this code. DARTS architectures must be retrained with the external evaluation repository named in the README README.md.
• Expressivity and trainability are measured at initialisation only. The paper's appendix argues that these indicators do not predict generalisation directly, and that NAS-Bench-201 training accuracy tracks test accuracy in its space app. E fig. 10. The method's claim therefore rests on that benchmark property and may not transfer to other search spaces.
• Correlations are from one benchmark. The Kendall-tau values of -0.42 (condition number) and 0.5 (region count) are measured on NAS-Bench-201 §3.1.1 §3.1.2 fig. 1 fig. 3. The paper does not report them for DARTS.
• Not stated in the paper: the rounding precision, the scoring network sizes, the active-operator threshold, tie handling, the NTK batch sampling scheme, and how repeated NTK scores are aggregated.

Lineage

• Builds on: the neural tangent kernel and its finite-width form §2 §3.1.1; linear-region counting of ReLU networks §2 §3.1.2; pruning from scratch (SNIP and related work) §3.2.
• Compared against: NAS w.o. Training, the training-free Jacobian method of Mellor et al. §2 tab. 1; DARTS, PC-DARTS, GDAS, P-DARTS, ProxylessNAS and others tab. 1 tab. 2 tab. 3.
• Superseded by: not stated in this card's sources.