Quantization-Aware NAS (QA-NAS)

The defining mechanism of QA-NAS is the integration of simulated quantization into the architecture evaluation loop. During the search, candidate models are trained or evaluated with fake quantization nodes that mimic the effects of converting weights and activations to low-precision integers (e.g., INT8). This exposes architectures to quantization noise early, allowing the search algorithm to select for quantization robustness.

• Key Technique: Uses Quantization-Aware Training (QAT) principles within the NAS supernet.
• Benefit: Discovers architectures that maintain high accuracy post-deployment after true integer quantization, unlike standard NAS models which can suffer significant accuracy drops.

QA-NAS treats bit-width as a searchable parameter alongside traditional architectural choices like operations and connections. The search space is expanded to include mixed-precision configurations, allowing the algorithm to discover optimal precision assignments per layer or per tensor.

• Example: A QA-NAS search might assign 8-bit precision to most layers but allocate 4-bit to non-critical layers to save memory, or use higher precision (16-bit) for sensitive layers to preserve accuracy.
• Outcome: Achieves a superior accuracy-efficiency Pareto frontier compared to applying uniform post-training quantization to a model found by standard NAS.

QA-NAS employs a hardware cost model that accurately reflects the benefits of quantization. Lower precision directly translates to reduced model size, memory bandwidth, and energy consumption on target hardware. The search objective penalizes architectures based on their quantized cost, not their floating-point cost.

• Critical Metric: Estimates inference latency and peak memory usage for the quantized version of the candidate model.
• Result: Discovers models that are not only accurate when quantized but are also Pareto-optimal for the quantized hardware deployment scenario (e.g., < 100KB model size, < 10ms latency on an Arm Cortex-M7).

In one-shot or differentiable QA-NAS methods, the supernet is trained with a quantization-aware objective. The shared weights are optimized to perform well across many sub-architectures under simulated quantization conditions. This ensures the performance rankings of sub-networks are predictive of their final quantized accuracy.

• Core Challenge: Mitigating the bias introduced by weight sharing in a quantized setting.
• Solution: Techniques like in-place distillation or sandwich rule sampling are often used to stabilize training and improve the correlation between supernet rankings and final stand-alone quantized performance.

The search algorithm inherently learns to favor architectural patterns that exhibit low quantization sensitivity. This often leads to the discovery of networks with specific traits:

• Avoidance of extreme activation ranges that cause clipping loss.
• Preference for operations with lower quantization error, such as depthwise separable convolutions over standard convolutions in certain contexts.
• Automatic insertion of quantization helpers, like clipping or requantization layers, at optimal points in the graph.

These patterns are emergent properties of the search under quantization noise, not manually designed.

QA-NAS is a cornerstone methodology for TinyML and microcontroller deployment. It directly addresses the core constraint of these environments: extremely limited memory. By co-designing the architecture for low-precision execution, QA-NAS produces models that are deployment-ready for kilobyte-scale memory budgets.

• Primary Output: A family of Pareto-optimal models at different quantized size-accuracy trade-offs.
• Downstream Benefit: Eliminates the trial-and-error process of quantizing a standard NAS model, providing deterministic guarantees that the discovered architecture will meet strict size and latency constraints after final conversion to TensorFlow Lite Micro or similar runtimes.

This method adapts the Differentiable Architecture Search (DARTS) framework to be quantization-aware. A continuous relaxation of the search space allows architecture parameters and network weights to be jointly optimized via gradient descent.

• Quantization Simulation: During search, operations like convolutions use fake quantization nodes that simulate the effects of low-bit integer arithmetic (e.g., INT8) in the forward pass, while maintaining full precision gradients for the backward pass.
• Objective: The search optimizes for both task accuracy (e.g., ImageNet classification) and a quantization robustness loss, often measured as the accuracy drop between the full-precision and quantized versions of the candidate sub-network.
• Outcome: Produces a model whose architecture choices (e.g., preferring ReLU6 over ReLU, selecting more quantization-friendly operations) are inherently stable when later converted to fixed-point for deployment.

This approach builds a weight-sharing supernet that is trained with quantization simulation from the start. The entire supernet undergoes Quantization-Aware Training (QAT), making all its shared weights robust to low-precision representation.

• Process: First, a large supernet encompassing all candidate operations is constructed. This supernet is trained on the target task using standard QAT techniques, where weights and activations are quantized and de-quantized during forward passes.
• Architecture Selection: After training, various sub-architectures (sub-networks) are sampled from the supernet. Their performance is evaluated by inheriting the robust, quantized weights and measuring accuracy on a validation set—without any fine-tuning.
• Advantage: Drastically reduces search cost, as only one QAT process for the supernet is required. Frameworks like Once-For-All (OFA) can be extended with QAT to enable this.

This method uses a Reinforcement Learning (RL) agent as the search strategy, where the reward function explicitly penalizes architectures that degrade significantly after quantization.

