Neural Architecture Search (NAS) for Compression

The search space defines the universe of possible architectures the NAS algorithm can explore. For compression, this space is constrained to include operations and connectivity patterns known to be efficient.

• Key components: Choice of building blocks (e.g., inverted residual blocks, depthwise separable convolutions), kernel sizes, expansion ratios, and skip connections.
• Design principle: The space must be large enough to contain high-performing solutions but constrained enough for the search to be computationally feasible. Overly large spaces make search intractable.
• Example: A mobile-focused search space might exclude standard convolutions in favor of depthwise separable convolutions and limit the maximum number of channels per layer to control the parameter count directly.

The search strategy is the algorithm used to navigate the defined search space to find optimal architectures. Efficiency is critical, as evaluating each candidate by full training is prohibitively expensive.

• Reinforcement Learning (RL): Uses a controller network (an RNN or transformer) to generate architecture descriptions, which are then trained and evaluated. The controller's policy is updated based on the reward (e.g., high accuracy, low latency).
• Evolutionary Algorithms: Maintains a population of candidate architectures. Through cycles of selection, mutation (e.g., changing an operation type), and crossover, the population evolves toward better performers.
• Differentiable NAS (DNAS): Formulates the search as a supernet containing all possible operations. Architecture parameters are relaxed into continuous variables, allowing the use of gradient descent to jointly optimize network weights and architecture choices. This is often the most sample-efficient strategy.

Accurately predicting the performance (accuracy, latency, size) of a candidate architecture without fully training it is the core challenge that makes NAS feasible.

• Weight Sharing / Supernet: Multiple candidate architectures share weights from a single, over-parameterized supernet. A candidate's performance is estimated by activating its specific path in the supernet, avoiding the cost of training from scratch.
• Proxy Tasks: Candidates are evaluated on a smaller dataset (e.g., CIFAR-10), trained for fewer epochs, or with a reduced model width. While not perfectly correlated with final performance, it provides a fast, rank-order estimate.
• Hardware-in-the-Loop Profiling: For latency or energy targets, the candidate's architecture code is compiled and profiled on the target device (or an accurate simulator) to get real hardware metrics. This is essential for Hardware-Aware NAS (HW-NAS).

The objective function quantifies what an 'optimal' architecture means, guiding the search strategy. For compression, it is always a multi-objective optimization balancing accuracy with efficiency metrics.

• Standard Form: Objective = Accuracy(L) - λ * [Cost(L)], where Cost(L) can be FLOPs, parameter count, measured latency, or estimated energy. The Lagrange multiplier λ controls the trade-off.
• Hard Constraints: The search can be formulated to find the most accurate architecture under a strict hardware constraint (e.g., latency < 10ms, model size < 5MB).
• Pareto-Optimal Search: The goal is to discover a set of architectures that form the compression-accuracy Pareto frontier, where no architecture is better in both metrics. This provides multiple optimal options for different deployment scenarios.

A critical specialization where the search directly optimizes for metrics specific to the target deployment hardware, moving beyond abstract proxies like FLOPs.

• Why FLOPs are insufficient: FLOPs count ignores critical hardware factors like memory bandwidth, cache effects, and operator-level kernel efficiency. A model with lower FLOPs can be slower on real hardware.
• Integrated Cost Models: The search uses a pre-built cost model (e.g., a neural network) that predicts latency/energy for a given architecture on target hardware (e.g., a specific mobile SoC or Neural Processing Unit).
• Direct Measurement: The most accurate method involves building each candidate kernel and profiling it on the actual device or a cycle-accurate simulator during the search, though this is computationally intensive.

The dominant paradigm for efficient NAS, where the entire search space is encoded into a single, trainable supernet (or one-shot model).

• Mechanism: All possible operations (e.g., 3x3 conv, 5x5 conv, identity) exist in parallel at each layer choice. Architecture parameters (alphas) act as attention weights over these operations. Training involves two stages:
  1. Supernet Training: Jointly train the weights of all operations and the architecture parameters.
  2. Architecture Selection: Freeze the alphas, select the strongest operation at each choice (e.g., via argmax), and extract the final, discrete sub-network.
