TuNAS: Tunable NAS for Latency-Aware CNNs

• TuNAS is a tunable NAS algorithm that uses an RL-based controller and weight-sharing techniques to efficiently navigate enormous CNN search spaces with latency constraints.
• It employs a one-shot supernetwork with aggressive parameter sharing methods such as operation collapsing and channel masking to reduce computational redundancy.
• Empirical evaluations on ImageNet and MS COCO show TuNAS achieves 1–2.5% higher accuracy with lower compute cost than random search in large-scale settings.

TuNAS (“Tunable NAS”) is a scalable, weight-sharing neural architecture search (NAS) algorithm designed to identify high-quality, resource-aware convolutional neural network (CNN) architectures within extremely large search spaces while adhering to strict latency constraints. TuNAS employs a reinforcement learning (RL)–based controller paired with a supernetwork, leveraging aggressive parameter sharing and latency-aware reward mechanisms to efficiently explore vast architectural regimes. Comprehensive experimental analyses on ImageNet and COCO benchmarks demonstrate that TuNAS yields architectures outperforming those found by random search, particularly as the search space size escalates (Bender et al., 2020).

1. Algorithmic Foundations

TuNAS maintains a supernetwork (one-shot model) parameterized by shared weights $W$ and an RL-based controller $\pi$ that defines a multinomial policy over architectural choices. Each step in the search alternates between two core operations:

• Policy Update: The controller $\pi$ is updated using a policy-gradient approach (REINFORCE) based on a scalar reward function evaluated over validation performance and latency compliance.
• Weight Sharing: The shared weights $W$ are optimized (SGD) using the loss from a single subnetwork, sampled according to $\pi$.

The algorithm can be summarized as follows:

1. Initialize shared weights $W$ and controller $\pi$ (uniform distribution).
2. Repeat for $N_\text{steps}$:
   • Sample architecture $\alpha \sim \pi$.
   • Compute training loss $L_\text{train}(W; \alpha)$.
   • Evaluate validation quality $\pi$0.
   • Calculate latency $\pi$1 and reward $\pi$2, for $\pi$3.
   • Update $\pi$4 via REINFORCE: $\pi$5, where $\pi$6 is a moving baseline.
   • Update $\pi$7 by gradient descent on $\pi$8.
3. Select the final architecture $\pi$9 by maximizing $\pi$0 over each decision.

This framework transforms the constrained optimization (maximizing $\pi$1 subject to $\pi$2) into an unconstrained RL objective using a reward function that enables direct and stable targeting of the latency constraint.

2. Search Space Construction and Architectural Choices

TuNAS operates over three primary mobile-CNN search spaces, each defined as sequences of inverted-bottleneck blocks in the MobileNetV2/V3 style. Each block allows categorical variation along these axes:

Search Space	Expansion Ratios	Depthwise Kernel Sizes	Output Channels	Squeeze-and-Excite	Activation
ProxylessNAS	3, 6	3, 5, 7	base MobileNetV2/V3	No	ReLU
ProxylessNAS-Enlarged	3, 6	3, 5, 7	Multiplier ×$\pi$3	No	ReLU
MobileNetV3-Like	1–6	3, 5, 7	Same as above	Yes/No	ReLU/Swish

No novel operators beyond the MobileNetV2/V3 family are introduced. Rather, TuNAS's contribution is the ability to explore these vast combinations with significant computational efficiency. Cardinalities of the evaluated search spaces are approximately $\pi$4, $\pi$5, and $\pi$6 respectively.

3. Weight-Sharing Strategies

TuNAS introduces aggressive parameter sharing to make supernetwork training tractable:

• Operation Collapsing: The same 1×1 expand convolutional weights are used across differing depthwise kernel sizes, reducing redundancy.
• Channel Masking: Instead of allocating unique weight tensors for each output channel configuration, convolutional layers are implemented up to a maximum width $\pi$7. Narrower configurations are realized by masking off excess channels, making the supernetwork's parameter count linear rather than exponential in categorical options.

