[Paper Review] AtomNAS: Fine-Grained End-to-End Neural Architecture Search

Jieru Mei, Yingwei Li+5 more|arXiv (Cornell University)|Dec 20, 2019

Advanced Neural Network Applications33 references42 citations

TL;DR

AtomNAS introduces a fine-grained atomic-block search space and end-to-end differentiable NAS with dynamic pruning and FLOPs-aware regularization, achieving state-of-the-art mobile ImageNet results with reduced search cost.

ABSTRACT

Search space design is very critical to neural architecture search (NAS) algorithms. We propose a fine-grained search space comprised of atomic blocks, a minimal search unit that is much smaller than the ones used in recent NAS algorithms. This search space allows a mix of operations by composing different types of atomic blocks, while the search space in previous methods only allows homogeneous operations. Based on this search space, we propose a resource-aware architecture search framework which automatically assigns the computational resources (e.g., output channel numbers) for each operation by jointly considering the performance and the computational cost. In addition, to accelerate the search process, we propose a dynamic network shrinkage technique which prunes the atomic blocks with negligible influence on outputs on the fly. Instead of a search-and-retrain two-stage paradigm, our method simultaneously searches and trains the target architecture. Our method achieves state-of-the-art performance under several FLOPs configurations on ImageNet with a small searching cost. We open our entire codebase at: https://github.com/meijieru/AtomNAS.

Motivation & Objective

• Motivate the need for a finer-grained NAS search space beyond traditional block-level choices.
• Propose atomic blocks as the minimal search unit to enable mixed operations and flexible channel allocation.
• Develop an end-to-end NAS framework that jointly trains network weights and atomic-block importance while enforcing FLOPs-aware regularization.
• Introduce dynamic network shrinkage to accelerate search by pruning dead atomic blocks on the fly.
• Demonstrate state-of-the-art performance on ImageNet under mobile FLOP constraints and show transferability to detection/segmentation.

Proposed method

• Define a fine-grained search space where a building block is decomposed into atomic blocks formed by two convolutions connected by a channel-wise operation.
• Represent the overall network as a sum over atomic blocks with learnable importance factors alpha to enable end-to-end selection of blocks (Eq. 4).
• Apply FLOPs-aware regularization (L1 penalty on alpha weighted by cost coefficients) to encourage low-cost blocks to be preferred (Eq. 5–6).
• Use dynamic network shrinkage to remove dead atomic blocks during training based on evolving alpha values (Algorithm 1).
• Recalculate BN statistics after pruning to ensure accurate inference performance.
• Train the supernet end-to-end with end-to-end differentiable search and pruning, achieving final architectures without post-hoc fine-tuning.

Experimental results

Research questions

• RQ1 Does a finer-grained atomic-block search space enable better trade-offs between accuracy and FLOPs compared to traditional block-level NAS?
• RQ2 Can end-to-end differentiable search with FLOPs-aware regularization yield architectures that outperform mobile- and compact-design baselines under fixed FLOP budgets?
• RQ3 Does dynamic on-the-fly pruning (shrinking) reduce search cost without sacrificing final accuracy?
• RQ4 How well do AtomNAS-generated models transfer to detection and instance segmentation tasks?
• RQ5 What is the impact of BN recalibration on final accuracy after pruning?

Key findings

• AtomNAS-C achieves 75.9% top-1 accuracy at 360M FLOPs on ImageNet, surpassing several state-of-the-art mobile models.
• AtomNAS-C+ reaches 77.6% top-1 accuracy at 363M FLOPs, outperforming MixNet-M and setting a new mobile-scenario SOTA under 400M FLOPs.
• AtomNAS-A and AtomNAS-B also attain strong results under their respective FLOPs budgets (e.g., AtomNAS-A 74.6% at 258M FLOPs, AtomNAS-B 75.5% at 326M FLOPs).
• In ablations, resource-aware regularization improves FLOPs efficiency and accuracy trade-offs compared to equal-penalty baselines.
• BN recalibration after pruning yields notable gains in top-1 accuracy (AtomNAS-A +1.4, AtomNAS-B +1.7, AtomNAS-C +1.2).
• Dynamic shrinkage reduces search cost by pruning dead blocks early, with total search+train time for AtomNAS-C around 25.5 hours (final from-scratch training ≈22 hours).
• AtomNAS models transfer well to COCO detection/segmentation, outperforming baselines on detection and segmentation metrics.