FedRolex: Model-Heterogeneous Federated Learning with Rolling Sub-Model Extraction

Samiul Alam; Luyang Liu; Ming Yan; Mi Zhang

Source-linked AI summary

FedRolex: Model-Heterogeneous Federated Learning with Rolling Sub-Model Extraction

Samiul Alam, Luyang Liu, Ming Yan, Mi Zhang

arXiv:2212.01548v2Dec 3, 2022cs.LGcs.CRcs.CVcs.DC

TL;DR

Model-homogeneous cross-device FL can exclude low-end clients and prevent training server models larger than client capacities. FedRolex uses rolling sub-model extraction to support heterogeneous models, and it consistently outperforms partial-training baselines while improving inclusiveness.

Problem
Model-homogeneous cross-device FL excludes low-end clients and restricts server models to capacities that clients can train.
Method
FedRolex applies partial training with a rolling sub-model extraction window so different parts of the global model are trained across communication rounds.
Results
FedRolex consistently outperforms state-of-the-art partial-training methods across models and datasets at both small and large scales.
Takeaways & Limitations
FedRolex can train a global server model larger than the largest client model and contributes to making federated learning more inclusive.
Takeaways & Limitations
A full convergence analysis is left for future work, and deployment of client models after global training is a separate task.

Abstract

from arXiv · show

Most cross-device federated learning (FL) studies focus on the model-homogeneous setting where the global server model and local client models are identical. However, such constraint not only excludes low-end clients who would otherwise make unique contributions to model training but also restrains clients from training large models due to on-device resource bottlenecks. In this work, we propose FedRolex, a partial training (PT)-based approach that enables model-heterogeneous FL and can train a global server model larger than the largest client model. At its core, FedRolex employs a rolling sub-model extraction scheme that allows different parts of the global server model to be evenly trained, which mitigates the client drift induced by the inconsistency between individual client models and server model architectures. We show that FedRolex outperforms state-of-the-art PT-based model-heterogeneous FL methods (e.g. Federated Dropout) and reduces the gap between model-heterogeneous and model-homogeneous FL, especially under the large-model large-dataset regime. In addition, we provide theoretical statistical analysis on its advantage over Federated Dropout and evaluate FedRolex on an emulated real-world device distribution to show that FedRolex can enhance the inclusiveness of FL and boost the performance of low-end devices that would otherwise not benefit from FL. Our code is available at: https://github.com/AIoT-MLSys-Lab/FedRolex

1 Introduction

Model-homogeneous cross-device FL excludes resource-constrained clients and limits server models to client-device capacities. FedRolex addresses these constraints with rolling sub-model extraction that evenly trains the global model across rounds.

Cross-device FL commonly uses identical server and client models, excluding low-end devices with insufficient resources.
Model homogeneity also restricts the server from training larger models when clients cannot support the same capacity.
FedRolex uses rolling sub-model extraction, advancing a window across the global model so different parameter regions are trained in successive rounds.
Evenly training global-model parameters mitigates client drift caused by mismatched client and server architectures.
FedRolex is evaluated in both small-model small-dataset and large-model large-dataset regimes.
FedRolex consistently outperforms state-of-the-art partial-training methods and reduces the heterogeneous-versus-homogeneous accuracy gap, especially at large scale.

2 Related Work

Related work addresses model-heterogeneous FL mainly through knowledge distillation or partial training. These approaches differ in how heterogeneous client information or server sub-models are handled.

Knowledge Distillation (KD)-based Model-Heterogeneous FL: Knowledge-distillation methods aggregate client knowledge at the server using client outputs or models, typically with public data.
Partial Training (PT)-based Model-Heterogeneous FL: Partial-training methods avoid full-model client training by extracting sub-models from the global server model.

3 Methodology

FedRolex formulates model-heterogeneous FL around client-specific sub-model capacities and trains those sub-models through rolling extraction and selective aggregation. Its rolling scheme contrasts with random and static extraction by evenly covering the global model across rounds while supporting heterogeneous clients.

