ProxylessNAS: Direct Neural Architecture Search on Target Task and Hardware

In plain words

Designing a neural network by hand takes expert time. Neural architecture search (NAS) automates that by trying many candidate designs and keeping the best. The catch is cost: most NAS methods are so expensive that they cannot run on a large dataset such as ImageNet. Researchers usually work around this by searching on a smaller stand-in task and then reusing the result, but the design that wins on the stand-in is not guaranteed to be best on the real task.

The paper's idea is to search on the real task and the real device. It builds one large network that contains every candidate choice at each position, then trains it. The trick is that at any moment only one candidate is active, chosen by a random draw weighted by learned scores. That keeps memory use close to what training one normal network would need. After training, the highest-scoring candidate at each position is kept and the rest are removed.

The search also takes the speed of the target device into account. Measured latency cannot be differentiated, so the paper replaces it with a smooth estimate built from a table of per-operation times. That estimate is added to the loss, so the search can trade accuracy against speed on a chosen phone, GPU or CPU. The authors report that architectures tuned for one kind of hardware are often slow on another.

TL;DR

ProxylessNAS runs architecture search on the target dataset and target hardware, without a proxy task. Each candidate path in a mixed operation is sampled as a binary gate, so only one path is held in memory during training. Architecture scores are updated with a gradient rule, or with REINFORCE for non-differentiable rewards. Latency enters the loss as an expected-latency term. Author-reported results are competitive on CIFAR-10 and ImageNet mobile and GPU settings.

Problem

• Conventional NAS costs thousands of models trained on the target task, so searching ImageNet directly is impractical §1.
• Cheaper NAS methods search on proxy tasks (fewer epochs, a smaller dataset, or fewer blocks). Blocks found this way are not guaranteed to be optimal on the target task, especially for hardware latency §1.
• Proxy-based methods must reuse one motif stacked many times for transfer, which limits block diversity §1.
• Differentiable NAS methods (One-Shot, DARTS) need about N times the GPU memory and GPU hours of training one compact model, where N is the number of candidate paths. Memory therefore grows linearly with the candidate set, and on ImageNet this can exceed hardware limits §1 §2 §3.1 eq. 1.
• Latency is not differentiable, so it cannot be optimized by gradient descent on the loss §3.3.

Method

Over-parameterized network §3.1. Each edge of the search DAG is a mixed operation with N parallel candidate paths fig. 2. One-Shot sums the path outputs; DARTS takes a softmax-weighted sum over real-valued architecture parameters eq. 1.

Binarized paths §3.2. Each mixed operation has N architecture parameters α. A binary gate vector g is sampled with probability p = softmax(α) and exactly one entry is 1 eq. 2. The mixed output is then the chosen path's output alone eq. 3. Because only one path is active, only its activations are stored §3.2.

Training §3.2.1 fig. 2. Weights and architecture parameters are trained alternately:

1. Freeze α. For each training batch, sample new gates and update the weights of the active paths with standard gradient descent on the training set.
2. Freeze the weights. Reset the gates and update α on a validation set.
3. Repeat. After training, derive the architecture by keeping the path with the highest path weight in each edge §3.2.1.

Gradient for α §3.2.1 eq. 4 app. D. The gradient with respect to α_i is approximated by replacing ∂L/∂p_j with ∂L/∂g_j and applying the softmax Jacobian eq. 4. Computing ∂L/∂g_j needs the output of every path, which would bring back the memory cost. Appendix D describes how to compute these terms in the backward pass, so the unused outputs need not be stored app. D. The main text also describes a cheaper approximation: sample two paths from p, mask the rest, update only the two sampled α values using eq. 4, then rescale them so the unsampled paths keep their probabilities §3.2.1.

Latency as a differentiable cost §3.3.1 fig. 3. The expected latency of a mixed operation is the sum over candidates of path probability times that candidate's predicted latency eq. 5. Blocks run in sequence, so the network's expected latency is the sum over blocks eq. 6. The loss adds this term with a scale λ2 that sets the accuracy-latency trade-off eq. 7. The latency of each candidate comes from a predictor described in app. B fig. 5.

