Hardware-Aware NAS

Hardware-Aware NAS treats hardware metrics as primary optimization objectives alongside accuracy. The search algorithm navigates a Pareto frontier, balancing trade-offs between competing goals.

• Common Objectives: Latency (inference time), energy consumption (milliwatt-hours), peak memory usage, model size (parameters/FLOPs).
• Search Formulation: Often framed as a constrained optimization problem (e.g., maximize accuracy subject to latency < 10ms) or a multi-objective loss function (e.g., Loss = CrossEntropy + λ * Latency).
• Example: ProxylessNAS directly incorporates the target device's measured latency into the training loss, enabling the discovery of architectures like MobileNetV3 that are Pareto-optimal for mobile CPUs.

A core challenge is efficiently and accurately predicting hardware metrics for candidate architectures without full deployment. Methods vary in fidelity and cost.

• Look-Up Tables (LUTs): Pre-measure latency/energy for each primitive operation (convolution, pooling) on target hardware. Total cost is estimated by summing costs of all ops in the candidate graph. Fast but can miss system-level interactions.
• Analytical Cost Models: Use mathematical formulas based on hardware specs (memory bandwidth, peak FLOPS) and operator characteristics to estimate execution time. More generalizable but requires deep hardware understanding.
• Neural Predictors: Train a small regression model (e.g., a multi-layer perceptron) on architecture encodings to predict latency or power. Achieves good accuracy after an initial profiling phase and can capture complex, non-linear interactions.

The set of possible architectures (the search space) is explicitly constrained by hardware considerations from the outset, making the search more efficient and relevant.

• Hardware-Primitive Alignment: The space includes operations known to be efficient on the target hardware. For an NPU, this might favor depthwise separable convolutions over standard convolutions due to lower memory bandwidth requirements.
• Constraint Embedding: The space definition itself can enforce limits, such as a maximum number of channels per layer to stay within on-chip SRAM capacity or restricting to power-of-two tensor dimensions for optimal tiling.
• Example: FBNet's search space was designed around efficient building blocks for mobile GPUs, excluding operations with poor kernel support or high latency.

Modern Hardware-Aware NAS often uses differentiable search methods, which are more sample-efficient than reinforcement learning or evolutionary algorithms.

• Supernet & Architecture Parameters: A weight-sharing supernet encompasses all candidate operations. Architecture parameters (α) control the contribution of each operation. Hardware cost is made differentiable with respect to α.
• Gradient-Based Co-Optimization: The search performs gradient descent on both network weights (w) and architecture parameters (α), simultaneously learning which operations are accurate and hardware-efficient.
• Practical Benefit: This allows the search to evaluate millions of architectures by training only one supernet, with hardware cost directly guiding the gradient updates.

The approach must adapt to vastly different hardware profiles, from data center GPUs to microcontrollers. The optimal architecture is highly platform-dependent.

• Compute-Bound vs. Memory-Bound Hardware: On a high-FLOPS GPU, a compute-heavy architecture may be optimal. On a memory-bandwidth-limited mobile NPU, an architecture that minimizes data movement is better.
• Specialized Hardware Units: Searches can target specific hardware features, such as tensor cores (optimizing for 4x4 or 8x8 matrix multiplies) or dedicated activation hardware (favoring ReLU over SiLU).
• Cross-Platform Deployment: Techniques like Once-For-All (OFA) train a single supernet from which efficient subnetworks can be extracted for multiple target devices (e.g., iPhone, Pixel, Raspberry Pi) without retraining.

Hardware-Aware NAS does not operate in isolation; it relies on and feeds into the broader ML compilation and deployment stack.

• Compiler-Aware Search: The search uses the same graph compiler (e.g., TVM, XLA) that will be used for deployment to ensure cost estimations reflect final, optimized kernels, including benefits from operator fusion.
• Quantization-Aware Search: The search can be conducted in a quantization-aware setting, where candidates are evaluated at INT8 precision to find architectures robust to precision loss.
• Deployment Pipeline: The output is an architecture description (e.g., a ONNX model or a network definition) that flows directly into standard model compilation, quantization, and runtime pipelines for the target hardware.

This is the most prevalent application, where the search directly targets latency and power consumption on specific mobile System-on-Chips (SoCs) like the Apple Neural Engine or Qualcomm Hexagon NPU. The search space is constrained to operations highly optimized for these accelerators (e.g., depthwise-separable convolutions).

• Key Metric: Inference latency (ms) measured directly on the target device or an accurate cycle-accurate simulator.
• Example: Searching for an image classification model that achieves <15ms latency on a Snapdragon 8 Gen 3's NPU while maintaining >75% ImageNet accuracy.

Here, NAS balances high accuracy for safety-critical tasks (object detection, segmentation) with the strict thermal and power budgets of automotive-grade hardware (e.g., NVIDIA DRIVE Orin, Tesla FSD Chip). The search must produce architectures that meet real-time frame rate requirements (e.g., 30 FPS) under worst-case thermal conditions.

• Key Constraint: Power consumption (watts) and memory bandwidth utilization, as sensor fusion pipelines run multiple models concurrently.
• Outcome: Architectures that use efficient attention variants or hybrid CNN-transformer blocks optimized for the vehicle's specific compute fabric.

Targets microcontrollers (MCUs) and ultra-low-power AI accelerators (e.g., Arm Ethos-U55). The search objective minimizes energy per inference (microjoules) to enable year-long battery life. The search space is severely constrained by tiny SRAM memory (e.g., 256KB), forcing novel layer designs that avoid costly off-chip memory accesses.

• Key Metric: Peak memory footprint and energy consumption, often evaluated using hardware-in-the-loop (HIL) profiling.
• Example: Discovering a keyword spotting model for a hearing aid that uses <100KB of RAM and consumes <10 µJ per inference.

Focuses on reducing the total cost of ownership (TCO) for large-scale cloud inference. The search optimizes for throughput (inferences/second) and memory bandwidth utilization on specific server-grade accelerators like NVIDIA T4/TensorRT, Google TPU v4, or AWS Inferentia. The goal is to maximize queries per second per dollar.

• Key Objective: Throughput under a target latency Service Level Agreement (SLA), often using batch size as a key search variable.
• Outcome: Architectures that maximize hardware utilization, often favoring regular, batch-friendly operations over irregular, memory-intensive ones.

Applies NAS to design surrogate models (emulators) for physics simulations that must run efficiently on specialized supercomputing hardware (e.g., AMD MI300X, Intel Gaudi). The search incorporates hardware-specific floating-point throughput (TFLOPS) and inter-GPU communication costs for model-parallel execution.

• Key Consideration: Support for mixed precision (FP16, BF16, TF32) and the ability to fuse custom, domain-specific operators into efficient kernels.
• Example: Searching for a neural PDE solver architecture optimized for the tensor core layout and high-bandwidth memory of an HPC cluster's primary accelerator.

Addresses the 'hardware heterogeneity' problem in federated learning, where a global model must perform well across a fleet of diverse devices (old phones, new tablets, IoT sensors). Hardware-aware NAS can search for a Pareto-optimal set of sub-networks (e.g., via a Once-For-All supernet) from which an optimal architecture can be extracted for each device class without retraining.

• Key Challenge: Creating a search space and training regimen that yields high-performing sub-networks across vastly different latency and compute profiles.
• Benefit: Enables efficient personalized deployment without maintaining hundreds of separate model variants.