• Advantage: It amortizes the cost of training across all candidates, making search orders of magnitude faster than training each candidate independently.
• Challenge: Requires careful design to avoid architecture parameter bias and ensure the supernet ranking correlates with stand-alone performance.

This is the primary application of NAS for compression. The search algorithm is directly constrained by or optimized for target hardware metrics, discovering architectures that are efficient on specific silicon.

• Key Metrics: The search objective incorporates real hardware measurements or accurate proxies for latency, energy consumption, memory footprint, and FLOPs.
• Example: Searching for a vision model that must achieve <20ms inference time on a specific mobile Neural Processing Unit (NPU). The algorithm evaluates candidate architectures using a latency lookup table or an on-device profiler.
• Outcome: Produces a model that is co-designed with its deployment environment, avoiding the inefficiency of compressing a generic model post-hoc.

NAS is used to push the boundaries of on-device AI, creating models that can run on microcontrollers (MCUs) with severe constraints (e.g., <512KB of RAM).

• Focus: Architectures with extremely low parameter counts and activations that fit in tiny SRAM. This often involves novel operations and macro-architecture choices.
• Techniques: Search spaces may include depthwise separable convolutions, squeeze-and-excitation blocks, and efficient activation functions like ReLU6.
• Real-World Use: Keyword spotting on smart home devices, anomaly detection in industrial sensors, and wake-word detection on wearables where cloud connectivity is unavailable or prohibitively expensive.

Differentiable Neural Architecture Search formulates the search as a continuous optimization problem, making it efficient to discover optimal, compact building blocks (cells) for larger networks.

• Mechanism: A supernet contains all possible candidate operations (e.g., 3x3 conv, 5x5 depthwise conv, skip connection). Architecture parameters are learned via gradient descent, and the final model is derived by selecting the strongest paths.
• Output for Compression: Discovers novel, hardware-friendly layers that outperform hand-designed counterparts like MobileNetV2's inverted residual blocks.
• Advantage: More computationally efficient than reinforcement learning-based NAS, enabling the search for compression-optimal cells within practical resource budgets.

Advanced NAS frameworks don't just search for operations and connections; they simultaneously search for the optimal quantization policy per layer or block.

• Search Space Expansion: The algorithm chooses both the operation type and its numerical precision (e.g., 8-bit, 4-bit, binary).
• Objective: Maximize accuracy under a combined constraint of model size (bits) and compute (integer ops).
• Result: A model that is born quantized, with different layers potentially using different precisions (mixed-precision) in an optimal configuration determined by the search, not heuristic rules.

Instead of pruning a pre-defined architecture, this approach searches for architectures that are inherently sparse-friendly or that already incorporate an optimal sparse pattern.

• Connection to Lottery Ticket Hypothesis: The search aims to find the high-performing sparse subnetwork from the start, rather than discovering it through iterative pruning of a dense network.
• Method: The search space may include sparsity masks or allow the algorithm to learn which connections are critical during the architecture exploration phase.
• Benefit: Produces models with a native sparse structure that can leverage sparse tensor cores or dedicated inference engines for maximal speedup.

This paradigm trains a single, large supernet that contains many sub-networks of varying sizes and computational costs. It decouples training from search, enabling instant specialization.

• Process: The supernet is trained with techniques like weight sharing and progressive shrinking. After training, an evolutionary search or predictor quickly finds the best sub-network for any given device constraint (latency, model size).
• Use Case for Compression: A developer can query this system for an optimal model that meets the exact memory budget of a new IoT chip, receiving a ready-to-deploy, high-accuracy sub-model in seconds without retraining.
• Efficiency: Eliminates the need to run a full NAS cycle for every new deployment target.

Hardware-Aware NAS is a search methodology that directly incorporates target hardware performance metrics—such as latency, energy consumption, or memory footprint—into the objective function when evaluating candidate architectures. Unlike standard NAS, which often optimizes only for accuracy or parameter count, HW-NAS searches for networks that are Pareto-optimal for both accuracy and on-device efficiency.

