FedESO: Exploration-to-Stabilization Optimization for Model-Heterogeneous Federated Learning

Model-heterogeneous federated learning (MH-FL) enables resource-constrained clients to train and deploy submodels of different capacities extracted from a shared global model. Achieving strong submodel performance requires both identifying promising parameter subsets and sufficiently optimizing their parameters. However, existing MH-FL methods either provide limited exploration, risking suboptimal selections, or continually vary selected submodels without explicitly regulating their stabilization, hindering consistent optimization. To bridge this gap, we propose FedESO, an exploration-to-stabilization framework that promotes early exploration through Gumbel-perturbed learned submodel selection and progressively anneals the perturbations to stabilize promising submodels for consistent optimization. We formulate dynamic MH-FL as a time-varying optimization problem and establish a nonconvex convergence guarantee governed by the cumulative variation of the induced submodel objectives. We show that sublinear cumulative objective variation is sufficient for asymptotic stationarity and, under finite-time mask stabilization, further establish stationary convergence for the resulting fixed aggregate submodel objective. Our analysis directly characterizes the resource-compatible client submodels relevant to cross-device deployment, rather than relying solely on convergence of the full-model objective. Extensive experiments across vision and natural language processing tasks under diverse resource constraints show that FedESO consistently outperforms state-of-the-art approaches, with gains of at least 10 and 2 percentage points in training-from-scratch and fine-tuning settings, respectively, while also improving optimization efficiency.