NeFL: Nested Model Scaling for Federated Learning with System Heterogeneous Clients

TL;DR

NeFL introduces a flexible nested model scaling framework for federated learning that improves performance and resource efficiency across heterogeneous clients by training multiple submodels with decoupled parameters.

Contribution

It proposes a novel nested federated learning framework that allows flexible model subdivision and parameter decoupling to handle system heterogeneity in FL.

Findings

Achieves 7.63% performance improvement on CIFAR-100 for the worst-case submodel.

Effectively enables resource-constrained devices to participate in FL.

Aligns with recent FL advances like pre-trained models and statistical heterogeneity.

Abstract

Federated learning (FL) enables distributed training while preserving data privacy, but stragglers-slow or incapable clients-can significantly slow down the total training time and degrade performance. To mitigate the impact of stragglers, system heterogeneity, including heterogeneous computing and network bandwidth, has been addressed. While previous studies have addressed system heterogeneity by splitting models into submodels, they offer limited flexibility in model architecture design, without considering potential inconsistencies arising from training multiple submodel architectures. We propose nested federated learning (NeFL), a generalized framework that efficiently divides deep neural networks into submodels using both depthwise and widthwise scaling. To address the inconsistency arising from training multiple submodel architectures, NeFL decouples a subset of parameters from those being trained for each submodel. An averaging method is proposed to handle these decoupled parameters during aggregation. NeFL enables resource-constrained devices to effectively participate in the FL pipeline, facilitating larger datasets for model training. Experiments demonstrate that NeFL achieves performance gain, especially for the worst-case submodel compared to baseline approaches (7.63% improvement on CIFAR-100). Furthermore, NeFL aligns with recent advances in FL, such as leveraging pre-trained models and accounting for statistical heterogeneity. Our code is available online.