REINFORCE alternative §3.3.2 eq. 8. The architecture parameters can instead be trained with the REINFORCE policy-gradient estimator, using a reward R of the sampled compact network. This needs no differentiable reward, so it can handle non-differentiable objectives such as latency. The two update rules can be combined §3.3.2.

Search spaces.

• CIFAR-10: a tree-structured cell space from prior work, applied to a PyramidNet backbone. Each cell has depth 3 and two branches per node; B blocks per stage and F final channels are the depth and width controls. A full architecture is 648 decisions §4.1. The seven candidate operations are listed in app. A.
• ImageNet: a MobileNetV2-style backbone in which each block chooses a depthwise-separable MBConv with kernel size in {3, 5, 7} and expansion ratio in {3, 6}, plus a zero (skip) option where a residual connection exists §4.2.

# One joint search step, gradient variant  [§3.2] [§3.2.1]
# Input:  train_batch, valid_batch, edges (each holds N alphas)
# Output: updated weights and alphas

def search_step(train_batch, valid_batch, edges, step):
    # 1. Sample one path per edge from softmax(alpha)  [§3.2] [eq. 2]
    for edge in edges:
        probs = softmax(edge.alphas)                  # (N,)
        edge.active = sample_one_index(probs)
    # 2. Update weights of the active paths on training data  [§3.2.1]
    loss = cross_entropy(model(train_batch.images), train_batch.labels)
    weights = sgd_step(weights, gradient(loss, weights))
    # 3. Every k steps, update alphas on validation data  [§3.2.1]
    if step % k == 0:
        # 3a. Sample two paths per edge without replacement  [§3.2.1]
        pairs = sample_two_indices(probs)
        # 3b. Gradient of the loss w.r.t. each sampled gate  [eq. 4]
        gate_grads = gate_gradients(valid_batch, pairs)
        # 3c. Chain rule through the softmax  [eq. 4]
        alphas = softmax_chain_update(alphas, gate_grads, pairs)
        # 3d. Rescale the pair so other paths keep their weight  [§3.2.1]
        alphas = rescale_pair(alphas, pairs)
    return weights, alphas

# Expected latency and training objective  [§3.3.1]
# Input:  path probabilities per block, latency of each candidate path
# Output: scalar loss

def expected_block_latency(probs, path_latency):
    # Probability-weighted latency of the candidates in one block  [eq. 5]
    total = 0
    for j in range(len(probs)):
        total = total + probs[j] * path_latency[j]
    return total

def search_loss(ce_loss, blocks, lambda_2):
    # Sum the expected latency over the sequential blocks  [eq. 6]
    expected = 0
    for block in blocks:
        expected = expected + expected_block_latency(block.probs, block.lat)
    # Add the latency term, scaled by lambda_2  [eq. 7]
    return ce_loss + lambda_2 * expected

Key equations (numbers as they appear in the compiled paper):

\text{One-Shot: } m^{\text{One-Shot}}(x) = \sum_{i=1}^N o_i(x), \qquad m^{\text{DARTS}}(x) = \sum_{i=1}^N \frac{\exp(\alpha_i)}{\sum_j \exp(\alpha_j)}\, o_i(x) eq. 1

g = \text{binarize}(p_1,\dots,p_N), \quad g = [1,0,\dots,0] \text{ w.p. } p_1, \dots, [0,\dots,0,1] \text{ w.p. } p_N eq. 2

\frac{\partial L}{\partial \alpha_i} \approx \sum_{j=1}^{N} \frac{\partial L}{\partial g_j}\, p_j\, (\delta_{ij} - p_i) eq. 4

\mathbb{E}[\text{latency}_i] = \sum_j p_j^i \times F(o_j^i), \qquad \mathbb{E}[\text{latency}] = \sum_i \mathbb{E}[\text{latency}_i] eq. 5 eq. 6

\text{Loss} = \text{Loss}_{CE} + \lambda_1 \lVert w \rVert_2^2 + \lambda_2\, \mathbb{E}[\text{latency}] eq. 7

\nabla_\alpha J(\alpha) \approx \frac{1}{M} \sum_{i=1}^{M} R(\mathcal{N}_{g^i})\, \nabla_\alpha \log p(g^i) eq. 8