Model-Heterogeneous FL: Client models have heterogeneous capacities, with each capacity defined as the proportion of nodes extracted from each global-model layer.The extracted sub-model can change across rounds, so the client objective may also change.
FedRolex: FedRolex trains each client on a sub-model extracted from the global server model and sends the corresponding updates back for aggregation.The server extracts sub-models with different capacities for participating clients.
Rolling Extraction: A rolling window advances each communication round and sequentially loops over all global-model parts so different parts are evenly trained.The window advances by a configurable step size and continues until convergence.
Rolling Extraction: Across three example rounds, the large- and small-capacity client sub-models shift through different parameter groups as the rolling window advances.The examples move the large model from {a, b, c, d} to {b, c, d, e} to {c, d, e, a}, while the small model moves from {c, d, e} to {d, e, a} to {e, a, b}.
Aggregation: FedRolex selectively averages updates for each global parameter using only clients that updated that parameter, leaving untouched parameters unchanged.Parameters shared by multiple client sub-models receive the average of their updates, while singly updated parameters use that client’s update.
Comparison with Random Extraction: Random extraction produces uneven update frequencies and greater vulnerability to client drift, whereas FedRolex reaches repeated coverage of all I sub-models in mI rounds rather than a larger expected order.Federated Dropout’s expected number of rounds is O(I log(I) + I(m −1) log log I).
Comparison with Static Extraction: Static extraction reuses designated sub-models, requires full-size clients to cover the global model, and trains different model parts on different client-data distributions.Consequently, static methods restrict the global model to the largest client and can degrade global training quality.

4 Experiments

FedRolex is evaluated across small-model/small-dataset and large-model/large-dataset regimes using heterogeneous client capacities, multiple baselines, and global/local accuracy metrics. Across these settings, it consistently outperforms PT-based methods, narrows the gap to model-homogeneous FL, and improves performance under realistic device distributions.

Datasets and Models: Experiments cover CIFAR-10 and CIFAR-100 with PreResNet18, plus Stack Overflow with a modified 3-layer Transformer.The evaluation spans small-model/small-dataset and large-model/large-dataset regimes.
Data Heterogeneity: Data heterogeneity is varied using label restrictions: CIFAR-10 uses L = 2 or 5, while CIFAR-100 uses L = 20 or 50.These settings represent high and low heterogeneity levels, roughly corresponding to DirK(α) values of 0.1 and 0.5.
Model Heterogeneity: Client model capacities use β = {1, 1/2, 1/4, 1/8, 1/16}, generated by varying convolutional kernels or attention-head hidden-layer nodes.For example, β = 1/2 denotes half the largest client model capacity.
Performance Comparison: FedRolex consistently outperforms HeteroFL and Federated Dropout across both evaluation regimes and under low and high data heterogeneity.Under high data heterogeneity, random sub-model extraction performs worse than the deterministic extraction used by FedRolex and HeteroFL.
Model-Homogeneous Comparison: FedRolex narrows the global-accuracy gap to upper-bound model-homogeneous FL and matches it on Stack Overflow, while HeteroFL and Federated Dropout fall below the smallest-model homogeneous case.The upper-bound case uses β = {1}, and the lower-bound case uses β = {1/16}.
Real-World Device Distribution: Under an emulated real-world device distribution, FedRolex improves global accuracy over the lower-bound homogeneous case across all three datasets and improves local accuracy on both CIFAR datasets.The local-accuracy gains hold under both low and high data heterogeneity.

5 Conclusion

FedRolex is a partial-training approach for model-heterogeneous federated learning that combines rolling sub-model extraction with statistical analysis and evaluation across scales and device distributions. The authors report consistent advantages over state-of-the-art partial-training methods and improved inclusiveness, while identifying convergence analysis and post-training deployment as future work.

FedRolex trains a global server model larger than the largest client model using partial training and rolling sub-model extraction.The rolling scheme evenly trains global-model parameters to mitigate client drift induced by model heterogeneity.
FedRolex consistently outperforms state-of-the-art partial-training methods across models and datasets at both small and large scales.
Evaluation on an emulated real-world device distribution shows that FedRolex contributes to making federated learning more inclusive.
A full convergence analysis remains future work, and determining which models to deploy after training the global server model is a separate task.The deployment question is especially relevant when the global server model is large.

