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

TL;DR

FedGMR introduces a gradual model restoration approach in federated learning to improve convergence and accuracy in heterogeneous, bandwidth-constrained environments by progressively increasing client sub-model density.

Contribution

The paper proposes FedGMR, a novel federated learning framework with a gradual model restoration technique that adapts sub-model sizes during training to handle heterogeneity effectively.

Findings

FedGMR accelerates convergence compared to fixed sub-model approaches.

It achieves higher final accuracy on multiple datasets under heterogeneity.

Theoretical analysis confirms convergence guarantees and error bounds.

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 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.