Source-linked AI summary

ProxylessNAS: Direct Neural Architecture Search on Target Task and Hardware

Han Cai, Ligeng Zhu, Song Han

arXiv:1812.00332v2cs.LGcs.CVstat.ML

TL;DR

NAS methods are often too costly for direct search on large-scale tasks, and proxy-task architectures are not guaranteed to be optimal for the target task. ProxylessNAS directly searches architectures on target tasks and hardware using path-level pruning and binarization, achieving strong CIFAR-10 and ImageNet results while enabling hardware specialization.

  • Problem

    Conventional NAS is prohibitively expensive on large-scale target tasks, while proxy-task architectures are not guaranteed to be optimal for those tasks and can restrict block diversity.

  • Method

    ProxylessNAS directly learns target-task and target-hardware architectures using path-level pruning and binarized architecture parameters, with latency regularization for hardware objectives.

  • Results

    ProxylessNAS achieves 2.08% CIFAR-10 test error with 5.7M parameters and 75.1% ImageNet top-1 accuracy, 3.1% higher than MobileNetV2 while being 1.2× faster.

  • Takeaways & Limitations

    Direct search supports large candidate sets without proxy transfer and enables architectures specialized for CPU, GPU, and mobile hardware.

Abstract

from arXiv · show

Neural architecture search (NAS) has a great impact by automatically designing effective neural network architectures. However, the prohibitive computational demand of conventional NAS algorithms (e.g. $10^4$ GPU hours) makes it difficult to \emph{directly} search the architectures on large-scale tasks (e.g. ImageNet). Differentiable NAS can reduce the cost of GPU hours via a continuous representation of network architecture but suffers from the high GPU memory consumption issue (grow linearly w.r.t. candidate set size). As a result, they need to utilize~\emph{proxy} tasks, such as training on a smaller dataset, or learning with only a few blocks, or training just for a few epochs. These architectures optimized on proxy tasks are not guaranteed to be optimal on the target task. In this paper, we present \emph{ProxylessNAS} that can \emph{directly} learn the architectures for large-scale target tasks and target hardware platforms. We address the high memory consumption issue of differentiable NAS and reduce the computational cost (GPU hours and GPU memory) to the same level of regular training while still allowing a large candidate set. Experiments on CIFAR-10 and ImageNet demonstrate the effectiveness of directness and specialization. On CIFAR-10, our model achieves 2.08\% test error with only 5.7M parameters, better than the previous state-of-the-art architecture AmoebaNet-B, while using 6$\times$ fewer parameters. On ImageNet, our model achieves 3.1\% better top-1 accuracy than MobileNetV2, while being 1.2$\times$ faster with measured GPU latency. We also apply ProxylessNAS to specialize neural architectures for hardware with direct hardware metrics (e.g. latency) and provide insights for efficient CNN architecture design.

1 INTRODUCTION

ProxylessNAS addresses the cost and proxy-task limitations of NAS by directly optimizing architectures on target tasks and hardware while reducing search memory and computation to training-level costs. It reports strong CIFAR-10 and ImageNet results and supports hardware-specific specialization.

  • Motivation: Conventional NAS is too computationally expensive for direct search on large-scale tasks, while proxy-optimized architectures are not guaranteed optimal for the target task.Proxy methods can also restrict block diversity by repeatedly stacking a few searched motifs.
  • ProxylessNAS: ProxylessNAS directly learns architectures on the target task and hardware, removes repeated-block restrictions, and allows each block to be specified.This enlarges the search space while avoiding proxy transfer.
  • Efficient search: Path-level pruning trains one over-parameterized network and prunes redundant paths, avoiding a meta-controller during architecture search.Architecture parameters identify redundant paths while the network is trained once.
  • Efficient search: Path-level binarization activates one candidate path at runtime, reducing GPU memory to the level of compact-model training while retaining a large candidate set.The method saves memory by one order of magnitude.
  • Results: 2.08% test error on CIFAR-10 uses 5.7M parameters, while ImageNet reaches 75.1% top-1 accuracy, 3.1% higher than MobileNetV2 and 1.2× faster.These results are reported for the paper’s direct and specialized architecture search.
  • Hardware specialization: Latency regularization enables gradient-based optimization for hardware objectives and specialization across CPU, GPU, and mobile platforms.The approach incorporates hardware considerations directly into architecture learning.