A.1 Statistical Analysis

The statistical analysis characterizes how many randomly selected rounds are needed to cover all indices, including repeated coverage, using coupon-collection and Poisson-process arguments. For one selection per index, the expected coverage time is asymptotically I log(I), while repeated selection requires an additional m-dependent term.

The proof for one-time coverage defines E(i) recursively as the expected rounds to choose exactly i indices at least once.It uses E(1) = 1 and adds I/(I + 1 − i) rounds to obtain a new index.
I log(I) is the asymptotic expected number of rounds needed to choose all I indices at least once.
The Poisson-process proof is introduced because the recursive proof for one-time coverage cannot be generalized to m ≥ 2.
For repeated coverage, the analysis models index selections as independent Poisson processes with rate 1/I and defines Xj as the time to select index j m times.
The expected number of rounds to choose all indices at least m times is asymptotically I log(I) + I(m − 1) log log I as I tends to infinity.

A.2 Formal Definition of Selective Aggregation Scheme

Selective aggregation averages each global parameter over the selected clients that update it, with client weights determining each client's contribution. The method allows weights to reflect model capacity, data volume, or other factors, although equal weighting is used by default.

Each global parameter is aggregated from the selected clients that update that parameter.For parameter θ[i,k], the relevant client set is Mk, the clients updating its kth parameter in layer i.
Client weights pm can reflect factors such as model capacity and the number of data points held by a client.
FedRolex assumes equal client weights by default, with pm = 1/N unless otherwise stated.

A.3 Ablation Study: Impact of Different Weighing Schemes

The ablation study compares model-size-based, update-based, and hybrid client weighting with non-weighting under high data heterogeneity. None of the three weighting schemes is significantly better, so FedRolex uses non-weighting by default.

The study evaluates model-size-based, model-update-based, and hybrid client weighting schemes against non-weighting.The schemes weight clients by kernels, updates, or both model size and model updates.
The three weighting schemes are not significantly better than the non-weighting scheme in model accuracy.
FedRolex therefore uses the non-weighting scheme.

A.4 Ablation Study: Impact of Overlapping Kernels

The study evaluates how overlap between rolling-window strides affects global model accuracy on CIFAR-10 and CIFAR-100 under different data-heterogeneity levels. Overlap influences accuracy, but the relationship is non-linear and inconsistent.

The ablation varies rolling-window kernel overlap for ResNet-18 on CIFAR-10 and CIFAR-100.The evaluation considers both low and high data heterogeneity.
FedRolex uses complete overlap, r = 1, advancing convolutional kernels by one position between iterations.The rolling window advances through each convolution layer in strides; FedRolex sets the overlap to its maximum value.
The overlap value affects global model accuracy, but its impact is non-linear and inconsistent.Figure 6 reports this pattern across the CIFAR-10 and CIFAR-100 settings.

A.5 Ablation Study: Impact of Client Participation Rate

FedRolex is evaluated against HeteroFL and Federated Dropout at 5%, 10%, and 20% client participation rates on CIFAR-10. It consistently outperforms both baselines across all tested participation rates.

The experiments compare FedRolex, HeteroFL, and Federated Dropout under 5%, 10%, and 20% client participation.The main paper uses a 10% participation rate, with 5% and 20% added to examine sensitivity.
Table 6 summarizes performance under the different client participation rates.

A.6 Communication and Computation Costs of FedRolex

FedRolex’s communication and computation costs are measured using average participating-client model size, FLOPs, and parameter counts per round. The reported results compare these costs with homogeneous upper and lower bounds.

Communication cost is the average model size sent by participating clients per round.Computation overhead is measured as the average FLOPs and parameter counts across participating clients per round.
The upper and lower bounds correspond to homogeneous settings using the largest and smallest models, respectively.These bounds contextualize FedRolex’s communication and computation costs.
Compared with the upper bound, FedRolex significantly reduces communication cost and computation overhead while achieving comparable model accuracy.Compared with the lower bound, FedRolex incurs higher communication and computation costs but achieves much higher model accuracy.