Breaking the Capacity Bottleneck in Model-Heterogeneous Federated Learning via Gradual Model Restoration
Chengjie Ma, Seungeun Oh, Jihong Park, Seong-Lyun Kim
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.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 64f278fe-605e-4b78-b05c-155c885cc788Builds on12
- Ensemble Distillation for Robust Model Fusion in Federated LearningTao Lin, Lingjing Kong, Sebastian U. Stich, Martin JaggiNeurIPS 2020 · 1,615 citations
- Once-for-All: Train One Network and Specialize it for Efficient DeploymentHan Cai, Chuang Gan, Tianzhe Wang, Zhekai Zhang et al.ICLR 2020 · 1,522 citations
- Universally Slimmable Networks and Improved Training TechniquesJiahui Yu, Thomas S. HuangICCV 2019 · 444 citations
- FjORD: Fair and Accurate Federated Learning under heterogeneous targets with Ordered DropoutSamuel Horváth, Stefanos Laskaridis, Mário Almeida, Ilias Leontiadis et al.NeurIPS 2021 · 390 citations
- Training Neural Networks with Fixed Sparse MasksYi-Lin Sung, Varun Nair, Colin RaffelNeurIPS 2021 · 295 citations
Related papers
- FedLMT: Tackling System Heterogeneity of Federated Learning via Low-Rank Model Training with Theoretical GuaranteesJiahao Liu, Yipeng Zhou, Di Wu, Miao Hu et al.ICML 2024 · 8 citations
- Heterogeneity-Aware Federated Learning with Adaptive Client Selection and Gradient CompressionZhida Jiang, Yang Xu, Hongli Xu, Zhiyuan Wang et al.INFOCOM 2023 · 43 citations
- FedEcover: Fast and Stable Converging Model-Heterogeneous Federated Learning with Efficient-Coverage Submodel ExtractionJuntao Liang, Lan Zhang, Xiangmou Qu, Jun WangICDE 2025
- FedProto: Federated Prototype Learning across Heterogeneous ClientsYue Tan, Guodong Long, Lu Liu, Tianyi Zhou et al.AAAI 2022 · 851 citations
- FedLAGC: Towards High Performance System-Heterogeneous Federated Learning via Layer-Adaptive Submodel Extraction and Gradient CorrectionQing Hu, Tianchi Liao, Shuyi Wu, Lei Yang et al.AAAI 2026
