Reinforcement Learning NAS (RL-NAS)

The Controller Network is a recurrent neural network (often an LSTM or Transformer) that acts as the agent in the RL framework. Its function is to sequentially generate a string of tokens that defines a candidate neural network architecture.

• Action Space: Each token corresponds to a specific architectural choice, such as layer type, kernel size, or number of filters.
• Policy: The controller's parameters define a stochastic policy for sampling architectures.
• Training: It is trained via policy gradient methods (e.g., REINFORCE) to maximize the expected reward from the generated architectures.

The Search Space is the finite set of all possible neural network architectures the controller can generate. It is a critical constraint that defines the scope of discovery.

• Cell-Based Search: A common approach where the controller designs a repeating computational cell (normal and reduction cells), which are then stacked to form the full network.
• Hierarchical Search: The controller makes high-level macro-architecture decisions (e.g., number of blocks) and low-level operator choices within each block.
• Micro-Architecture Parameters: The space includes choices for operators (convolution, pooling, identity), kernel sizes, channel widths, and skip connection patterns.

The Reward Function R(A) is the signal that guides the controller's learning. It quantitatively evaluates a candidate architecture A, typically based on its performance after training.

• Primary Objective: Often the validation accuracy of the child network after a short training cycle (a proxy task).
• Hardware-Aware Objectives: For RL-NAS targeting microcontrollers, the reward is a composite function, e.g., R(A) = Accuracy(A) - λ * Latency(A), where λ is a Lagrange multiplier balancing accuracy and a hardware cost.
• Sparse and Noisy: The reward is computationally expensive to obtain and can be noisy due to the stochastic nature of neural network training.

The Child Network (or candidate architecture) is the neural network instantiated from the controller's tokens. Its performance is used to compute the reward.

• Training Protocol: Due to cost, child networks are trained for a reduced number of epochs on a subset of data. This performance estimation strategy is a major bottleneck.
• Weight Sharing: To accelerate evaluation, most modern RL-NAS methods use a supernet. All child networks share weights from this single, over-parameterized network, allowing instant evaluation without full training.
• Validation: The child network's accuracy on a held-out validation set is the primary performance metric fed into the reward function.

A Hardware Cost Model is a predictive function, essential for Hardware-Aware RL-NAS, that estimates the on-device resource consumption of a candidate architecture.

• Purpose: Provides the latency, memory footprint, or energy consumption metrics for the reward function without needing hardware-in-the-loop profiling for every candidate.
• Implementation: Can be a pre-characterized lookup table of operator costs, a small neural network regressor, or an analytical model based on FLOPs and memory access counts.
• Microcontroller Focus: For TinyML, models predict metrics like SRAM usage, flash storage, and inference latency in milliseconds, ensuring designs fit within kilobyte-scale budgets.

Policy Gradient Optimization is the reinforcement learning algorithm used to update the controller network's parameters (θ) based on the rewards from sampled architectures.

• REINFORCE Algorithm: The foundational method. The gradient is estimated as ∇θ J(θ) ≈ Σ ∇θ log P(A_t; θ) * R_t, where A_t is a sampled architecture and R_t its reward.
• Baselines: A moving average baseline of rewards is subtracted from R_t to reduce variance and stabilize training.
• Goal: Adjust θ to increase the probability (P(A; θ)) of generating high-reward architectures and decrease the probability of low-reward ones over many iterations.

The core adaptation of RL-NAS for TinyML is the design of a multi-objective reward function. Instead of rewarding only validation accuracy, the controller is trained to maximize a composite score that heavily penalizes violations of hardware limits.

• Key Reward Components: Reward = Accuracy - λ₁ * (Latency - Target)² - λ₂ * (Peak RAM - Budget)
• Constraint Types: Includes hard constraints (e.g., model must fit in 512KB Flash) and soft constraints (e.g., lower latency is better).
• Search Goal: Discovers architectures that are not just accurate, but are feasible for deployment on devices like the Arm Cortex-M series or Espressif ESP32.

The RL controller's action space and state representation are modified to be hardware-aware, enabling it to learn the cost of architectural choices.

• Action Space Includes: Layer type (e.g., standard conv vs. depthwise separable conv), kernel size, expansion ratios, and skip connections—all with known hardware footprints.
• State Embedding: The controller's state includes running totals of estimated FLOPs, parameter count, and memory access cost (MAC) for the partial architecture.
• Outcome: The agent learns to avoid operations that are prohibitively expensive on microcontrollers, such as large dense layers or standard convolutions with high channel counts.

