Differentiable NAS

The core innovation of Differentiable NAS is the continuous relaxation of the discrete choice between candidate operations (e.g., 3x3 convolution, 5x5 convolution, skip connection). Each edge in the computational super-network is represented as a weighted sum of all possible operations. The mixing weights, known as architecture parameters (alpha), are continuous variables learned via gradient descent alongside the model weights. This transforms architecture search from a combinatorial problem into a differentiable one.

Differentiable NAS is framed as a bi-level optimization problem. Two sets of parameters are optimized in an alternating fashion:

• Model Weights (w): Optimized on the training data to minimize loss, given the current architecture.
• Architecture Parameters (α): Optimized on a held-out validation split to find the architecture that generalizes best. This separation prevents the architecture search from overfitting to the training data. The process uses standard gradient-based optimizers, making it significantly faster than reinforcement learning or evolutionary-based NAS methods.

After the joint optimization of weights and architecture parameters concludes, a discrete architecture must be derived from the continuous super-network. This is typically done by retaining the operation with the highest learned architecture parameter (alpha) value on each edge and pruning away the others. For cell-based search spaces, the top-k strongest operations per node are selected to form the final computational cell, which is then stacked to form the complete network.

Differentiable NAS is highly relevant for TinyML and embedded neural network architectures due to its efficiency and direct constraint optimization:

• Search Speed: Completes in GPU-days versus the GPU-years of early NAS methods, making iterative design feasible.
• Hardware-Aware Search: Can directly incorporate latency or memory footprint as differentiable or proxied terms in the loss function, enabling hardware-aware NAS (HW-NAS). This is critical for finding networks that fit within the severe SRAM/Flash limits of microcontrollers.
• Discovery of Novel Blocks: Has led to the discovery of efficient building blocks like DARTS cells, which can be more parameter-efficient than hand-designed mobile blocks for certain tasks.

Despite its advantages, Differentiable NAS presents several challenges:

• Memory Overhead: The over-parameterized super-network requires significant GPU memory, which can limit the search space size.
• Performance Collapse: The continuous relaxation can lead to architecture parameter convergence issues, where skip connections or zero operations are overly favored, degrading final performance. Techniques like partial channel connections and architecture regularization are used to mitigate this.
• Proxy Gap: Searches are often conducted on smaller datasets (CIFAR-10) or fewer epochs as a proxy for large-scale tasks (ImageNet), which may not perfectly translate.
• Discretization Gap: The performance of the derived discrete architecture can be lower than that of the continuous super-network during search.

Differentiable NAS sits within a broader ecosystem of automated architecture discovery:

• Vs. Reinforcement Learning (RL) NAS (e.g., NASNet): DARTS uses gradients instead of policy gradients, offering orders-of-magnitude speedup.
• Vs. Evolutionary NAS: DARTS is more sample-efficient, not relying on population-based mutation and selection.
• Vs. One-Shot / Weight-Sharing NAS: DARTS is a specific, gradient-based instance of weight-sharing NAS, where the shared super-network's weights are used to estimate architecture gradients.
• ProxylessNAS: An evolution that performs direct search on the target task and hardware without proxies, using path-level binarization but still leveraging gradient signals, bridging DARTS and direct hardware measurement.

DARTS is used to discover ultra-compact neural networks for microcontrollers (MCUs) with severe memory (<512KB SRAM) and power constraints. The search directly optimizes for metrics like peak memory usage and latency on target hardware (e.g., ARM Cortex-M series). This results in custom architectures like MCUNet, which co-designs the network (TinyNAS) and inference engine (TinyEngine) to run ImageNet-scale models on devices with 256KB of SRAM.

• Key Constraint: Search loss includes hardware feedback (e.g., measured latency).
• Outcome: Networks under 100KB that avoid operations incompatible with MCU accelerators.

Differentiable NAS tailors convolutional networks for fixed-function hardware accelerators like Google's Edge TPU or Intel's Movidius VPU. The search space is constrained to operators supported by the target's instruction set (e.g., avoiding SE blocks with non-linearities for Edge TPU). This hardware-aware optimization leads to families like EfficientNet-Lite, which provide high accuracy while ensuring full compatibility and peak throughput on dedicated AI chips.

• Search Target: Maximize FLOPS/utilization on the target accelerator.
• Benefit: Eliminates manual redesign, delivering latency-optimized models ready for production compilation.

The differentiable supernetwork from DARTS serves as the foundation for Once-For-All (OFA) training. A single, large over-parameterized network is trained to contain many high-performing sub-networks of varying depths, widths, and kernel sizes. This allows instant specialization for diverse deployment scenarios—from a powerful gateway to a tiny sensor—without retraining, enabling efficient hardware-aware NAS across an entire product line.

• Mechanism: Architecture parameters define a continuous distribution of sub-networks.
• Use Case: A single OFA model deployment can be dynamically shrunk or expanded based on real-time device resources.