2 RELATED WORK

Prior NAS work reduces evaluation cost through hypernetworks or weight sharing, while One-Shot and DARTS formulate search as training an over-parameterized network. ProxylessNAS is also connected to pruning but changes topology through path-level architecture selection.

  • Efficient NAS evaluation: Hypernetwork and weight-sharing methods reduce architecture-evaluation cost by generating or sharing weights across sampled networks.These approaches remain within a meta-learning or standard NAS evaluation framework.
  • Differentiable and one-shot NAS: One-Shot and DARTS eliminate the meta-controller by modeling NAS as training a single over-parameterized network containing candidate paths.One-Shot uses DropPath with a fixed probability to drop paths during training.
  • NAS and pruning: ProxylessNAS resembles network pruning by starting from an over-parameterized network and removing redundant parts, but it can modify network topology rather than only filter or unit counts.This connects architecture search to model compression while preserving broader structural changes.

3 METHOD

ProxylessNAS constructs an over-parameterized network, trains binarized path choices, and incorporates hardware latency into architecture optimization. These mechanisms reduce search memory and support hardware-aware specialization.

  • 3.1 CONSTRUCTION OF OVER-PARAMETERIZED NETWORK: The search network contains all candidate operations as parallel paths within mixed operations, allowing architectures to be represented in one over-parameterized network.Each edge in the DAG is replaced by a mixed operation whose paths correspond to candidate primitive operations.
  • 3.1 CONSTRUCTION OF OVER-PARAMETERIZED NETWORK: The mixed-operation formulation stores every candidate path during computation, requiring roughly N times the GPU memory and GPU hours of compact-model training.This motivates path binarization for large search spaces.
  • 3.1 CONSTRUCTION OF OVER-PARAMETERIZED NETWORK: Path binarization activates only one candidate path at run-time, reducing over-parameterized-network memory to the level of compact-model training.Unlike weight binarization, ProxylessNAS binarizes entire paths using binary gates.
  • 3.2.1 TRAINING BINARIZED ARCHITECTURE PARAMETERS: Architecture parameters are trained separately from weights: active-path weights use standard gradient descent, while binarized architecture parameters use gradient estimates based on BinaryConnect.Architecture parameters are not directly part of the computation graph, so they require a separate gradient-based update procedure.
  • 3.2.1 TRAINING BINARIZED ARCHITECTURE PARAMETERS: The architecture update factorizes choosing one path among N candidates into binary comparisons by sampling two paths and updating their architecture parameters.This reduces the temporary candidate set from N to 2 during an architecture-parameter update.
  • 3.3.1 MAKING LATENCY DIFFERENTIABLE: Latency is optimized by predicting each candidate operation’s latency, summing expected latencies across sequential mixed operations, and adding the result to the loss with a trade-off coefficient.The latency regularization loss combines cross-entropy, weight decay, and expected latency.
  • 3.3.2 REINFORCE-BASED APPROACH: REINFORCE provides an alternative update for binarized architecture parameters that handles non-differentiable rewards without requiring the reward to be differentiable with respect to binary gates.The update uses sampled binary gates, their sampling probabilities, and the corresponding compact networks.

4 EXPERIMENTS AND RESULTS