Sampling proceeds by executing and updating weights along a single path $\pi$8 per search iteration. Only the active path's weights receive gradient signal, promoting greater parameter reuse and diminishing interference between rare operations.

4. Hyperparameter Specification and Quality Improvement Techniques

Key hyperparameter values are as follows:

• Shared-weight training: RMSProp (momentum=0.9, decay=0.9, $\pi$9), max learning rate 2.64 (warm-up 2.5%), cosine decay, batch size 4096 (utilizing 128×32 TPU cores), and L2 regularization $W$0.
• Controller: Adam (learning rate $W$1 base, $W$2, $W$3, $W$4). Controller learning rate zero for the first 25% of steps (RL disabled for warmup), then either constant or exponentially increased to $W$5.
• Search duration: $W$6–$W$7 hours on TPU-v3-32, corresponding to $W$8 epochs.

To safeguard training signal propagation to rare operations and filters, warmup schedules are implemented:

• Filter Warmup: During the first 25%, all $W$9 filters are stochastically enabled (linear anneal $\pi$0), ensuring rare channel slices receive updates.
• Operation Warmup: Similarly, all operations in each block are enabled with linearly annealed probability, averaging their outputs during initial training steps.
• Absolute-Value Reward: Direct reward $\pi$1 eliminates the need for extensive tuning of cost exponents $\pi$2, with the same value performing robustly across spaces and hardware.

These innovations reduce manual hyperparameter tuning and stabilize optimization, particularly in early training epochs.

5. Empirical Evaluation and Comparative Analysis

TuNAS was benchmarked on ImageNet-1K classification and MS COCO object detection. Architectures are searched using 90 epochs of weight-sharing training; final selected architectures are retrained from scratch for up to 360 epochs.

ImageNet-1K Classification

Search Space	Random Search (N=50)	TuNAS (90 epochs)	TuNAS (360 epochs)	Reference Model (training setup)
ProxylessNAS	$\pi$3	$\pi$4	–	$\pi$5
ProxylessNAS-Enl.	$\pi$6	$\pi$7	–	$\pi$8
MobileNetV3-Like	$\pi$9	$W$0	$W$1	$W$2 (at 58 ms)

TuNAS yields a $W$3–$W$4 percentage point absolute improvement in top-1 accuracy compared to random search, with gains magnifying for larger search spaces (up to $W$5 at $W$6 models). Notably, random search incurs $W$7–$W$8 greater compute cost per result, as each trial involves full stand-alone training.

MS COCO Detection

Backbone	mAP (%)	Latency (ms)
MobileNetV2+SSDLite	20.7	126
MnasNet+SSDLite	21.3	129
ProxylessNAS+SSDLite	21.8	140
MobileNetV3+SSDLite	22.0	106
TuNAS+SSDLite	22.5	106

TuNAS+SSDLite achieves the highest mAP (22.5) at a 106 ms latency target.

6. Key Insights and Practical Recommendations

TuNAS provides evidence that weight-sharing NAS—with the proposed architectural parameter sharing, RL-based controller, absolute-value latency reward, and warmup schedules—is especially advantageous in large-scale search spaces where computational resources are constrained.

• On search spaces up to $W$9 models, TuNAS finds CNNs that are $\pi$0–$\pi$1 percentage points absolutely more accurate than those found by random search under identical hardware and latency constraints.
• Hyperparameter robustness is improved: the same settings (notably $\pi$2) are directly transferable between search spaces, mitigating the need for costly per-task grid search.
• For practical use: set the target latency $\pi$3, reuse the published hyperparameters for $\pi$4 and $\pi$5, and execute five independent search runs to estimate variance, selecting the final architecture by argmax-probability under $\pi$6 and retraining from scratch.

A plausible implication is that as search space complexity and resource constraints become more salient, the advantages of weight-sharing approaches—when properly tuned—are magnified, reinforcing the rationale for their deployment in industrial NAS applications (Bender et al., 2020).