DARTS automates the design of novel blocks for non-vision tasks prevalent in IoT. This includes searching for optimal 1D convolutional or recurrent structures for time-series data from accelerometers, gyroscopes, or audio sensors. The search can incorporate temporal resolutions and filter lengths as continuous variables, producing models that outperform hand-designed counterparts for applications like keyword spotting, predictive maintenance, and anomaly detection.

• Example: Searching for an optimal temporal convolution stack for a 20-class keyword spotting task.
• Advantage: Discovers efficient, task-specific feature extractors that minimize compute per inference.

Differentiable NAS can be extended to jointly search the network topology and its quantization strategy. The search space includes the bit-width (e.g., 8-bit, 4-bit, binary) for each layer's weights and activations as continuous, optimizable parameters. This results in heterogeneously quantized networks where sensitivity varies per layer, maximizing accuracy under a strict model size budget (e.g., < 50KB) for deployment on microcontrollers.

• Outcome: A single search produces both the architecture and the integer-only inference graph.
• Benefit: Achieves a superior size-accuracy trade-off compared to post-training quantization of a fixed architecture.

In real-time embedded systems (e.g., drones, robotics), meeting a strict inference latency deadline is critical. Differentiable NAS incorporates a latency loss term into the optimization, using either a pre-built hardware latency lookup table or on-device measurements. The search penalizes architectures that exceed the latency target, directly discovering networks that guarantee performance on the specific hardware, be it a microcontroller, DSP, or low-power GPU.

• Constraint: Search optimizes: Loss = Task_Loss + λ * (Latency - Target)^2.
• Result: Architectures that are provably fast enough for closed-loop control at 30 FPS or higher.

Neural Architecture Search (NAS) is the overarching automated process for discovering optimal neural network topologies for a given task and set of constraints. It treats architecture design as a search problem within a pre-defined space of possible operations (e.g., convolution types, kernel sizes).

• Core Methods: Includes reinforcement learning, evolutionary algorithms, and gradient-based approaches like Differentiable NAS.
• Objective: Automatically balances model accuracy against constraints like latency, model size, and energy consumption.
• Significance: Eliminates manual, trial-and-error network design, which is critical for finding efficient architectures for embedded systems.

Hardware-Aware Neural Architecture Search (HW-NAS) is a specialized variant of NAS where the search algorithm directly optimizes for performance metrics on specific target hardware.

• Direct Optimization: The search cost function incorporates measured or predicted inference latency, memory usage, and power consumption on the target device (e.g., a specific microcontroller or NPU).
• Beyond FLOPs: Recognizes that computational FLOPs are a poor proxy for real-world hardware efficiency due to factors like memory bandwidth and cache effects.
• Application: Essential for TinyML, where the search must find architectures that fit within severe SRAM/Flash limits (e.g., < 256KB) of microcontrollers.

A Once-For-All (OFA) network is a weight-sharing super-network trained to contain a vast number of sub-networks of varying depths, widths, and kernel sizes.

• Efficient Search: Enables hardware-aware NAS without retraining. After the single OFA network is trained, optimal sub-networks for different hardware constraints can be extracted via fast search.
• Unified Training: Mitigates the prohibitive cost of training each candidate architecture from scratch.
• TinyML Impact: Pioneered by the MCUNet framework, where OFA principles are used to co-design networks and inference engines for microcontrollers, enabling ImageNet-scale inference on devices with under 1MB of memory.

ProxylessNAS is a neural architecture search method that conducts the search directly on the target task and hardware without using a proxy (like a smaller dataset or fewer training epochs).

• Direct Path Learning: Employs path-level binarization, where only one path in the over-parameterized network is activated at a time, making memory usage manageable during search.
• Gradient-Based: Like DARTS, it uses gradient descent to learn architecture parameters but applies it directly to the full dataset and incorporates a hardware latency loss term.
• Advantage: Eliminates the accuracy gap that can occur when a proxy task is used, ensuring the final architecture is optimized for both accuracy and true on-device latency.

MCUNet is a groundbreaking TinyML co-design framework that jointly optimizes the neural network architecture (via TinyNAS) and the underlying inference engine (TinyEngine) for microcontroller units (MCUs).

• Two-Sided Optimization: TinyNAS performs hardware-aware search under the precise memory and latency profile dictated by the TinyEngine scheduler.
• Breakthrough Result: Enabled ImageNet-scale image classification on commercially available microcontrollers with as little as 256KB of SRAM and 1MB of Flash.
• Relation to DARTS: Demonstrates the practical application of differentiable and hardware-aware NAS principles in the most constrained embedded environments.

TinyNAS is the neural architecture search component within the MCUNet framework, specifically designed for microcontrollers.

• Memory-Centric Search: The primary constraint is the peak SRAM usage during inference, not just FLOPs or parameter count. The search space is designed to avoid operations that cause large activation maps.
• Efficiency Focus: Generates families of extremely small models (e.g., Micro-DNNs) under 100KB that are Pareto-optimal for the accuracy-memory trade-off.
• Evolution from DARTS: Applies the principles of differentiable search but within a drastically constrained design space tailored for MCU memory hierarchies and integer-only arithmetic.