[News] ICML 2026 Acceptance: Breaking the Capacity Bottleneck in Model-Heterogeneous Federated Learning via Gradual Model Restoration

RAMO Lab is pleased to announce that our paper, “Breaking the Capacity Bottleneck in Model-Heterogeneous Federated Learning via Gradual Model Restoration,” has been accepted to ICML 2026.

This work addresses a fundamental challenge in model-heterogeneous federated learning (MHFL). In heterogeneous client environments, bandwidth-constrained clients are typically assigned smaller sub-models so that they can participate in training with lower communication and computation costs. However, our study shows that while these small sub-models are effective in the early stage of training, they gradually become capacity-limited later on and lose their ability to make meaningful contributions to the global model.

To address this issue, we propose FedGMR, a federated learning framework centered on Gradual Model Restoration (GMR). FedGMR progressively restores the model density of bandwidth-constrained clients during training, allowing them to preserve early-stage efficiency while regaining late-stage effectiveness. To support this restoration process, the framework further integrates asynchronous coordination and mask-aware aggregation to maintain training stability under heterogeneous conditions.

From a theoretical perspective, the paper analyzes the bias and variance introduced by fixed low-density sub-model training in MHFL, highlights the importance of average model density and coverage in convergence behavior, and shows that gradual restoration can continuously narrow the optimization gap toward full-model federated learning. The paper also examines aggregation strategies during restoration and shows that mask-aware aggregation better preserves gradient stability, making it more compatible with GMR than naive alternatives.

Experimentally, the proposed method is validated on FEMNIST, CIFAR-10, ImageNet-100, and StackOverflow, covering both vision and language tasks. FedGMR consistently improves over multiple MHFL baselines, with especially clear gains under highly heterogeneous and non-IID settings. Ablation studies further confirm the central intuition of the paper: smaller sub-models are more suitable in the early stage, whereas larger-capacity models become increasingly important later in training. Additional cross-method experiments show that the benefit of restoration is not tied to one specific pruning rule or one specific MHFL implementation, but reflects a more general mechanism for re-activating weak clients in late-stage training.

Overall, this work suggests that the goal in model-heterogeneous federated learning is not only to help weak clients train faster with smaller models early on, but also to restore their model capacity at the right time so that they remain effective contributors throughout global optimization.

Links

- arXiv: https://arxiv.org/pdf/2512.05372v2

- Code: https://github.com/machengjie321/ICML2026-Breaking-the-Capacity-Bottleneck-in-Model-Heterogeneous-FL-via-Gradual-Model-Restoration