FedSPM: Routing-Enabled Federated Learning under Dual Heterogeneity via Semiparametric Mixture

Abstract: Routing-prediction federated learning has emerged as a new paradigm that reframes inter-client heterogeneity as a resource for system-level intelligence: at inference time, the server routes each external query to the best-matched client for prediction. Existing approaches, however, typically treat each client as internally homogeneous, overlooking latent subpopulations within local data. For example, patients with the same diagnosis at one hospital may exhibit morphologically distinct disease subtypes. The coexistence of inter-client and intra-client heterogeneity, which we call dual heterogeneity, can impair both routing and prediction. To address this challenge, we propose FedSPM, a routing-enabled semiparametric mixture framework that represents each client using client-specific latent components. Each component combines a predictive distribution for classification with a feature distribution for routing. To flexibly model feature distributions while effectively sharing information across clients, FedSPM models their density ratios relative to a common nonparametric measure estimated via empirical likelihood. We develop a federated expectation-maximization algorithm that optimizes a tractable surrogate and prove convergence of the exact profiled objective at the standard $\mathcal{O}(1/\sqrt{T})$ rate when the surrogate errors are properly controlled. Experiments on controlled benchmarks and real-world medical data demonstrate consistent improvements in routing and prediction under dual heterogeneity. Code is available at https://github.com/zijianwang0510/FedSPM.

• The paper introduces the FedSPM framework that models dual heterogeneity using a semiparametric mixture, capturing both inter- and intra-client variations.
• It employs a federated EM algorithm with a density ratio model and empirical likelihood to effectively manage covariate and label shifts.
• Empirical evaluations demonstrate significant accuracy gains over existing methods on benchmarks and real-world datasets, including healthcare scenarios.

FedSPM: Semiparametric Mixture Modeling for Routing-Enabled Federated Learning Under Dual Heterogeneity

Motivation and Background

Federated Learning (FL) is a distributed learning paradigm allowing collaboration among clients without centralizing their raw data. Traditionally, FL treats inter-client heterogeneity—as differences in data distributions across clients—as an obstacle, often tackled through global regularization, aggregation reweighting, or optimization refinements. Recently, routing-prediction FL reframed this heterogeneity constructively, enabling server-side query routing to the best-matched client model, thus leveraging client-specialized expertise. However, these approaches typically assume intra-client homogeneity, ignoring latent subpopulations within local data. Real-world scenarios, especially in healthcare, are rife with dual heterogeneity: substantial variation both between and within clients. Existing mixture-model-based personalization methods only partially address intra-client heterogeneity, often assuming identical component distributions across clients and failing to align with both practical and theoretical requirements in complex FL environments.

FedSPM Framework

FedSPM (Federated Semiparametric Mixture) introduces a principled semiparametric approach to model dual heterogeneity in routing-enabled FL. Each client is represented as a mixture of latent components, allowing both intra- and inter-client variation. Each component is defined by (i) a predictive distribution for classification, and (ii) a feature distribution for routing. The predictive distribution parameterizes shared representations with client-specific heads, capturing concept shift. The feature distribution employs a density ratio model (DRM) relative to a shared nonparametric baseline (empirical likelihood, EL), capturing covariate shift. Label shift is addressed through flexible mixing weights.

Notably, FedSPM recasts both predictive and feature modeling using a semiparametric mixture, thus enabling effective cross-client information sharing while maintaining flexibility for client-specific modeling. The baseline distribution $G$ in the DRM is estimated nonparametrically via EL over pooled samples, providing both theoretical guarantees and practical effectiveness without parametric restrictions.

Federated EM Algorithm and Optimization Guarantees

FedSPM employs a federated expectation-maximization (EM) algorithm tailored for the mixture model under EL constraints. The server broadcasts common parameters, and each client performs local E-steps and M-steps: responsibility-based computation of component memberships, followed by parameter updates for mixing weights, supervised block (shared representation and client-specific heads), and DRM block (feature distribution). A critical innovation is that summaries required for cross-client dependence (component responsibilities) are efficiently communicated as aggregate statistics ($m \times C$ scalars) to the server, maintaining privacy and communication efficiency.

For optimization, FedSPM proves convergence of local SGD with momentum for the DRM block under mild regularity assumptions. The system attains the standard $O(1/\sqrt{T})$ convergence rate, up to negligible surrogate errors from profile likelihood optimization. Precise analysis of the momentum parameter reveals a trade-off: as momentum increases, initial gap terms improve but stochastic noise is amplified. Empirical experiments validate convergence and control of surrogate error terms.

Empirical Evaluation

FedSPM is extensively evaluated across synthetic and real-world datasets. Controlled experiments on FMNIST, CIFAR-10, and CIFAR-100 simulate dual heterogeneity through structured label and covariate shifts using Dirichlet-based partitioning and image transformations. Baselines encompass global-model FL (FedAvg, FedProx), personalization (fine-tuning, Ditto, ClusterFL, FedBABU), mixture models (FedEM, FedGMM), and routing-enabled methods (FedDRM).

FedSPM consistently achieves the highest system and average accuracies across all benchmarks, regardless of model capacity scaling, and is robust with respect to heterogeneity intensity. Increasing the number of mixture components confers additional gains, especially when encoding over-specification for predictive flexibility. FedSPM demonstrates stark improvements over mixture-model-based methods, indicating that client-specific component modeling is critical for handling concept shift.

On Fed-ISIC2019, a real medical dataset from 6 clinical centers with naturally arising dual heterogeneity, FedSPM surpasses all baselines in system, average, and routing accuracy. Mixture modeling is empirically necessary: mixture-method baselines outperform those that fail to account for intra-client heterogeneity, and FedSPM's semiparametric approach outperforms parametric alternatives (FedGMM), validating theoretical claims and practical efficacy in complex real-world FL settings.

Implications and Future Directions

FedSPM advances the routing-prediction FL paradigm by providing a unified semiparametric mixture framework tailored to dual heterogeneity. The integration of client-specific mixture components with flexible EL-based DRM delivers accurate and robust routing and prediction, transforming heterogeneity from an impediment to a resource for system-level intelligence. This framework is especially pertinent for health informatics and domains with naturally hierarchical population structures, and enables explicit query routing with expertise-aware prediction.

Further theoretical developments may explore identifiability and interpretability of learned latent structures, generalization bounds for semiparametric mixture models in FL, and privacy-preserving variants of EL estimation. Practical enhancements may include adaptive component selection, hierarchical mixture modeling, and integration with multi-modal data. Future work will likely investigate scaling FedSPM to larger FL networks, more complex tasks, and time-evolving heterogeneity.

Conclusion

FedSPM introduces a semiparametric mixture model and federated EM methodology to routing-enabled FL under dual heterogeneity. By jointly modeling intra- and inter-client variation using a flexible mixture and DRM-based empirical likelihood, FedSPM achieves both strong empirical performance and theoretical guarantees. The paradigm enables adaptive query routing and robust client prediction, offering a scalable solution for expertise-aware FL systems across heterogeneous, distributed domains (2607.04085).