• Key Objective: Minimize a joint loss function: Loss = Task_Loss + λ * Hardware_Cost.
• Search Space: Includes operations and layer configurations known to be efficient on the target silicon (e.g., depthwise convolutions for mobile NPUs).
• Example: Searching for a model that achieves <10ms latency on a specific smartphone SoC while maintaining >95% ImageNet top-5 accuracy.

Differentiable NAS is a gradient-based search method that formulates the selection of operations and connections within a super-network as a continuous optimization problem. Instead of using reinforcement learning or evolutionary algorithms, DNAS employs a relaxed architecture representation where architectural choices are modeled by continuous parameters (e.g., architecture weights). The search is performed via standard gradient descent, making it significantly faster than discrete search methods.

• Core Mechanism: Uses a softmax over candidate operations per layer. The probability of selecting an operation is learned via Gumbel-Softmax or similar techniques.
• Application to Compression: DNAS is highly effective for channel pruning and mixed-precision quantization search, as it can smoothly explore the space of sparsity patterns and bit-widths.
• Outcome: Produces a single, optimized sub-network from the super-network after the search concludes.

Automated Model Compression is a reinforcement learning-based framework that automatically determines the optimal compression policy (e.g., pruning ratio, quantization bit-width) for each layer of a network. AMC treats the compression process as a sequential decision-making problem, where an agent (typically an RNN controller) interacts with the environment (the model being compressed) to maximize a reward function balancing accuracy and compression ratio.

• Search Space: Defines actions like 'prune 50% of filters in layer 5' or 'quantize layer 10 to 4-bit'.
• Reward Signal: Reward = -1 * Accuracy_Drop * log(Compression_Ratio).
• Distinction from NAS: While NAS searches for novel architectures, AMC searches for the optimal application of standard compression techniques (pruning, quantization) to a given architecture. It is a policy search for compression scheduling.

One-Shot NAS is a family of techniques where the search is conducted by training a single, over-parameterized super-network (or one-shot model) just once. After training, efficient sub-networks are extracted and evaluated without requiring additional training. This dramatically reduces the computational cost of NAS, which is critical when the search target is compression for resource-constrained devices.

• Key Principle: Weight sharing across all candidate sub-networks within the super-network.
• Process:
  1. Train the super-network with a uniform path sampling or dropout strategy.
  2. Use a search algorithm (e.g., evolutionary search, gradient-based) to find high-performing sub-networks using the shared weights.
  3. Optionally fine-tune the best-found sub-network.
• Use in Compression: Ideal for searching over large spaces of pruned architectures or channel-width multipliers, as evaluating a candidate requires only a forward pass through the super-network.

Neural Architecture Transfer is a strategy that leverages architectures or architectural principles discovered via NAS on a large, well-studied dataset (e.g., ImageNet) and transfers them to a new, often smaller, target domain or task. For compression, this means using a pre-searched efficient backbone as a starting point for a specialized on-device model, rather than running a full NAS from scratch for each new application.

• Workflow:
  1. Perform a comprehensive NAS for efficiency on a proxy task/dataset.
  2. Extract the discovered efficient building blocks (e.g., mobile inverted bottleneck convolutions from MobileNetV3).
  3. Compose or fine-tune these blocks for the target task with limited data and compute.
• Benefit: Dramatically reduces the cost and time of deploying compressed models by amortizing the expensive search phase. It is a cornerstone of practical NAS for industry deployment.

Multi-Objective NAS extends the search objective beyond a single metric (like accuracy) to simultaneously optimize for multiple, often competing, constraints essential for on-device deployment. For compression, the typical objectives form a Pareto frontier across accuracy, model size (parameters), computational cost (FLOPs), latency, and energy use.

• Search Algorithms: Employs methods like NSGA-II (a genetic algorithm), Bayesian Optimization, or gradient-based multi-task learning to explore the trade-off surface.
• Output: Not a single model, but a family of models representing the best possible compromises (the Pareto frontier). A developer can then select the model that best fits their specific device constraints.
• Critical for Compression: Directly addresses the core challenge of compression scheduling: finding the optimal point in the high-dimensional space of accuracy-efficiency trade-offs.