Neural Architecture Search

The search space defines the universe of possible neural network architectures the NAS algorithm can explore. For embedded systems, this space is heavily constrained to ensure the final model is deployable on target hardware.

• Cell-based spaces define repeating, parameterized building blocks (cells) that are stacked to form the full network.
• Macro spaces allow the search over the entire network's macro-architecture, including the number of layers, layer types (e.g., depthwise convolution, inverted residual block), and connection patterns.
• Hardware-aware constraints are baked in, limiting operations to those supported by the microcontroller (e.g., excluding unsupported activations) and setting bounds on total parameters or FLOPs.

The search strategy is the optimization algorithm that navigates the search space to find high-performing architectures. Different strategies offer trade-offs between search efficiency and computational cost.

• Reinforcement Learning (RL): Uses a controller (e.g., an RNN) to generate architectures, which are then trained and evaluated; the resulting accuracy serves as a reward to update the controller (e.g., NASNet, MnasNet).
• Evolutionary Algorithms: Treat architectures as individuals in a population, using mutation and crossover operations to evolve better designs over generations (e.g., AmoebaNet).
• Differentiable NAS (DARTS): Formulates the search as a continuous optimization problem where architecture choices are relaxed into soft, learnable parameters, enabling efficient gradient-based search.
• One-Shot / Weight-Sharing: Trains a single, over-parameterized supernet that encompasses all sub-architectures; the search involves evaluating sub-networks by sampling paths from this shared weight pool, drastically reducing search cost (e.g., ProxylessNAS, Once-For-All).

Evaluating every candidate architecture by fully training it is prohibitively expensive. Performance estimation strategies provide low-cost proxies for final model accuracy.

• Lower Fidelity Estimates: Train candidates for fewer epochs, on a subset of data, or with reduced resolution.
• Weight Sharing: As used in one-shot NAS, sub-networks inherit weights from the supernet, and their performance is estimated via a single forward pass or very few epochs of fine-tuning.
• Learning Curve Extrapolation: Predicts final accuracy from the initial phases of training.
• Surrogate Models: Train a separate predictor model (e.g., a regressor) to map an architecture's encoding to its expected performance.
• For hardware metrics like latency, a look-up table or a latency predictor model is often pre-built by profiling each operation type on the target device (e.g., microcontroller or NPU).

A defining feature of NAS for embedded systems is the direct optimization for hardware-specific constraints, moving beyond just accuracy. The search objective becomes a multi-criteria optimization problem.

• Primary Constraints: Latency (inference time), memory footprint (SRAM/Flash usage), and energy consumption.
• Objective Functions: Often combine accuracy with a penalty for violating constraints (e.g., Accuracy * (Latency/Target)^-w).
• On-Device Measurement: The most reliable method integrates direct hardware profiling into the search loop, using tools like MCU-based profilers or simulator estimates for cycle-accurate performance.
• Platform-Specific Operations: The search space is limited to kernels well-supported by the target hardware, such as depthwise separable convolutions for efficient ARM CPUs or binary operations for ultra-low-power microcontrollers.

Differentiable Architecture Search (DARTS) is a prominent gradient-based NAS method that significantly accelerates the search process. It transforms the discrete search space into a continuous one for efficient optimization.

• Continuous Relaxation: Each choice (e.g., which operation connects two nodes in a cell) is represented by a set of continuous architecture parameters (alpha). All candidate operations are executed and mixed via a softmax-weighted sum.
• Bi-Level Optimization: The process alternates between:
  1. Updating network weights (w) by standard gradient descent on the training set.
  2. Updating architecture parameters (alpha) by gradient descent on a validation set.
• Discretization: After search, a final architecture is derived by replacing the mixed operations with the strongest one (e.g., the operation with the highest alpha value).
• Advantage: DARTS can complete a search in a few GPU days, compared to thousands for earlier RL methods.

One-Shot NAS methods dramatically reduce computational cost by training a single, weight-shared supernetwork (or one-shot model) that encompasses all architectures in the search space.

