Once-For-All Network

The core of an OFA network is a single, over-parameterized supernet trained to contain a vast number of sub-networks within its shared weight parameters. This eliminates the need to train thousands of independent models from scratch. The supernet is trained with a progressive shrinking algorithm that sequentially introduces architectural flexibility (e.g., depth, width, kernel size), allowing sub-networks of varying sizes to inherit well-trained weights.

• Key Benefit: Drastically reduces the search cost for optimal sub-networks from thousands of GPU-days to near zero.

OFA networks are elastic across multiple architectural dimensions, allowing sub-networks to be sampled by adjusting:

• Depth: Number of layers or blocks.
• Width: Number of channels (filters) per layer.
• Kernel Size: Spatial size of convolutional filters (e.g., 3x3, 5x5, 7x7).
• Resolution: Input image size.

This multi-dimensional elasticity creates a continuous spectrum of model sizes and computational costs (e.g., from 5M to 600M FLOPs) from one supernet, enabling precise matching to any target device's memory and latency budget.

OFA decouples training from search. After the supernet is trained once, a zero-shot search can be performed for any new hardware target. The search algorithm evaluates candidate sub-networks by:

1. Sampling a sub-network configuration from the elastic space.
2. Inheriting weights directly from the supernet (no fine-tuning).
3. Measuring its accuracy on a validation set and its latency/energy on the actual target hardware (e.g., a specific microcontroller).

This process finds the Pareto-optimal sub-network for the exact trade-off between accuracy and on-device performance required.

Training the supernet is non-trivial; naively training all sub-networks simultaneously leads to interference and poor performance. OFA uses progressive shrinking:

• Stage 1: Train the largest possible sub-network (full depth, width, kernel size) as a baseline.
• Stage 2: Finetune this network while allowing smaller kernel sizes (e.g., 3,5,7) to be sampled, making weights good for all kernels.
• Stage 3: Finetune further while allowing reduced depth (fewer layers) to be sampled.
• Stage 4: Finetune finally while allowing reduced width (fewer channels) to be sampled.

This sequential training stabilizes learning and ensures high accuracy across all sub-networks.

For TinyML, OFA provides a practical path to custom, efficient models. A developer can:

1. Start with a pre-trained OFA supernet (e.g., OFA-ResNet50, OFA-MobileNetV3).
2. Profile their target microcontroller's SRAM, flash, and latency limits.
3. Run the hardware-aware search to find the best sub-network for those exact constraints.
4. Extract and deploy that sub-network as a standalone, compact model file.

This enables ImageNet-scale vision models to run on devices with under 256KB of SRAM, as demonstrated by frameworks like MCUNet which builds upon OFA principles.

OFA addresses critical limitations of earlier Neural Architecture Search methods:

• Cost: Traditional NAS trains each candidate model independently (prohibitively expensive). Differentiable NAS (DARTS) is faster but often finds architectures that perform poorly when trained standalone (the discrepancy gap).
• Hardware Ignorance: Many NAS methods optimize only for FLOPs or parameters, which are poor proxies for real on-device latency and energy use.
• Flexibility: A model found by traditional NAS is fixed. If hardware changes, the entire NAS process must be rerun.

OFA provides a reusable, hardware-aware model zoo from a single training run, offering unparalleled efficiency for embedded system design.

OFA enables dynamic performance-accuracy trade-offs on a single device. After deployment, the device can switch between different sub-networks extracted from the super-network based on real-time system conditions. For example:

• A battery-powered sensor could use a larger, more accurate sub-network when mains power is connected, and automatically switch to a tiny, ultra-efficient sub-network when running on battery to extend operational life.
• A device under high CPU load from other tasks could temporarily downgrade to a faster, lighter sub-network to maintain overall system responsiveness, then revert to the high-accuracy model when resources are free. This provides a form of adaptive inference without requiring multiple, separately trained models to be stored on the device.

OFA synergizes with other model compression techniques to push the boundaries of what's possible on microcontrollers. The super-network training process inherently encourages weight sharing and robustness across architectures. This creates a strong foundation for applying post-training techniques:

• Quantization: The shared weights of the OFA super-network can be quantized (e.g., to INT8), and this quantization is inherited by all sub-networks, ensuring consistent behavior.
• Pruning: The super-network can be globally pruned, removing redundant connections that are unimportant across all sub-networks, leading to a more compact base for subsequent sub-network extraction.
• Knowledge Distillation: The large super-network can act as the teacher model for distilling knowledge into even smaller, non-shared student models for ultra-low-footprint deployment.

Neural Architecture Search (NAS) is an automated process for discovering optimal neural network architectures for a given task and set of constraints. Instead of manual design, NAS uses algorithms—such as reinforcement learning, evolutionary strategies, or gradient-based methods—to explore a vast space of possible layer types, connections, and hyperparameters.

