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

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.

Federated learning allows many devices to train a shared AI model without sending their private data to a central server. In practice, however, devices are often very different from one another. Some have fast network connections and enough resources to train large models, while others are slower and can only handle much smaller models. A common solution is to give weaker devices smaller models so they can still participate. This works well at the beginning of training, but later on these small models often become too limited to provide useful updates, which slows overall progress. Our paper studies this overlooked problem and proposes a simple solution: instead of keeping weaker devices on small models for the entire training process, we gradually restore their model capacity over time. In this way, these devices can train efficiently early on and still make meaningful contributions later. We also design a training and aggregation strategy that keeps this gradual restoration stable in realistic heterogeneous settings. We support this idea through both theory and experiments. Across several image and language tasks, our method improves both training speed and final accuracy, especially when devices and data are highly heterogeneous. These results suggest that in real-world federated learning, it is not enough to make weaker devices participate quickly at the start; it is also important to help them remain useful throughout training.