Figures, described in words: fig. 1 is an overview of the approach and its headline mobile result. fig. 2 shows, for one edge, the parallel paths, the single sampled path used during weight training, and the alternating architecture-update step. fig. 3 shows the expected-latency term added to the loss. fig. 4 plots accuracy against mobile latency for width-rescaled models, with ProxylessNAS above MobileNetV2 across the latency range. fig. 5 shows predicted against measured latency, close to the identity line. fig. 6 shows the searched GPU, CPU and mobile architectures.

Hyperparameters

Values marked "paper" are stated in the paper. Values marked "code" are defaults in the official repository and are not stated in the paper. Where the code and paper disagree, see Paper vs code.

Results

Author-reported. The paper reports single headline numbers without error bars or run counts. Each table below is reproduced from the paper, with its citation.

CIFAR-10 test error tab. 1, with Cutout where marked c/o.

ImageNet, mobile latency ≤ 80 ms tab. 2, search cost in GPU hours.

LL is the latency regularization loss. Details of MnasNet's search cost are in app. C.

ImageNet, GPU latency on V100 tab. 3

Specialization across hardware tab. 4

Setup: ImageNet classification, with CPU latency at batch 1 and GPU latency at batch 8 on V100 §4.2. Headline comparisons stated in the text: 3.1 points above MobileNetV2 and 1.2× faster on GPU §4.2; 2.6 points above MobileNetV2 at similar mobile latency, and 1.83× faster at matched accuracy §4.2 fig. 4; 0.6 points above MnasNet at slightly lower mobile latency §4.2. Compared with MnasNet, the search uses 200 times fewer GPU hours §4.2 tab. 2.

CIFAR-10 setup §4.1. The text reports Proxyless-G at 2.08% error, which is slightly better than AmoebaNet-B at 2.13% with 6 times more parameters. Proxyless-G and Proxyless-R reach similar or lower error than PathLevel EAS, with 5.7M and 5.8M parameters against 14.3M §4.1 tab. 1.

Paper vs code