To avoid the intractable cost of deploying every candidate on real hardware during search, RL-NAS systems use pre-characterized hardware cost models.

• Look-Up Tables (LUTs): Pre-measured latency and energy costs for each operator (e.g., 3x3 Conv, ReLU) on the target MCU are stored in a LUT. The total cost of a candidate network is estimated by summing the costs of its constituent operations.
• Neural Predictors: Small MLPs trained to predict latency/memory from a graph representation of the architecture.
• ProxylessNAS Inspiration: Methods like ProxylessNAS demonstrate direct hardware measurement, where a few top candidates are profiled on-device to refine the cost model, creating a hardware-in-the-loop search loop.

RL-NAS can be used to discover architectures that maintain high accuracy after post-training quantization (PTQ) or are amenable to quantization-aware training (QAT).

• Reward Signal: The accuracy used in the reward is the quantized accuracy—the performance of the model after simulating 8-bit or even 4-bit integer quantization during search.
• Architectural Bias: The search learns to favor operations with low quantization error, such as ReLU activations over more complex ones like Swish, and architectures with redundant channels that can be pruned after quantization.
• Synergy with QA-NAS: This approach aligns closely with Quantization-Aware NAS (QA-NAS), creating models that are born ready for efficient integer inference on MCUs.

A single RL-NAS search can be directed to find a Pareto frontier of optimal models, providing multiple deployment options for a range of devices.

• Multi-Objective Optimization: The controller searches for architectures that represent the best trade-offs between accuracy, latency, and model size.
• Result: A suite of models, from a larger, more accurate version for a high-end MCU (e.g., Cortex-M7) to a tiny, less accurate version for a severely constrained device (e.g., Cortex-M0).
• Efficiency: This is more efficient than running separate searches for each constraint value and is analogous to the Once-For-All (OFA) training paradigm but discovered via RL.

Neural Architecture Search (NAS) is the overarching automated machine learning technique that algorithmically discovers optimal neural network architectures for a given task and dataset, replacing manual, expert-driven design. It consists of three core components: a search space defining possible architectures, a search strategy to explore it, and a performance estimator to evaluate candidates. RL-NAS is one specific implementation of a search strategy.

Hardware-Aware Neural Architecture Search is a specialized variant of NAS that incorporates hardware-specific constraints—such as inference latency, peak memory usage, and energy consumption—directly into the search objective. Unlike standard NAS which often optimizes only for accuracy, this approach discovers models optimized for a target deployment platform like a mobile phone or microcontroller, ensuring they are both accurate and efficient.

In Neural Architecture Search, the search space defines the universe of all possible neural network architectures the algorithm can explore. It is parameterized by variables such as:

• Layer types (e.g., convolution, pooling, attention)
• Connectivity patterns (e.g., skip connections, branching)
• Operator hyperparameters (e.g., kernel size, number of channels) A well-designed search space balances expressiveness with tractability, constraining the search to plausible architectures and preventing combinatorial explosion.

The search strategy is the core algorithm that navigates the NAS search space to find high-performing architectures. RL-NAS is one such strategy. Other prominent strategies include:

• Evolutionary Algorithms: Use mutation and crossover on a population of architectures.
• Gradient-Based Methods (e.g., DARTS): Relax the search space to be continuous and use gradient descent.
• Bayesian Optimization: Model the performance landscape to guide the search. The choice of strategy directly impacts search efficiency, cost, and the quality of the final architecture.

A performance estimator is a method for predicting the final accuracy or hardware efficiency of a candidate neural architecture without the prohibitive cost of full training and evaluation. Techniques include:

• Weight Sharing: Training a single supernet to approximate the performance of all sub-architectures.
• Low-Fidelity Estimators: Training on a smaller dataset or for fewer epochs.
• Zero-Cost Proxies: Using metrics derived from the network's initial state (e.g., gradient norms). These estimators are critical for making NAS computationally feasible.

Microcontroller NAS (MCU-NAS) refers to hardware-aware neural architecture search methods specifically tailored for the extreme constraints of microcontrollers (MCUs). It optimizes for metrics critical to embedded deployment:

• Kilobyte-scale memory footprints (RAM and Flash)
• Microjoule-level energy consumption per inference
• Real-time latency on single-core, low-clock-speed CPUs This represents the most constrained end of the hardware-aware NAS spectrum, directly relevant to TinyML deployment.