• Core Mechanism: A controller (often an RNN or optimizer) proposes candidate architectures, which are trained and evaluated. The performance feedback is used to improve subsequent proposals.
• Primary Goal: To automate and often surpass human-designed architectures, especially under specific constraints like model size, latency, or energy consumption.
• Relation to OFA: OFA is a supernetwork training paradigm that enables highly efficient hardware-aware NAS. It provides a weight-shared search space, allowing the evaluation of thousands of sub-networks without individual training.

Hardware-Aware Neural Architecture Search (HW-NAS) is a specialized variant of NAS where the search algorithm directly optimizes for performance metrics on target hardware. The search objective incorporates direct measurements or accurate proxies for:

• Inference Latency (ms)
• Energy Consumption (mJ)
• Peak Memory Usage (KB)
• Compatibility with specific accelerators (e.g., NPU instruction sets)

Implementation Methods:

• Look-Up Tables (LUTs): Pre-measured latency/energy costs for each operator on the target device.
• On-Device Profiling: Deploying and measuring candidate models directly on the hardware during search.
• Analytical Models: Mathematical models predicting performance from operator FLOPs and memory access patterns.

Relation to OFA: OFA is a quintessential HW-NAS enabler. By training a single supernet, it allows for instant, cost-free evaluation of sub-networks against hardware constraints, making the search for a device-specific optimal model practical.

MCUNet is a pioneering co-design framework that jointly optimizes the Tiny Neural Network architecture (TinyNAS) and the underlying memory-aware inference engine (TinyEngine) to run ImageNet-scale vision models on microcontrollers.

• TinyNAS: A hardware-aware NAS that searches for networks operating within a tight memory budget (e.g., < 256KB SRAM). It uses a two-stage scaling method to first find a good base layer pattern, then scales it up optimally.
• TinyEngine: A latency-driven inference engine that generates tailored, in-place memory scheduling code to eliminate peak memory usage. It acts as a "compiler-aware" NAS, optimizing the graph for the engine's allocation patterns.

Key Achievement: Demonstrated that co-design—optimizing the model and the runtime system together—is essential for pushing the limits of on-device AI. While OFA provides the search space, MCUNet exemplifies the full-stack system design needed for real microcontroller deployment.

ProxylessNAS is a neural architecture search method that directly learns architectures on the target task and hardware without using a proxy (like training on a smaller dataset or for fewer epochs).

• Core Innovation: It removes the proxy, which can introduce bias, by performing the search directly on the full ImageNet dataset and hardware. It uses path-level binarization and gradient-based optimization to make this computationally feasible.
• Hardware Latency Loss: It directly incorporates a latency loss term into the training objective, allowing the search to optimize for real device speed (e.g., on a mobile CPU or GPU).
• Contrast with OFA: While both are hardware-aware, ProxylessNAS typically searches for a single optimal architecture for one constraint. OFA, in contrast, trains one supernet to produce many optimal architectures for diverse constraints. ProxylessNAS is "one-shot" per search; OFA is "once-for-all."

Differentiable Architecture Search (DARTS) is a gradient-based NAS method that formulates the search space as a continuous, over-parameterized network. In this supernet, the choice between operations (e.g., 3x3 conv, 5x5 conv, skip connect) is represented by a set of continuous architecture parameters.

• Search Process: The supernet weights and the architecture parameters are optimized jointly via gradient descent on a validation set. The search concludes by selecting the operation with the highest learned architecture parameter weight.
• Advantage: Significantly faster (order of days on a single GPU) than earlier reinforcement learning or evolutionary-based NAS methods.
• Relation to OFA: DARTS introduced the influential concept of the differentiable supernet. OFA adopts and extends this weight-sharing principle but focuses on multi-dimension elasticity (depth, width, kernel size) and decouples the supernet training from the search phase, enabling the efficient extraction of many sub-networks post-training.

A Supernet (or One-shot Model) is a single, over-parameterized neural network that encompasses many possible sub-networks (architectures) within its structure through weight sharing.

• Mechanism: During training, only a subset of the supernet's weights are activated and updated per training step, often via random path sampling. This encourages all shared weights to generalize well across many potential sub-networks.
• Purpose: To amortize the massive cost of training each candidate architecture from scratch in NAS. After the supernet is trained, sub-networks can be evaluated by simply selecting the corresponding paths without any additional training.
• Foundation of OFA: The Once-For-All Network is a specialized supernet with predefined, structured elasticity dimensions. Its training process (progressive shrinking) is specifically designed to maintain the performance of the largest sub-network while ensuring high accuracy for the smallest, most constrained sub-networks—a key challenge in supernet training.