Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

SuperFedNAS: Cost-Efficient Federated Neural Architecture Search for On-Device Inference

Alind Khare, Animesh Agrawal, Aditya Annavajjala, Payman Behnam, Myungjin Lee, Hugo Latapie, Alexey Tumanov

TL;DR

This work tackles the challenge of scalable NAS in Federated Learning for on-device inference where deployment targets vary across hardware, latency, and energy constraints. It introduces SuperFedNAS, which decouples training and search by co-training a massive supernet containing approximately $5\cdot 10^8$ architectures within FL, followed by a training-free local NAS to extract specialized subnets for any deployment target, achieving $O(1)$ search cost per target. A key contribution is MaxNet, a multi-objective FL training algorithm that minimizes the losses of the worst-performing subnets on each client data partition, reducing interference and improving convergence with weight-sharing-aware subnet sampling and decayed aggregation. The paper demonstrates that SuperFedNAS attains up to $37.7\%$ higher accuracy for the same MACs or up to $8.13\times$ MAC reduction for the same accuracy compared to existing federated NAS methods, while reducing training cost by up to $11\times$ to satisfy 20 deployment targets, across CIFAR, CINIC, and Shakespeare datasets. Overall, the approach provides a practical, scalable path to deployment-aware NAS in FL, enabling diverse, hardware-tailored models for on-device inference without retraining for each target.

Abstract

Neural Architecture Search (NAS) for Federated Learning (FL) is an emerging field. It automates the design and training of Deep Neural Networks (DNNs) when data cannot be centralized due to privacy, communication costs, or regulatory restrictions. Recent federated NAS methods not only reduce manual effort but also help achieve higher accuracy than traditional FL methods like FedAvg. Despite the success, existing federated NAS methods still fall short in satisfying diverse deployment targets common in on-device inference like hardware, latency budgets, or variable battery levels. Most federated NAS methods search for only a limited range of neuro-architectural patterns, repeat them in a DNN, thereby restricting achievable performance. Moreover, these methods incur prohibitive training costs to satisfy deployment targets. They perform the training and search of DNN architectures repeatedly for each case. SuperFedNAS addresses these challenges by decoupling the training and search in federated NAS. SuperFedNAS co-trains a large number of diverse DNN architectures contained inside one supernet in the FL setting. Post-training, clients perform NAS locally to find specialized DNNs by extracting different parts of the trained supernet with no additional training. SuperFedNAS takes O(1) (instead of O(N)) cost to find specialized DNN architectures in FL for any N deployment targets. As part of SuperFedNAS, we introduce MaxNet - a novel FL training algorithm that performs multi-objective federated optimization of a large number of DNN architectures ($\approx 5*10^8$) under different client data distributions. Overall, SuperFedNAS achieves upto 37.7% higher accuracy for the same MACs or upto 8.13x reduction in MACs for the same accuracy than existing federated NAS methods.

SuperFedNAS: Cost-Efficient Federated Neural Architecture Search for On-Device Inference

TL;DR

This work tackles the challenge of scalable NAS in Federated Learning for on-device inference where deployment targets vary across hardware, latency, and energy constraints. It introduces SuperFedNAS, which decouples training and search by co-training a massive supernet containing approximately architectures within FL, followed by a training-free local NAS to extract specialized subnets for any deployment target, achieving search cost per target. A key contribution is MaxNet, a multi-objective FL training algorithm that minimizes the losses of the worst-performing subnets on each client data partition, reducing interference and improving convergence with weight-sharing-aware subnet sampling and decayed aggregation. The paper demonstrates that SuperFedNAS attains up to higher accuracy for the same MACs or up to MAC reduction for the same accuracy compared to existing federated NAS methods, while reducing training cost by up to to satisfy 20 deployment targets, across CIFAR, CINIC, and Shakespeare datasets. Overall, the approach provides a practical, scalable path to deployment-aware NAS in FL, enabling diverse, hardware-tailored models for on-device inference without retraining for each target.

Abstract

Neural Architecture Search (NAS) for Federated Learning (FL) is an emerging field. It automates the design and training of Deep Neural Networks (DNNs) when data cannot be centralized due to privacy, communication costs, or regulatory restrictions. Recent federated NAS methods not only reduce manual effort but also help achieve higher accuracy than traditional FL methods like FedAvg. Despite the success, existing federated NAS methods still fall short in satisfying diverse deployment targets common in on-device inference like hardware, latency budgets, or variable battery levels. Most federated NAS methods search for only a limited range of neuro-architectural patterns, repeat them in a DNN, thereby restricting achievable performance. Moreover, these methods incur prohibitive training costs to satisfy deployment targets. They perform the training and search of DNN architectures repeatedly for each case. SuperFedNAS addresses these challenges by decoupling the training and search in federated NAS. SuperFedNAS co-trains a large number of diverse DNN architectures contained inside one supernet in the FL setting. Post-training, clients perform NAS locally to find specialized DNNs by extracting different parts of the trained supernet with no additional training. SuperFedNAS takes O(1) (instead of O(N)) cost to find specialized DNN architectures in FL for any N deployment targets. As part of SuperFedNAS, we introduce MaxNet - a novel FL training algorithm that performs multi-objective federated optimization of a large number of DNN architectures () under different client data distributions. Overall, SuperFedNAS achieves upto 37.7% higher accuracy for the same MACs or upto 8.13x reduction in MACs for the same accuracy than existing federated NAS methods.
Paper Structure (12 sections, 6 equations, 6 figures, 6 tables)

This paper contains 12 sections, 6 equations, 6 figures, 6 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Method
  4. Problem Formulation
  5. Naive Supernet FL-Training Algorithms
  6. MaxNet: Proposed Supernet FL-Training Algorithm
  7. SuperFedNAS's Search Stage
  8. Experiments
  9. Setup
  10. Comparison with Baselines
  11. Ablation Study
  12. Conclusion

Figures (6)

  • Figure 1: SuperFedNAS vs Existing Federated NAS methods fednasfedpnas
  • Figure 2: Supernet and its Naive FL-Training Methods. (a) Supernet consists of subnets within it that differ in depth/width. (b) Multi-Stage method sends entire supernet to clients. (c) Single-Stage Supernet FL-Training sends subnets to clients.
  • Figure 3: Naive Supernet FL-Training Cost/Convergence Comparison. Communication/Computational cost (left,middle) are compared for naive supernet FL approaches with FedAvg training of largest/smallest subnet. The right plot compares convergence over 1500 rounds. Naive-methods have slow convergence and high training cost.
  • Figure 4: Training Cost Comparison.
  • Figure 5: SuperFedNAS's DNN Arch. Specialization. Specialized DNNs found by SuperFedNAS's search stage on different hardware/latency targets. SuperFedNAS finds a more accurate DNN for RTX 2080Ti GPU (91.56%) compared to AMD CPU (85.25%). SuperFedNAS's specialized DNNs are: shallow/thin for AMD CPU, wide/deep for GPU.
  • ...and 1 more figures