• Controller: An RNN controller proposes a candidate architecture description (e.g., a string of tokens defining layers and operations).
• Reward Signal: The candidate is trained briefly, then quantized (either via Post-Training Quantization or a short QAT cycle). The reward is a composite function: R = Accuracy(quantized_model) - λ * [Accuracy(fp_model) - Accuracy(quantized_model)]. This directly rewards high post-quantization accuracy and penalizes the quantization gap.
• Hardware Cost Integration: The reward can be further combined with hardware metrics (e.g., latency of the quantized model on a target MCU) to perform joint QA-NAS and hardware-aware search. Early works like MNasNet can be extended with this quantization-aware reward.

A foundational implementation detail is the design of the NAS search space itself to favor operations and structures known to be quantization-friendly.

• Operation Choices: The search space is limited to or biased towards layers with low quantization error. This includes:
  • Depthwise Convolutions: Lower parameter count and activation range reduces quantization sensitivity.
  • ReLU6 Activation: Bounds activations to [0,6], providing a stable, known range for quantization scaling factors, unlike unbounded ReLU.
  • Skip Connections with Add (not Concatenate): Add operations are more efficient and stable for integer arithmetic than concatenation, which increases channel dimensions.
• Explicit Bit-Width Search: Some advanced QA-NAS formulations expand the search space to include the bit-width for weights or activations per layer as a searchable parameter, seeking a mixed-precision model.

This is QA-NAS applied under the extreme constraints of microcontrollers (MCUs), where the search is directly constrained by the memory footprint of the quantized model.

• Constraint Modeling: A hardware cost model predicts the RAM, flash, and latency of a candidate architecture after it is quantized to a target precision (e.g., INT8). This quantized cost is used during the search.
• Multi-Objective Optimization: The search algorithm (e.g., an evolutionary algorithm) optimizes for Pareto-optimal solutions balancing quantized accuracy and quantized model size (in kilobytes).
• Implementation Loop: Candidate architectures are compiled for a target MCU (e.g., ARM Cortex-M) using a framework like TensorFlow Lite for Microcontrollers to obtain accurate memory profiles, closing the loop between the NAS and the final deployment pipeline.

An evolution of ProxylessNAS, this method performs direct QA-NAS on the target task and dataset without proxy stages, and includes quantization simulation throughout.

• Direct Hardware Measurement: Instead of relying on latency lookup tables, the architecture is directly deployed and profiled on the target hardware (e.g., a mobile phone or MCU dev board) in its quantized form during the search.
• Binarized Path Sampling: The method uses a binary gate to select one path at a time in the over-parameterized network. These paths are trained and evaluated with fake quantization enabled.
• Gradient-Based Architecture Adjustment: The binary gates are updated via gradient descent using the straight-through estimator, where the gradient incorporates the performance feedback from the quantized forward pass. This results in an architecture optimized end-to-end for quantized performance on the actual device.

The foundational category for QA-NAS. Hardware-Aware NAS incorporates platform-specific constraints—such as latency, memory footprint, and energy consumption—directly into the search objective. Unlike standard NAS, which optimizes only for accuracy, it treats hardware metrics as first-class optimization targets, often using a hardware cost model to predict performance without on-device measurement for every candidate.

The core compression technique that QA-NAS must accommodate. Quantization reduces the numerical precision of a model's weights and activations, typically from 32-bit floating-point to 8-bit integers or lower.

• Post-Training Quantization (PTQ): Applied after a model is fully trained. QA-NAS can search for architectures inherently robust to this process.
• Quantization-Aware Training (QAT): Simulates quantization during training to learn compensated weights. QA-NAS often integrates QAT loops into the supernet training phase.

The computational backbone of efficient QA-NAS. A supernet is an over-parameterized network that embeds all possible architectures defined by the search space. Weight sharing allows all candidate sub-networks (subnets) to share parameters from this single supernet.

• This enables the evaluation of thousands of architectures without individually training each one, a necessity for assessing quantization robustness across many designs.
• QA-NAS trains the supernet with simulated quantization to ensure shared weights are already adapted to low-precision arithmetic.

A predictive function that estimates the on-device performance of a candidate architecture. In QA-NAS, this model must account for the effects of quantization.

• It can be a lookup table of measured kernel latencies or a small neural network regressor.
• Predicts metrics like inference latency, peak memory usage, and energy for a given (architecture, quantization-scheme, hardware-target) triplet.
• Using a cost model avoids the prohibitive overhead of deploying and profiling every candidate on real hardware during search.

A prominent gradient-based NAS methodology often adapted for QA-NAS. DARTS relaxes the discrete search space (e.g., choice of operation) into a continuous one using architecture parameters.

• These parameters are optimized via gradient descent alongside network weights.
• For QA-NAS, the relaxation can be extended to include quantization policy choices (e.g., bit-width per layer), allowing the joint search of architecture and quantization strategy through gradients.

The mathematical framework governing QA-NAS. The search must balance competing objectives: model accuracy, quantized model size, inference latency, and sometimes energy use.

• The goal is to find architectures on the Pareto frontier, where improving one metric degrades another.
• Search strategies use techniques like weighted sum loss or evolutionary algorithms to navigate this trade-off space and deliver a set of optimal models for different deployment constraints.