Experiments evaluate ProxylessNAS by directly searching task- and hardware-specific architectures on CIFAR-10 and ImageNet, using broad spaces and latency-aware objectives. The method achieves strong accuracy, parameter efficiency, latency, and hardware specialization across platforms.

  • Experimental Setup: ProxylessNAS directly learns architectures on each target task rather than transferring repeatedly stacked cells from a smaller setting.The CIFAR-10 and ImageNet experiments use target-task architecture search, while the CIFAR-10 space removes repeated-cell constraints.
  • CIFAR-10: CIFAR-10 searches require 648 independent architectural decisions because cells are not repeated.The count is 12 learnable edges across 18 blocks and 3 stages.
  • CIFAR-10: 2.08% test error with 5.7M parameters slightly improves on AmoebaNet-B while using 6× fewer parameters.ProxylessNAS also matches or improves PathLevel EAS with half as many parameters.
  • ImageNet: ImageNet experiments optimize accuracy and latency across mobile phone, GPU, and CPU hardware platforms.Mobile latency is measured on Google Pixel 1, GPU latency on a V100 with batch size 8, and CPU latency on a server with batch size 1.
  • ImageNet: The mobile latency predictor closely tracks measured latency, with a reported latency RMSE of 0.75ms.The authors use it during architecture search to avoid an expensive mobile farm infrastructure.
  • ImageNet: On mobile, ProxylessNAS improves MobileNetV2 top-1 accuracy by 2.6% at similar latency and reaches around 74.6% accuracy at 78ms versus 143ms.The mobile model also achieves 0.6% higher top-1 accuracy than MnasNet with slightly lower latency and uses 200× fewer GPU hours for search.
  • ImageNet: Latency regularization improves the resulting model: removing it leaves 71.8% top-1 accuracy, 2.4% below Proxyless-G with latency regularization.The comparison supports treating latency as a direct optimization objective.
  • Hardware Specialization: On GPU, ProxylessNAS is 1.2× faster than MobileNetV2 while improving top-1 accuracy by 3.1%, and models optimized for one platform do not run fastest on others.The experiments motivate hardware-specific architecture specialization.

5 CONCLUSION

ProxylessNAS directly searches architectures on target tasks and hardware without proxies, while reducing NAS cost to normal-training levels through path binarization. It supports hardware specialization and reveals that different platforms favor different model designs.

  • ProxylessNAS directly learns architectures on target tasks and hardware without using proxy tasks.
  • Path binarization reduces NAS GPU-hours and GPU memory to the same level as normal training.
  • GPU favors shallow, wide models with early pooling, whereas CPU favors deep, narrow models with late pooling.
  • 200× lower search cost supports efficient specialization for target hardware architectures.
  • Directly incorporating measured hardware latency enables architectures specialized for different hardware platforms.

A THE LIST OF CANDIDATE OPERATIONS USED ON CIFAR-10

The CIFAR-10 search space uses seven candidate operations spanning depthwise-separable convolutions, identity, pooling, and dilation.

  • Seven operations define the CIFAR-10 candidate space.
  • Candidates include identity and 3 × 3, 5 × 5, and 7 × 7 depthwise-separable convolutions.
  • The space also includes 3 × 3 dilated depthwise-separable convolution and 3 × 3 average and max pooling.

B MOBILE LATENCY PREDICTION

On-device latency measurement is accurate but too slow and expensive for scalable NAS, so the paper builds a latency model from measured mobile architectures.

  • On-device latency measurement requires averaging hundreds of runs, taking approximately 20 seconds per measurement.
  • Mobile latency measurement is expensive because it requires many devices and substantial software engineering for an automated device-farm pipeline.
  • 5k architectures were sampled, with 4k used to build the latency model and the remainder reserved for testing.
  • Latency-model features include operator type, input and output feature-map sizes, and convolution attributes such as kernel size, stride, and expansion ratio.

C DETAILS OF MNASNET’S SEARCH COST

MnasNet’s search cost is estimated at roughly 40,000 GPU hours because it trains 8,000 mobile-sized ImageNet models for five epochs each.

  • 8,000 mobile-sized models are trained on ImageNet for five epochs each during MnasNet architecture search.
  • Roughly 40,000 GPU hours are required when those models are trained on V100 GPUs.

D IMPLEMENTAION OF THE GRADIENT-BASED ALGORITHM

The implementation avoids storing inactive-path outputs during the forward pass by recomputing them during backpropagation after obtaining the mixed-operation gradient. This preserves the exact gradient calculation while reducing GPU memory to compact-model training levels.

  • The naive implementation stores each path output during the forward pass to compute the corresponding architecture-gradient term during backpropagation.The mixed-operation gradient uses the output gradient, elementwise products, and a reduction over elements.
  • Inactive-path outputs need not be stored because they are only required for architecture gradients when their paths do not contribute to the mixed operation.
  • Recomputing inactive-path outputs after obtaining the output gradient preserves the non-approximated gradient calculation while reducing GPU memory to compact-model training levels.The outputs are released after computing the architecture-gradient term.
Loading 1812.00332v2…