One-Shot Neural Architecture Search

The foundational component of one-shot NAS is a single, over-parameterized supernet or one-shot model. This network contains all possible architectural choices (e.g., different layer types, kernel sizes, channel widths) within its structure. During a single training run, a shared set of weights is learned for all potential sub-networks. To evaluate a candidate architecture, a sub-network is carved from the supernet, inheriting the corresponding pre-trained weights for immediate assessment without standalone training. This weight-sharing mechanism is what makes the search computationally tractable.

Many one-shot NAS methods employ a differentiable search strategy. Architectural choices, such as which operation to use in a layer, are relaxed into continuous variables using a softmax over all options. This allows the architecture parameters (alphas) and the model weights (omegas) to be optimized jointly via gradient descent in a bi-level optimization loop:

• Inner Loop: The supernet weights are trained on the training data.
• Outer Loop: The architecture parameters are updated based on performance on a validation set. This approach, exemplified by DARTS (Differentiable Architecture Search), transforms the discrete search problem into a continuous optimization, enabling efficient search using standard backpropagation.

During supernet training or the final search phase, specific architectures are evaluated by sampling paths through the supernet. Only the weights along the sampled subnetwork's path are activated and used for a forward/backward pass. Common strategies include:

• Uniform Sampling: Each candidate operation or block has an equal chance of being selected, encouraging the supernet to learn generally useful weights.
• Gumbel-Softmax: A method for drawing discrete samples from a categorical distribution (defined by the architecture parameters) in a differentiable way.
• Single-Path methods: Further simplify evaluation by encoding the entire architecture in a single path, reducing memory overhead. The final performance ranking of architectures is based on metrics like validation accuracy or a neural predictor trained on subnetwork encodings.

A key characteristic is the strict separation of the search phase and the final training phase.

1. Search Phase: The supernet is trained once. The search algorithm (e.g., gradient-based, evolutionary) explores the space by evaluating thousands of subnetworks using the shared weights. The output is a discovered architecture description.
2. Training Phase (Retraining): The best-found architecture is constructed from scratch as a standalone model and trained independently with randomized weights. This retraining is necessary because the shared weights in the supernet are only an approximation; the final model achieves optimal performance when trained conventionally. This decoupling ensures a fair comparison and maximal performance.

A major challenge in one-shot NAS is search-training coupling bias or multi-model forgetting. Since all subnetworks share weights, optimizing for one path can degrade the performance of others. This leads to a discrepancy between the stand-alone performance of a retrained model and its estimated performance when evaluated as a subnetwork within the supernet. Factors contributing to this include:

• Operation Collapse: The supernet may over-specialize in a subset of operations.
• Optimization Interference: Gradients from different paths interfere with each other.
• Batch Normalization Statistics: Statistics are aggregated across all active paths, which may not be representative of any single subnetwork. Advanced techniques like sandwich rule sampling and path dropout are used to mitigate this bias.

One-shot NAS provides dramatic efficiency gains over earlier NAS methods:

• Computational Cost: Reduces search cost from thousands of GPU days to a few, often to < 10 GPU days, by amortizing training cost across all architectures.
• Search Space Size: Enables exploration of vast, complex search spaces (e.g., billions of architectures) that would be infeasible with reinforcement learning or evolutionary-based methods requiring independent training.
• Hardware-Aware Search: Easily integrates hardware feedback (e.g., latency, energy) by profiling subnetworks on target devices during the search, leading to Pareto-optimal architectures for specific deployment constraints. This makes it practical for designing models for edge AI and on-device inference.

One-shot NAS is used to automatically design adapter modules or LoRA rank configurations for PEFT. Instead of manually setting the adapter's bottleneck dimension or LoRA's rank, a NAS supernet explores a space of possible PEFT module architectures.

• Search Space: Defines variables for adapter placement (which transformer layers), adapter hidden size, and LoRA rank values.
• Supernet Training: A base model (e.g., LLaMA) is equipped with a supernet of all possible PEFT modules. The shared weights are trained on a downstream task.
• Result: The search identifies the most effective and compact PEFT configuration for a given task and compute budget, maximizing adaptation quality per trained parameter.

Neural architecture search (NAS) is the automated process of discovering high-performing neural network structures for a given task. It treats the network design—choices like layer types, connections, and widths—as a search problem.

• Core Goal: Replace manual architecture engineering with algorithmic optimization.
• Search Components: Defined by a search space (possible architectures), a search strategy (e.g., reinforcement learning, evolution), and a performance estimation strategy (e.g., full training, weight sharing).
• Application to PEFT: NAS is used to automatically discover optimal adapter placements, LoRA ranks, or sparse tuning masks, making PEFT configuration data-driven and efficient.

Weight sharing is the foundational efficiency technique behind one-shot NAS. A single, over-parameterized supernet is trained once, and all candidate sub-architectures inherit and share its weights for evaluation.

• Mechanism: Instead of training thousands of architectures from scratch, each is treated as a sub-graph of the supernet. Performance is estimated by activating only the corresponding paths.

• Key Benefit: Reduces the computational cost of NAS from thousands of GPU days to a few, enabling practical use.

• PEFT Connection: In automated PEFT, weight sharing allows for the rapid evaluation of different adapter configurations or pruning masks within a single trained supernet framework.

Differentiable NAS (DNAS) is a gradient-based search method that relaxes the discrete search space into a continuous one. Architecture parameters are optimized via gradient descent alongside model weights.

• How it Works: The supernet uses continuous relaxation (e.g., Gumbel-Softmax) for architectural choices. The search becomes a bi-level optimization: training network weights and architecture parameters.

• Advantage: More efficient than reinforcement learning or evolutionary-based NAS, converging faster due to gradient signals.

• Relevance: DNAS is a common implementation strategy for one-shot NAS, making it directly applicable to searching for optimal PEFT module architectures in a differentiable manner.

The Once-For-All (OFA) network is a training paradigm and supernet designed to support diverse depth, width, kernel size, and resolution configurations without retraining.

• Core Idea: Decouple training from search and deployment. Train one large network that subsumes many smaller ones.

• Process: 1) Train the OFA supernet. 2) Search for the best sub-network for a target hardware constraint (e.g., latency, model size) using the shared weights. 3) Deploy the specialized sub-network.

• PEFT Analogy: OFA's philosophy aligns with training a base model with a flexible PEFT superstructure (e.g., many potential adapter slots), then searching for the most efficient sub-configuration for a specific task or device.

In NAS and automated PEFT configuration, the search space is the predefined set of all possible architectural or configuration choices the algorithm can explore.

• For NAS: Defines allowable operations (convolution, pooling), connection patterns, number of layers, and channel widths.

• For Automated PEFT: Defines the tunable aspects of the adaptation method, such as:

  • Adapter insertion points (which layers), hidden dimension, and reduction factor.
  • LoRA rank (r) values and which weight matrices to target.
  • Sparse tuning masks (percentage and distribution of updated parameters).

• Importance: A well-designed search space balances expressiveness with tractability, crucial for finding effective, efficient adaptations.

These are performance estimation techniques used to rank candidate architectures without full training, critical for scaling NAS and PEFT search.

• Neural Predictor: A separate ML model (e.g., a regressor) trained to predict the final accuracy of an architecture from its encoded representation.

• Zero-Cost Proxy: An ultra-fast heuristic metric requiring minimal computation (e.g., a single forward/backward pass). Examples include gradient norm, synaptic saliency, or NASWOT (no training score).

• Use Case: After supernet training, these tools rapidly screen thousands of candidate PEFT configurations to identify the most promising few for final fine-tuning validation, making the search process orders of magnitude faster.