• Supernet Training: The supernet is trained once, with each training step sampling a different sub-architecture (path) and updating only the weights along that path.
• Search as Evaluation: After training, the search involves evaluating the performance of many sub-networks by inheriting weights from the supernet, requiring only inference or minimal fine-tuning.
• Efficiency: This decouples the expensive weight training from the architecture search, enabling NAS on limited compute budgets.
• Challenges: Weight coupling and co-adaptation can make the performance ranking of sub-networks less reliable. Methods like ProxylessNAS address this by directly training binarized paths on the target hardware and task.
• Once-For-All (OFA): Extends this concept by training a supernet to support a large range of sub-networks (different depths, widths, kernel sizes), enabling instant specialization to diverse hardware constraints without retraining.

General-purpose AutoML libraries that include NAS as a core capability, providing high-level APIs to automate model design. These frameworks abstract the search algorithm, allowing developers to specify constraints like model size or latency.

• Google's AutoKeras: An open-source library built on Keras/TensorFlow, using Bayesian optimization and network morphism for accessible NAS.
• Microsoft's NNI (Neural Network Intelligence): A comprehensive toolkit supporting multiple NAS strategies (e.g., ENAS, DARTS) and hyperparameter tuning across frameworks.
• AutoGluon: An AutoML library from Amazon that provides easy-to-use tabular, text, and image predictors, automating architecture and hyperparameter choices.

These tools lower the barrier to entry but may require customization for strict microcontroller constraints.

Specialized platforms that directly integrate hardware feedback into the search loop. They profile candidate architectures on the target device (or an accurate simulator) to optimize for metrics like latency, energy, and memory footprint—critical for TinyML.

• ProxylessNAS: Pioneered direct search on the target hardware (e.g., mobile phone) using gradient-based methods and a latency lookup table, eliminating proxy inaccuracies.
• Once-For-All (OFA): Trains a single supernet that subsumes many sub-architectures. Search involves selecting a sub-network tailored to specific device constraints without retraining, enabling efficient deployment across heterogeneous edge hardware.
• MCUNet's TinyNAS: Co-designs the neural architecture with the inference engine (TinyEngine). It searches for networks that fit within the SRAM and flash memory limits of microcontrollers (e.g., < 256KB) while maximizing accuracy.

Frameworks based on the Differentiable Architecture Search (DARTS) paradigm, which treats the search space as a continuous, over-parameterized network. Architecture parameters are learned via gradient descent alongside weights, making search orders of magnitude faster than reinforcement learning or evolutionary methods.

• DARTS: The foundational implementation, formulating the choice between operations (e.g., 3x3 conv, skip connect) as a continuous relaxation.
• GDAS (Gradient-based Search using Differentiable Architecture Sampler): Improves stability and reduces memory cost by sampling a single path in the supernet during search.
• SNAS (Stochastic NAS): Introduces a concrete distribution and policy gradient to maintain the differentiability advantage while yielding a discrete architecture in a single search run.

These are powerful for research but can be memory-intensive during search.

Platforms employing population-based evolutionary algorithms or controller-based reinforcement learning (RL) to explore the architecture space. These methods are highly flexible and can discover novel, non-differentiable structures.

• Google's AmoebaNet: Used regularized evolution (a genetic algorithm) to discover architectures that outperformed human-designed models on ImageNet.
• ENAS (Efficient NAS): Uses a controller RNN trained with policy gradient to generate model architectures, with weight sharing across child models to drastically reduce search cost.
• LEAF (Lifelong Evolutionary AF): An evolutionary framework focused on continual adaptation of architectures for changing tasks or data streams on edge devices.

These approaches often require substantial computational resources for the initial search phase.

Emerging frameworks specifically designed for the extreme constraints of microcontrollers. They incorporate microarchitecture-level details like SRAM/Flash budgets, fixed-point quantization effects, and available DSP instructions.

• MCUNet (TinyNAS + TinyEngine): The leading co-design framework for microcontrollers. TinyNAS performs hardware-aware search under the memory bottleneck, while TinyEngine is a memory-optimized inference library that executes the found model.
• µNAS (Micro-NAS): A hardware-aware NAS framework that uses a predictor-based approach, training a cheap latency/accuracy predictor for target MCU platforms to accelerate the search.
• SPOS (Single Path One-Shot NAS): A supernet-based method where the search is simplified to training a one-shot model and then evaluating sub-networks via uniform path sampling, adaptable to TinyML constraints.

