[Paper Review] ProxylessNAS: Direct Neural Architecture Search on Target Task and Hardware

Han Cai, Ligeng Zhu+1 more|arXiv (Cornell University)|Dec 2, 2018

Neural Networks and Applications1,280 citations

TL;DR

ProxylessNAS directly learns neural architectures on the target task and hardware with reduced memory and compute, enabling large search spaces and hardware-aware specialization without proxy tasks.

ABSTRACT

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$\ imes$ fewer parameters. On ImageNet, our model achieves 3.1\\% better top-1 accuracy than MobileNetV2, while being 1.2$\ imes$ 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.

Motivation & Objective

• Motivate NAS that eliminates proxy tasks to optimize directly on large-scale datasets (e.g., ImageNet).
• Enable large search spaces without repeating blocks and reduce memory/compute to regular training levels.
• Introduce path-level pruning via binarized architecture parameters to lower memory usage.
• Handle non-differentiable hardware metrics (e.g., latency) through differentiable latency modeling or REINFORCE-based methods.
• Demonstrate architecture specialization for different hardware (GPU, CPU, mobile) with latency-aware objectives.

Proposed method

• Construct an over-parameterized network containing all candidate paths as mixed operations.
• Binarize architecture parameters to activate only one path at run-time, reducing memory to standard training levels.
• Train weight parameters with frozen architecture parameters and update architecture parameters using a gradient-based estimate (BinaryConnect-inspired) via two-path sampling to keep memory small.
• Model latency as a differentiable regularization term by predicting per-path latency and adding its expected value to the loss with a latency weight λ2.
• Provide an alternative REINFORCE-based update for architecture parameters when using non-differentiable latency objectives.
• Evaluate on CIFAR-10 and ImageNet, including hardware-aware searches for mobile, GPU, and CPU.

Experimental results

Research questions

• RQ1 Can NAS directly optimize architectures on large-scale tasks (e.g., ImageNet) without proxy tasks?
• RQ2 Does allowing all blocks to be learned (no repeated motif restriction) improve performance and efficiency?
• RQ3 Can memory and computation be kept at regular-training levels while exploring large path-based search spaces?
• RQ4 How effectively can latency be integrated as a differentiable objective to produce hardware-aware architectures?
• RQ5 Do hardware-specialized architectures differ across GPUs, CPUs, and mobile devices, and can NAS capture these distinctions?

Key findings

• On CIFAR-10, ProxylessNAS achieves 2.08% test error with 5.7M parameters (surpassing AmoebaNet-B with ~6× more params).
• On ImageNet, Proxyless-G achieves 75.1% top-1 accuracy, 1.2× faster latency than MobileNetV2, with 200× lower search cost than prior methods.
• Proxyless-G (mobile) achieves top-1 74.6% with 78 ms mobile latency, outperforming MobileNetV2 under similar latency constraints.
• Proxyless-NAS finds hardware-specialized architectures for GPU, CPU, and mobile, showing distinct architectural preferences per platform (e.g., GPU: shallower/wider; CPU: deeper/narrower).
• Latency-aware search (with latency regularization) yields better accuracy-latency trade-offs than latency-naive approaches.
• The approach reduces memory by path binarization and enables large architectural search spaces without repeating blocks.