Source-linked AI summary
AttentiveNAS: Improving Neural Architecture Search via Attentive Sampling
Dilin Wang, Meng Li, Chengyue Gong, Vikas Chandra
TL;DR
Uniform sampling in two-stage NAS does not prioritize networks on the accuracy–efficiency Pareto front, which is central to search. AttentiveNAS uses Pareto-aware sampling and reports stronger ImageNet accuracy–FLOPs trade-offs than prior approaches.
Problem
Uniform sampling treats candidate networks equally even though two-stage NAS search focuses on the accuracy–efficiency Pareto front.
Source: §1 · p. 2Method
AttentiveNAS samples networks on the best or worst Pareto front and proposes approaches to identify them efficiently during training.
Source: §1 · p. 2Results
Takeaways & Limitations
Attentive sampling can improve accuracy over uniform sampling by taking the performance Pareto into account.
Source: §5 · p. 8Using negative validation loss as the performance estimator can be computationally expensive when the validation set is large.
Source: §3.2 · p. 4
Abstract
from arXiv · showhide
Neural architecture search (NAS) has shown great promise in designing state-of-the-art (SOTA) models that are both accurate and efficient. Recently, two-stage NAS, e.g. BigNAS, decouples the model training and searching process and achieves remarkable search efficiency and accuracy. Two-stage NAS requires sampling from the search space during training, which directly impacts the accuracy of the final searched models. While uniform sampling has been widely used for its simplicity, it is agnostic of the model performance Pareto front, which is the main focus in the search process, and thus, misses opportunities to further improve the model accuracy. In this work, we propose AttentiveNAS that focuses on improving the sampling strategy to achieve better performance Pareto. We also propose algorithms to efficiently and effectively identify the networks on the Pareto during training. Without extra re-training or post-processing, we can simultaneously obtain a large number of networks across a wide range of FLOPs. Our discovered model family, AttentiveNAS models, achieves top-1 accuracy from 77.3% to 80.7% on ImageNet, and outperforms SOTA models, including BigNAS and Once-for-All networks. We also achieve ImageNet accuracy of 80.1% with only 491 MFLOPs. Our training code and pretrained models are available at https://github.com/facebookresearch/AttentiveNAS.
1. Introduction
AttentiveNAS redirects two-stage NAS training from uniform sampling toward networks on the best or worst accuracy–efficiency Pareto front. The paper reports improved ImageNet accuracy–FLOPs trade-offs, including models with higher accuracy and fewer FLOPs than named baselines.
1. Introduction
- Two-stage NAS separates weight-sharing training of candidate networks from a second-stage search for models under resource constraints.This approach has delivered state-of-the-art empirical results with search efficiency.Source: §1 · p. 1
- Uniform sampling treats candidate networks equally, although the search stage focuses on accuracy–efficiency Pareto-front models.The authors say this misses opportunities to improve the accuracy of Pareto-front networks during training.Source: §1 · pp. 1–2
- AttentiveNAS samples networks more likely to produce a better Pareto front, exploring BestUp and WorstUp strategies.BestUp directs training budget toward improving the current best Pareto front; WorstUp targets the worst performance trade-offs as hard training examples.Source: §1 · p. 2
- The authors propose two approaches to guide efficient identification of networks on the best or worst Pareto front during training.They identify training loss and accuracy predicted by a pretrained predictor as proxies for accuracy comparison.Source: §1 · p. 2
- AttentiveNAS-A5 reaches 80.1% accuracy with only 491 MFLOPs; AttentiveNAS-A0 achieves 2.1% better accuracy than MobileNetV3 with fewer FLOPs.AttentiveNAS-A2 achieves 0.8% better accuracy than FBNetV3 with 10% fewer FLOPs.Source: §1 · p. 2
2. Related Work and Background
NAS seeks accurate networks under resource constraints, but conventional search can be prohibitively expensive. Two-stage NAS separates shared-weight pre-training from resource-constrained search, reducing that burden while supporting multiple deployment constraints.
2. Related Work and Background
- Earlier NAS methods train many architectures from scratch and are computationally expensive; weight-sharing reduces this burden by estimating candidates with inherited weights.Some weight-sharing approaches use differentiable relaxation, which can be sensitive to hyperparameters and trials; inherited weights may remain sub-optimal, requiring re-training.Source: §2 · pp. 2–3
2.1. Two-stage NAS
- At each training step, one or several subnetworks are sampled from the search space, with configurations choosing resolution, channel widths, depth, kernel size, and expansion ratio.The sampled subnetworks are optimized through one-step stochastic gradient descent.Source: Figure 2 · p. 3, §2.1 · p. 3
- The typical NAS formulation restricts search to small subnetworks and one resource constraint, limiting over-parameterization benefits and requiring separate searches for different constraints.This makes finding architectures across multiple resource constraints a challenging optimization problem.Source: §2.1 · p. 3
- Two-stage NAS first pre-trains a shared-weight network without resource constraints, then searches for the best-performing subnetworks under specified constraints.The pre-training objective jointly optimizes candidate subnetworks, while the search stage can accommodate a range of deployment constraints without further modification.Source: §2.1 · p. 3
- Pre-training uses shared weights and a regularization example that includes losses for the smallest and largest subnetworks, a rule called sandwich training.The smallest and largest candidates are jointly optimized with uniformly sampled architectures using SGD; the formulation can transfer knowledge from larger to smaller networks.Source: §2.1 · p. 3
3. NAS via Attentive Sampling
AttentiveNAS targets Pareto-optimal networks by changing which subnetworks receive training attention. It samples candidates at matched FLOPs and trains the best or worst performer, using efficient performance estimates to guide selection.
3. NAS via Attentive Sampling
- The paper motivates improving worst-Pareto models as a way to better optimize the weight-sharing graph and potentially obtain more informative gradients and architecture-space exploration.It relates worst-Pareto sampling to hard-example mining by treating those subnetworks as difficult examples.Source: §3 · p. 4
3.1. Sub-networks of Interest
- A best Pareto architecture has no lower-FLOPs or equal-FLOPs alternative with better validation performance.The paper also defines a worst Pareto architecture as one dominated in accuracy by architectures with the same or larger FLOPs.Source: §3.1 · p. 4
3.2. Pareto-attentive pre-training
- The Pareto-attentive objective shifts training focus from equally weighted candidates to best or worst Pareto architectures, while retaining a FLOPs-based sampling distribution.Setting the FLOPs and conditional architecture distributions to their priors recovers the uniform-sampling objective.Source: §3.2 · p. 4
- For each target FLOPs value, the algorithm samples k subnetworks and selects the highest- or lowest-performing candidate according to estimator P.The objective approximates the FLOPs expectation with n Monte Carlo samples; the selection rule uses an indicator for the best or worst estimated performance among candidates.Source: §3.2 · p. 4
- BestUp trains the best-performing sampled candidate, whereas WorstUp trains the worst-performing one; both are AttentiveNAS sampling strategies.The algorithm samples target FLOPs from the search-space prior and draws candidates under that FLOPs constraint.Source: §3.2 · pp. 4–5
- Performance estimation can use a pretrained accuracy predictor or mini-batch loss; validation loss is an ideal estimator but can be expensive on a large validation set.The paper presents these efficient surrogates as alternatives for comparing candidate performance during training.Source: §3.2 · pp. 4–5
4. Experimental Results
AttentiveNAS evaluates a constrained ImageNet search space, efficient performance estimators, and sampling strategies that focus training on Pareto-relevant subnetworks. Its selected strategy outperforms prior NAS baselines on accuracy–FLOPs trade-offs.
4.1. Search Space
- The search space follows FBNetV3’s meta-architecture with reduced ranges and spans 203 to 1,939 MFLOPs.The authors report that its deeper, narrower subnetworks achieve higher accuracy than those from the BigNAS search space under similar FLOPs constraints.Source: §4.1 · p. 5
4.2. Training and Evaluation
- FLOPs-constrained sampling approximates the conditional architecture distribution empirically, then uses rejection sampling from that approximation and asynchronous CPU sampling.The approximation is intended to improve efficiency over rejection sampling from the entire search space; FLOPs are discretized in 25-MFLOP steps.Source: §4.2 · p. 5
- The evaluation uses evolutionary search with 5,248 architectures and calibrates batch-normalization statistics before evaluation.The search fixes an initial population of 512, mutation and crossover populations of 128, and 20 iterations.Source: §4.2 · p. 6
4.3. Attentive Sampling with Efficient Performance Estimation
- The accuracy predictor reaches Kendall’s τ of 0.89 on held-out subnetworks, 0.87 across training stages, and 0.88 across random data partitions.The predictor is trained after 30 epochs; rank correlation is evaluated against actual accuracy measured using weight-sharing parameters.Source: §4.3 · pp. 6–7
4.4. NAS with Efficient Attentive Sampling
- Around 200 ± 10 MFLOPs, WorstUp-50 (acc) exceeds Uniform by about 0.3% top-1 accuracy, and WorstUp-1M (acc) also improves on Uniform.WorstUp-50 and WorstUp-1M also raise the minimum and first-quartile accuracy over BestUp strategies and Uniform; BestUp-50 (acc) begins to outperform them above 400 MFLOPs.Source: §4.4 · p. 8
- BestUp-3 (loss) performs best overall, with results suggesting that effective sampling balances exploration and exploitation.Both loss-based strategies improve on Uniform; BestUp-3 (loss) explores better than the accuracy-based BestUp variants and exploits better than Uniform.Source: §4.4 · p. 8
4.5. Comparison with Prior NAS Approaches
- AttentiveNAS significantly outperforms all baselines on ImageNet accuracy–FLOPs trade-offs.The comparison uses the discovered best Pareto set, evaluated on the original ImageNet validation set after search on 200,000 training examples.Source: §4.5 · p. 8, Table 1 · p. 8
5. Conclusion
AttentiveNAS uses attentive sampling for two-stage NAS and reports improved ImageNet accuracy over prior-art NAS approaches under varied FLOPs constraints.
- AttentiveNAS outperforms prior-art NAS approaches on ImageNet, establishing new SOTA accuracy under various FLOPs constraints.Source: §5 · p. 8
- AttentiveNAS improves accuracy over uniform sampling by taking the performance Pareto into account.Source: §5 · p. 8
- The paper proposes a variety of attentive sampling strategies for training two-stage NAS.Source: §5 · p. 8
A. Training settings
Training uses sandwich sampling, in-place distillation, and fixed-size image patches rescaled to target resolutions, alongside SGD and the listed regularization settings.
- Each iteration evaluates four subnetworks: the smallest and largest plus two sampled subnetworks, with smaller ones supervised by the largest.The sandwich sampling rule always trains the smallest and biggest subnetworks, and n = 2 makes four total evaluations.Source: A. Training settings · p. 11
- For different input resolutions, fixed-size patches such as 224x224 on ImageNet are rescaled to each target resolution with bicubic interpolation.Source: A. Training settings · p. 11
- Training uses SGD with cosine learning-rate decay, 64 GPUs, batch size 32 per GPU, and base learning rate 0.1 scaled per 256 samples.Source: A. Training settings · p. 11
- Training uses AutoAugment, label smoothing 0.1, momentum 0.9, weight decay 10^-5, dropout 0.2, and stochastic layer dropout 0.2.Unless specified, training lasts 360 epochs; only the largest model uses weight decay and dropout, and synchronized batch normalization is not used.Source: A. Training settings · p. 11
B. Robustness of two-stage NAS
Stage 1 training shows reproducible rank correlations across data partitions and epochs, while the proposed FLOPs-conditional sampling reduces the reported sampling-trial count substantially.
B. Robustness of two-stage NAS
- The robustness evaluation trains on a random 90% partition of ImageNet and tests sampled subnetworks on the remaining 10%.After training, 1,024 subnetworks are sampled and evaluated on their corresponding test partition.Source: B. Robustness of two-stage NAS · p. 11
- Rank correlations are 0.94 between runs at epoch 30, 0.96 between trials at epoch 360, and 0.88 between epoch-30 and final performance.The study assesses robustness across data partitions, initializations, and training epochs; Figure 7 labels partitions by seeds 0 and 1 and epochs 30 and 360.Source: B. Robustness of two-stage NAS · p. 11, Figure 7 · p. 11
C. Sampling efficiency
- About 12 trials are needed on average to sample a targeted architecture from the proposal distribution, versus 50,878 with naive uniform sampling under a random FLOPs constraint.Source: C. Sampling efficiency · p. 11
D. Comparisons of search space
The search space is adapted from FBNetV3 and differs from BigNAS; the provided passages describe its components and note a constraint on one comparison.
- The search space is adapted from FBNetV3 and contains deeper, narrower subnetworks than the BigNAS search space.Source: D. Comparisons of search space · p. 11
- The uniform sampling strategy is compared across both search spaces, while the search space proposed in [3] is excluded because its training requires progressive shrinking and stage-specific tuned hyperparameters.Source: D. Comparisons of search space · p. 12
- Table 2 describes MBConv and MBPool blocks, SE layers, width, depth, kernel size, expansion ratio, and swish activation.MBConv is identified as an inverted residual block, MBPool as an efficient last stage, and SE as squeeze and excite.Source: Table 2 · p. 12
E. Comparisons of training and search time
AttentiveNAS reports the lowest training cost among the compared methods, despite back-propagating through four subnetworks per iteration. It also produces Pareto-optimal networks for multiple deployment scenarios in one training run, without fine-tuning or retraining.
E. Comparisons of training and search time
- AttentiveNAS builds on weight-sharing, alleviating the burden of training thousands of subnetworks from scratch.Source: E. Comparisons of training and search… · p. 12
- One training run finds Pareto-optimal networks for different FLOPs requirements, avoiding repeated NAS runs for each deployment scenario.The networks can be sampled from the weight-sharing graph without fine-tuning or retraining.Source: E. Comparisons of training and search… · p. 13
- AttentiveNAS has the lowest reported training cost after 360 epochs, despite back-propagating through four subnetworks per iteration.This makes per-batch time roughly 4× slower; the additional evolutionary search is less than 10% of total training time.Source: E. Comparisons of training and search… · p. 13, Table 3 · p. 13
F. Additional results on inference latency
AttentiveNAS’s attentive sampling can also be applied to inference latency. The reported comparisons show better latency–ImageNet accuracy trade-offs than EfficientNet models.
F. Additional results on inference latency
- AttentiveNAS models yield better latency versus ImageNet validation accuracy trade-offs than EfficientNet models.Source: F. Additional results on inference… · p. 13, Table 4 · p. 13
- The paper uses GPU latency as an example of adapting attentive sampling to metrics beyond FLOPs.Source: F. Additional results on inference… · p. 13
- The latency comparisons cover both 2080 Ti and V100 GPUs.Source: F. Additional results on inference… · p. 13, Table 4 · p. 13
G. Transfer learning results
AttentiveNAS-A1 and AttentiveNAS-A4 are evaluated on three transfer-learning benchmarks. The paper reports that its models achieve the best transfer-learning accuracy under the stated fine-tuning setup.
G. Transfer learning results
- The transfer-learning evaluation tests AttentiveNAS-A1 and AttentiveNAS-A4 on Oxford Flowers, Stanford Cars, and Food-101.Source: G. Transfer learning results · p. 13
- AttentiveNAS models yield the best transfer-learning accuracy on the reported benchmarks.The results cover Oxford Flowers, Stanford Cars, and Food-101.Source: G. Transfer learning results · p. 13, Table 5 · p. 13
- Fine-tuning runs for 150 epochs with batch size 64, while learning rate and weight decay are selected using a 20% hold-out subset.The setup also uses SGD with momentum 0.9, label smoothing 0.1, and dropout 0.5.Source: G. Transfer learning results · p. 13