DεpS: Delayed ε-Shrinking for Faster Once-For-All Training

TL;DR

D$\varepsilon$pS tackles the high cost of multi-target deployment optimization by delaying the shrinkage phase in once-for-all training. It introduces FM-Warmup to initialize the supernet with a partially trained full model, $\mathcal{E}$-Shrinking to gradually ramp subnet learning rates during shrinking, and IKD-Warmup to distill knowledge from progressively better partial full models to subnets. The approach yields faster training with lower GPU hours while maintaining or improving subnet accuracy across CIFAR and ImageNet datasets, across MobilenetV3 and ProxylessNAS spaces, achieving up to 1.83% improved ImageNet-1k top-1 accuracy or similar accuracy with 1.3x FLOPs reduction and up to 2.5x training-cost reduction compared to OFA/BigNAS baselines. These results, along with targeted ablations, demonstrate the effectiveness and robustness of the three components in making scalable, hardware-aware NAS practical for diverse deployment targets.

Abstract

CNNs are increasingly deployed across different hardware, dynamic environments, and low-power embedded devices. This has led to the design and training of CNN architectures with the goal of maximizing accuracy subject to such variable deployment constraints. As the number of deployment scenarios grows, there is a need to find scalable solutions to design and train specialized CNNs. Once-for-all training has emerged as a scalable approach that jointly co-trains many models (subnets) at once with a constant training cost and finds specialized CNNs later. The scalability is achieved by training the full model and simultaneously reducing it to smaller subnets that share model weights (weight-shared shrinking). However, existing once-for-all training approaches incur huge training costs reaching 1200 GPU hours. We argue this is because they either start the process of shrinking the full model too early or too late. Hence, we propose Delayed $ε$-Shrinking (D$ε$pS) that starts the process of shrinking the full model when it is partially trained (~50%) which leads to training cost improvement and better in-place knowledge distillation to smaller models. The proposed approach also consists of novel heuristics that dynamically adjust subnet learning rates incrementally (E), leading to improved weight-shared knowledge distillation from larger to smaller subnets as well. As a result, DEpS outperforms state-of-the-art once-for-all training techniques across different datasets including CIFAR10/100, ImageNet-100, and ImageNet-1k on accuracy and cost. It achieves 1.83% higher ImageNet-1k top1 accuracy or the same accuracy with 1.3x reduction in FLOPs and 2.5x drop in training cost (GPU*hrs)

Figures (7)

• Figure 1: D$\epsilon$pS reduces training time compared to existing approaches like OFAofa & BigNASbignas.
• Figure 2: Supernetwork initialization. D$\epsilon$pS provides better initialization for the supernetwork for smaller subnets compared to OFA due to FMWarmup. This validates the hypothesis that the supernet weights become specialized if the full model is trained to completion (OFA), resulting in poorer accuracy of subnetworks with increased training of the full model.
• Figure 3: Gradients w/ & w/o Shrinking on Mobilenet-Based Supernet. Delayed Shrinking causes gradients ($G_{\text{Shrink}}$) to differ from the full model gradient ($G_{\text{noShrink}}$) leading to rapid changes in the supernet's weights. $\mathcal{E}$-Shrinking's gradient ($G_{\text{Shrink}}(\mathcal{E})$) reduces such differences and avoids rapid weight changes.
• Figure 4: Gradient Magnitude Over Time. Gradient magnitude with ($\mathcal{G}_{shrink}(\mathcal{E}, t)$) and without ($\mathcal{G}_{shrink}(t)$) $\mathcal{E}$-shrinking is compared w.r.t the initial full model gradient ($\mathcal{G}_{Noshrink}$) over shrinking steps. $\mathcal{E}$-shrinking avoids sudden changes in the supernet parameters by lowering the gradient magnitude.
• Figure 5: D$\epsilon$pS's Accuracy Improvement across Datasets. The comparison of D$\epsilon$pS with the baselines is shown w.r.t. accuracy (of subnets) for CIFAR10/100, ImageNet-100, and ImageNet-1k datasets. D$\epsilon$pS consistently outperforms the baselines across all the datasets and achieves upto $2.1\%$ better accuracy for the same FLOPs or upto $2.3$x FLOP reduction at same accuracy.