Gradient-Based NAS

The fundamental innovation of Gradient-Based NAS is the continuous relaxation of the discrete architecture choice. Instead of selecting a single operation (e.g., convolution or pooling) for a layer, the method assigns a mixture or weighted sum of all possible operations. This creates a differentiable supernet where architecture parameters (often denoted as alpha, α) become continuous variables that can be optimized via gradient descent. The final discrete architecture is derived by selecting the operation with the highest learned weight.

Training involves solving a bi-level optimization problem with two sets of parameters:

• Model weights (w): Standard network parameters (e.g., convolution filters).
• Architecture parameters (α): Parameters defining the structure of the supernet.

The objective is to find α that minimizes the validation loss, where the model weights w are obtained by minimizing the training loss. The standard update procedure alternates between:

1. Updating w on the training set with α fixed.
2. Updating α on the validation set with w fixed, using gradients approximated via the chain rule. This separation prevents overfitting the architecture to the training data.

To make gradient-based search computationally feasible, a single over-parameterized supernet is trained. This supernet encompasses all possible operations and pathways defined by the search space. All candidate sub-architectures share the weights of their constituent operations from this supernet. This allows for the approximate evaluation of millions of architectures without training each from scratch, reducing search cost from thousands of GPU days to orders of magnitude less.

Unlike Reinforcement Learning (RL) or Evolutionary-based NAS, the search strategy is gradient-based optimization. The architecture parameters α are updated using gradients ∇αL_val, where L_val is the validation loss. This is typically done with standard optimizers like Adam or SGD. The direct use of gradients provides a more efficient and stable search path through the architecture space compared to discrete, sample-based methods.

After the joint optimization of α and w, a discretization step is required to obtain a deployable model. Common methods include:

• Argmax: For each layer, select the operation with the highest learned α value.
• Retraining: The selected discrete architecture is usually retrained from scratch (without weight sharing) on the full training dataset to obtain final, standalone weights. This step is crucial as the supernet weights are co-adapted and may not represent the optimal weights for the final, pruned architecture.

Gradient-Based NAS offers distinct benefits:

• Speed: Dramatically faster search (e.g., 1-4 GPU days) compared to RL-NAS (thousands of GPU days).
• Deterministic Search: Gradient descent provides a more direct, reproducible path to a solution versus stochastic RL controllers or evolutionary algorithms.
• Integration with Hardware-Aware Objectives: The framework can be extended to optimize for latency or energy by making hardware cost estimators differentiable, enabling direct gradient-based optimization of non-differentiable metrics via techniques like Gumbel-Softmax or REINFORCE.

The search space defines the universe of all possible neural network architectures an algorithm can explore. It is parameterized by variables such as:

• Layer types (e.g., convolution, depthwise convolution, pooling)
• Connectivity patterns (e.g., skip connections, branching)
• Channel counts and kernel sizes A well-designed search space balances expressiveness (ability to find high-performing models) with tractability (ease of search). For TinyML, the search space is tightly constrained by hardware limits, often excluding operations that are computationally prohibitive on microcontrollers.

A supernet is an over-parameterized neural network that embeds all possible operations and pathways defined by the search space. Weight sharing is the technique of training this single supernet, allowing the parameters of its many sub-networks (child architectures) to be co-trained. This is foundational to efficient Gradient-Based NAS (like DARTS) and One-Shot NAS methods. Instead of training thousands of models from scratch, candidate architectures are evaluated by activating a specific path through the shared weights of the supernet, reducing search cost from thousands of GPU days to a few.

A hardware cost model is a predictive function that estimates the real-world efficiency metrics of a candidate neural network architecture on target hardware. Instead of deploying every candidate, the NAS algorithm queries this model. Common implementations include:

• Look-up tables (LUTs) pre-measuring the latency/energy of each operation.
• Small neural network regressors trained to predict latency from architecture descriptors. For microcontroller NAS, the cost model must accurately predict peak memory usage, inference latency, and energy consumption, which are non-linear functions of the architecture due to memory bottlenecks and fixed-point arithmetic.

Multi-Objective Neural Architecture Search optimizes for several competing goals simultaneously, such as maximizing accuracy while minimizing latency and model size. The output is not a single "best" model but a set of architectures on the Pareto frontier. A model is Pareto-optimal if no other architecture exists that is better in all objectives. For TinyML, a typical Pareto frontier trades test accuracy against RAM footprint and inference time. Engineers select the final model from this frontier based on their specific hardware budget.

DARTS is the seminal Gradient-Based NAS method. It relaxes the discrete search space into a continuous one by introducing architecture parameters (alpha) that represent the strength of connections between operations. The supernet is trained via bilevel optimization: the model weights (w) are optimized on the training data, and the architecture parameters (alpha) are optimized on a validation set, both using gradient descent. After search, a discrete architecture is derived by retaining the operations with the highest alpha values. DARTS made NAS significantly more efficient but can be memory-intensive due to the supernet.

Once-For-All is a one-shot NAS and training paradigm that decouples training from search. A single progressively shrinking supernet is trained to support many sub-networks of varying depth, width, and kernel size. After training, a specialized search algorithm (e.g., evolutionary search) can extract thousands of hardware-tailored sub-networks in seconds without any retraining. OFA is highly practical for TinyML deployment, as one supernet can service a diverse fleet of microcontrollers with different resource constraints, simplifying the deployment pipeline.