Breaking the Capacity Bottleneck in Model-Heterogeneous Federated Learning via Gradual Model Restoration

Chengjie Ma, Seungeun Oh, Jihong Park, Seong-Lyun Kim
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:84807-84845, 2026.

Abstract

Federated learning (FL) enables distributed model training, yet in heterogeneous deployments, Bandwidth-Constrained Clients (BCCs) often contribute inefficiently due to limited uplink bandwidth. In model-heterogeneous FL with fixed small sub-models, BCCs with sub-models may improve quickly in early rounds but become under-parameterized later, resulting in slow convergence and poor generalization. To address this challenge, we propose FedGMR, a federated learning framework centered around Gradual Model Restoration (GMR), where GMR progressively increases each client’s sub-model density during training, allowing BCCs to remain effective contributors throughout optimization. To make GMR practical under real-world heterogeneity, FedGMR is realized as an end-to-end workflow with asynchronous coordination and stable, mask-aware aggregation. We further establish convergence guarantees, showing that the aggregation error scales with the average sub-model density across clients and rounds, and that GMR provably narrows the gap toward full-model FL. Extensive experiments on FEMNIST, CIFAR-10, ImageNet-100, and StackOverflow demonstrate that FedGMR improves both convergence speed and final accuracy, especially under severe heterogeneity and non-IID data distributions. Code is available at https://github.com/machengjie321/ICML2026-FedGMR.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-ma26ab, title = {Breaking the Capacity Bottleneck in Model-Heterogeneous Federated Learning via Gradual Model Restoration}, author = {Ma, Chengjie and Oh, Seungeun and Park, Jihong and Kim, Seong-Lyun}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {84807--84845}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/ma26ab/ma26ab.pdf}, url = {https://proceedings.mlr.press/v306/ma26ab.html}, abstract = {Federated learning (FL) enables distributed model training, yet in heterogeneous deployments, Bandwidth-Constrained Clients (BCCs) often contribute inefficiently due to limited uplink bandwidth. In model-heterogeneous FL with fixed small sub-models, BCCs with sub-models may improve quickly in early rounds but become under-parameterized later, resulting in slow convergence and poor generalization. To address this challenge, we propose FedGMR, a federated learning framework centered around Gradual Model Restoration (GMR), where GMR progressively increases each client’s sub-model density during training, allowing BCCs to remain effective contributors throughout optimization. To make GMR practical under real-world heterogeneity, FedGMR is realized as an end-to-end workflow with asynchronous coordination and stable, mask-aware aggregation. We further establish convergence guarantees, showing that the aggregation error scales with the average sub-model density across clients and rounds, and that GMR provably narrows the gap toward full-model FL. Extensive experiments on FEMNIST, CIFAR-10, ImageNet-100, and StackOverflow demonstrate that FedGMR improves both convergence speed and final accuracy, especially under severe heterogeneity and non-IID data distributions. Code is available at https://github.com/machengjie321/ICML2026-FedGMR.} }
Endnote
%0 Conference Paper %T Breaking the Capacity Bottleneck in Model-Heterogeneous Federated Learning via Gradual Model Restoration %A Chengjie Ma %A Seungeun Oh %A Jihong Park %A Seong-Lyun Kim %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-ma26ab %I PMLR %P 84807--84845 %U https://proceedings.mlr.press/v306/ma26ab.html %V 306 %X Federated learning (FL) enables distributed model training, yet in heterogeneous deployments, Bandwidth-Constrained Clients (BCCs) often contribute inefficiently due to limited uplink bandwidth. In model-heterogeneous FL with fixed small sub-models, BCCs with sub-models may improve quickly in early rounds but become under-parameterized later, resulting in slow convergence and poor generalization. To address this challenge, we propose FedGMR, a federated learning framework centered around Gradual Model Restoration (GMR), where GMR progressively increases each client’s sub-model density during training, allowing BCCs to remain effective contributors throughout optimization. To make GMR practical under real-world heterogeneity, FedGMR is realized as an end-to-end workflow with asynchronous coordination and stable, mask-aware aggregation. We further establish convergence guarantees, showing that the aggregation error scales with the average sub-model density across clients and rounds, and that GMR provably narrows the gap toward full-model FL. Extensive experiments on FEMNIST, CIFAR-10, ImageNet-100, and StackOverflow demonstrate that FedGMR improves both convergence speed and final accuracy, especially under severe heterogeneity and non-IID data distributions. Code is available at https://github.com/machengjie321/ICML2026-FedGMR.
APA
Ma, C., Oh, S., Park, J. & Kim, S.. (2026). Breaking the Capacity Bottleneck in Model-Heterogeneous Federated Learning via Gradual Model Restoration. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:84807-84845 Available from https://proceedings.mlr.press/v306/ma26ab.html.

Related Material