These tools are essential for pushing the boundaries of on-device AI capability.

Hardware-Aware Neural Architecture Search (HW-NAS) is a variant of NAS where the search algorithm directly optimizes the neural architecture for specific hardware performance metrics like latency, energy consumption, or memory usage on target devices such as microcontrollers or NPUs. Unlike standard NAS which focuses solely on accuracy, HW-NAS incorporates a hardware feedback loop.

• The search process is guided by a cost model or direct on-device measurements (e.g., latency lookup tables, energy profilers).
• It enables the discovery of models that are Pareto-optimal for the trade-off between accuracy and hardware efficiency.
• This is critical for TinyML, where a model's viability is determined by its ability to fit within a device's SRAM and flash memory constraints and meet real-time latency requirements.

Differentiable Architecture Search (DARTS) is a gradient-based NAS method that formulates the search space as a continuous, over-parameterized network where architecture parameters are learned via standard gradient descent alongside model weights. It replaces the discrete search over operations with a continuous relaxation.

• All candidate operations (e.g., 3x3 conv, 5x5 conv, skip connect) are represented in a mixed operation, weighted by architecture parameters (alpha).
• The search involves jointly optimizing the network weights (with standard gradient descent) and the architecture parameters (with gradient descent on the validation loss).
• After search, a discrete architecture is derived by retaining the operation with the highest weight for each node. This method is significantly faster than reinforcement learning or evolutionary-based searches but can be memory-intensive during the search phase.

A Once-For-All (OFA) network is a single, large neural network trained to support many sub-networks of different depths, widths, and kernel sizes through weight sharing. It decouples training from search to enable efficient hardware-aware NAS.

• The supernet is trained with progressive shrinking, where larger sub-networks are trained first, and then smaller ones are fine-tuned from the larger network's weights.
• After training, a specialized sub-network can be extracted for any target constraint (e.g., 50ms latency, 200KB memory) without retraining.
• This allows for rapid architecture specialization across a diverse fleet of edge devices with varying capabilities, as the search involves only evaluating candidate sub-networks from the pre-trained supernet.

ProxylessNAS is a neural architecture search method that directly learns architectures on the target task and hardware without using a proxy (like training on a smaller dataset or for fewer epochs). It brings the search cost closer to that of training a single network.

• It employs path-level binarization, where at any time only one path in the over-parameterized network is activated, reducing memory overhead during search.
• Gradients for the architecture parameters are estimated using binary connect, enabling gradient-based optimization.
• It can directly incorporate hardware latency, measured via a lookup table or on-device, into the loss function, making it a practical method for HW-NAS. This avoids the inaccuracy of proxy metrics when searching for microcontroller-optimal models.

MCUNet is a co-design framework that jointly optimizes Tiny Neural Network architecture (TinyNAS) and the underlying inference engine (TinyEngine) to enable ImageNet-scale classification directly on microcontrollers with severe memory constraints (e.g., 256KB SRAM).

• TinyNAS is a hardware-aware NAS that searches for models under the peak memory usage constraint of the target device.
• TinyEngine is a memory-aware inference library that performs in-place depthwise convolution and other optimizations to minimize memory footprint during execution.
• The co-design ensures the searched architecture is not only efficient in theory but can be executed efficiently by the inference engine on real hardware. This approach has demonstrated state-of-the-art accuracy on microcontrollers like the Arm Cortex-M7.

Evolutionary Neural Architecture Search uses evolutionary algorithms, inspired by biological evolution, to discover optimal network architectures. It maintains a population of candidate models which are iteratively improved through selection, mutation, and crossover operations.

• Each model (genotype) is evaluated for fitness (e.g., accuracy on a validation set, latency).
• The best-performing models are selected to produce offspring for the next generation through mutation (e.g., changing an operation type, adding/removing a layer) and crossover (combining parts of two parent architectures).
• While often computationally expensive, evolutionary methods are highly flexible and can navigate complex, non-differentiable search spaces. They are effective for multi-objective optimization, such as jointly optimizing for accuracy, latency, and model size.