The official repository is at MIT-HAN-LAB/ProxylessNAS (linked from the paper's abstract footnote), commit b23018c9c369d22931f7422b71ca6a7eaa354c46. The search code is under search/. The repository contains no CIFAR-10 search script, so every CIFAR-10 search setting in the paper is unverified.

• Latency loss form. The paper adds λ2 × expected latency to the loss eq. 7. The code default is a multiplicative form, the cross-entropy scaled by α × (log E / log T)^β with α = 0.2 and β = 0.3 nas_manager.py imagenet_arch_search.py. A linear alternative, λ × (E − T) / T, is also in the code nas_manager.py imagenet_arch_search.py. The paper does not say which form produced its ImageNet results.
• Gradient mode. The main text describes sampling two paths per update §3.2.1. The code default is full_v2, which evaluates the gradient for every gate in the backward pass imagenet_arch_search.py mix_op.py. The two-path variant is the two mode mix_op.py. The paper does not say which mode produced its results.
• Latency predictor. The paper describes a learned model fitted on 4,000 sampled architectures, with RMSE 0.75 ms app. B fig. 5. The code is a lookup table of mean times, keyed by input and output size, expansion, kernel, stride and skip flag. It is downloaded at runtime from an external URL latency_estimator.py. The repository contains no fitting or evaluation code for the RMSE claim.
• Expected-latency composition. Eq. 6 sums only the mixed operations. The code adds fixed latencies for the first convolution, the feature-mix layer and the classifier super_proxyless.py, and also weights each candidate by its path probability super_proxyless.py.
• REINFORCE reward. The paper's objective is ACC × [LAT/T]^w, with w not stated §4.2. The code computes ACC × (T/LAT)^0.1, with the exponent 0.1 as a default nas_manager.py imagenet_arch_search.py. In the paper's notation this corresponds to w = −0.1. This is an inference from the formulas, not a stated value.
• Warm-up. The paper does not mention warm-up. The ImageNet search script trains weights alone for 40 epochs before the joint search imagenet_arch_search.py imagenet_arch_search.py nas_manager.py.
• Architecture update frequency. The paper says the two updates alternate, with no frequency given §3.2.1. The code updates α every five batches and skips the first epoch imagenet_arch_search.py nas_manager.py nas_manager.py.
• Fixed parts of the ImageNet space. The first MBConv block is fixed to 3×3 kernel and expansion 1, and is not searched super_proxyless.py. The paper does not state this.
• Zero operation gradient. When the zero operation is the active path, the code sets the α gradient to None, so the architecture parameters of that edge are not updated in that step mix_op.py. The paper does not mention this case.
• Validation BatchNorm. During search-time validation the code runs BatchNorm in training mode nas_manager.py. The paper does not say this. It affects the validation accuracies logged during search.
• Final training. The paper says the compact network trains for 300 epochs instead of 200 §4.1 §4.2. The code default is 300 imagenet_run_exp.py, which matches the final setting.
• No discrepancy found for the ImageNet validation split size (50,000) §4.2 imagenet_arch_search.py, the 256 batch size §4.2 imagenet_arch_search.py, or the 0.001 architecture learning rate §4.2 imagenet_arch_search.py.

Pitfalls

• Memory claim depends on the mode. The memory saving holds only when one path is kept active. The full mode keeps the outputs of all paths, so it uses about N times the memory mix_op.py. Use full_v2 or two to get the stated saving.
• Biased gate gradient. The α update uses ∂L/∂g in place of ∂L/∂p eq. 4. This is an approximation, and the paper gives only an intuitive justification §3.2.1. The paper does not measure its bias.
• Two-path renormalization is ambiguous. The paper says the two sampled paths' probabilities are "reset accordingly" §3.2.1. The code computes the softmax over the two sampled α values only mix_op.py, and rescales them in log space mix_op.py. A reimplementation that normalizes over all N paths will differ.
• Logarithm in the latency term. The code's mul#log form takes log(E) and raises the ratio to power β = 0.3. This is defined only when E > 1 ms, since otherwise the base is zero or negative nas_manager.py. This is an inference from the formula. The reference of 80 ms is a unit choice that must match the latency table's units imagenet_arch_search.py.
• Unsampled weights get no update in that step. The code sets the gradients of all candidate operations to None before each step mix_op.py. PyTorch optimizers skip parameters whose gradient is None, so unsampled paths also receive no weight decay or momentum that step. This is an inference from PyTorch's behaviour, not something the paper states.
• Latency table dependency. The code downloads its latency table from an external URL and overwrites any cached copy on each construction latency_estimator.py. Results can change if that file changes. Shapes absent from the table are not handled latency_estimator.py.
• Hardware specificity. Latency numbers apply to the listed devices only: a V100 at batch 8, a Pixel 1 at batch 1, and a two-socket Xeon at batch 1 §4.2. In the paper's Table 4, the model searched for one device is slower on the other two tab. 4.
• Search cost is underspecified. The 200 GPU-hour figure does not state the GPU type or count, or whether warm-up is included. The paper says the MnasNet comparison assumes V100 GPUs "as done in our experiments" app. C, which only implies the GPU type used.
• Single-run comparisons. The tables give one number per configuration, with no error bars or seed counts. Some differences are small: for example, Proxyless-R (74.6) and Proxyless-G + LL (74.2) differ by 0.4 points tab. 2.
• Loss scale. The paper does not state λ2 eq. 7, so the latency weight must be tuned. The code's α and β defaults are the only concrete values, and the paper does not say whether they produced Table 2 imagenet_arch_search.py.

Lineage

• Builds on: BinaryConnect, for training binarized weights §3.2; REINFORCE, for the alternative update rule §3.3.2; DARTS and One-Shot, for the over-parameterized formulation and their memory cost §2 §3.1; the tree-structured space from PathLevel EAS for CIFAR-10 §4.1; and MobileNetV2 as the ImageNet backbone §4.2.
• Compared against: NASNet-A, AmoebaNet, MnasNet, DARTS, MobileNetV1, MobileNetV2, ShuffleNetV2 and ResNet-34 tab. 1 tab. 2 tab. 3.
• Superseded by: the repository README names Once-for-All as the next generation of the work README.md. This card does not cover that paper.