1 Introduction↩︎

Federated learning (FL) [1] enables multiple clients to collaboratively train a model without centralizing their raw data. In a typical FL system, a server distributes a shared model to participating clients, each client updates the model using its local data, and the server aggregates the resulting locally updated models into an improved global model. This distributed training paradigm allows knowledge to be shared across clients such as hospitals [2] and mobile devices [3], making FL particularly promising when data are sensitive, geographically dispersed, or impractical to centralize [4]–[6].

Traditionally, inter-client heterogeneity, where data distributions vary across clients, is viewed as an obstacle in FL. Since clients optimize different local objectives, their gradients may drift in conflicting directions [7], thereby slowing convergence [8] and degrading global model performance [9]. In contrast, the recent routing-prediction FL paradigm [10] reframes inter-client heterogeneity as a useful signal of client specialization. Like personalized FL [11]–[13], it learns specialized models for individual client domains. Beyond personalization, it further estimates how well an external query matches each client’s data distribution. At inference time, the server uses these distributional match scores to route the query to the most suitable client, whose specialized model then makes the final prediction, thereby turning client-specific expertise into system-level intelligence.

However, such a routing-prediction FL paradigm typically assumes that the data within each client are drawn from a homogeneous distribution, thereby overlooking intra-client heterogeneity. This assumption is often violated in practice, as local data may arise from a mixture of latent components [14], [15]. For example, within one hospital, dermoscopic images may involve different lesion types, anatomical sites, and patient age groups: cases with the same diagnosis may exhibit markedly different visual patterns [16], [17], whereas visually similar cases may correspond to different diagnoses [18], [19]. Since these factors and their interactions are rarely fully observed or annotated, the resulting subgroup memberships are latent, making it infeasible to fit a separate model to each predefined group. Ignoring such latent structure forces a homogeneous local model to fit a mixture of heterogeneous feature and predictive distributions, thereby degrading both routing and prediction accuracy. Taken together, intra-client and inter-client heterogeneity constitute what we call dual heterogeneity. This raises our key question: how can routing and prediction be jointly improved under dual heterogeneity?

No existing heterogeneous FL approach fully addresses this problem. Methods centered on a global model mitigate client drift through regularization [20], [21], aggregation reweighting [22], [23], or refined optimization [24], [25], but their shared predictor still struggles to adapt to client-specific distributions. Conventional personalized methods, including local fine-tuning [26], regularization [11], [27], client clustering [12], [28], and representation learning [13], [29], improve prediction within each client’s local domain, but provide no mechanism for server-side routing of external queries. Routing-based personalization [10] enables server-side routing, yet treats each client as internally homogeneous and therefore overlooks intra-client heterogeneity. Mixture-model-based personalization captures intra-client heterogeneity [14], [15], but its assumption of shared component distributions across clients limits flexibility in modeling client-specific latent structures.

To address this challenge, we propose FedSPM(Fig. 1), a routing-enabled FL framework that represents each client as a mixture of latent components rather than a single homogeneous population. This mixture discovers hidden variation within each client, while allowing its components to differ across clients captures inter-client distribution shifts. Each component contains two complementary parts: a predictive distribution that relates features to labels and a feature distribution that characterizes the inputs covered by the component. To balance model flexibility with effective information sharing, the predictive distributions combine shared representations with client-specific prediction heads, while the feature distributions use a density ratio model (DRM) [30] relative to a shared nonparametric baseline. We learn the model using empirical likelihood (EL) [31] and a federated expectation-maximization (EM) algorithm. At inference time, the server routes an external query to the most suitable client using the learned feature distributions, and the selected client combines its component-wise predictions to produce the final output. Our main contributions are summarized as follows:

• We formulate routing-prediction FL under dual heterogeneity and propose FedSPM, a semiparametric mixture framework that flexibly captures latent variation within clients and distribution shifts across clients while enabling effective information sharing.

• We develop a federated EM algorithm based on EL and establish convergence guarantees for the resulting nonconvex optimization under local stochastic gradient descent (SGD) with momentum.

• We evaluate FedSPM on controlled benchmarks and a real-world medical dataset, demonstrating consistent improvements in both routing and prediction over competitive FL baselines.

2 Method↩︎

Problem Formulation. Consider an FL system with \(m\) clients for \(K\)-class classification. Each client \(i \in [m]\mathrel{\vcenter{:}}=\{1,\dots, m\}\) has a local dataset \(\mathcal{D}_i\mathrel{\vcenter{:}}=\{(x_{ij},y_{ij})\}_{j\in [n_i]}\) with samples drawn independently from a client-specific distribution \(P_i^{X,Y}\). Let \(n\mathrel{\vcenter{:}}= \sum_{i=1}^m n_i\) be the total sample size, \(\rho_i\mathrel{\vcenter{:}}= n_i/n\) be the sample fraction of client \(i\), and \(\mathcal{D}\mathrel{\vcenter{:}}=\bigcup_{i=1}^m\mathcal{D}_i\) be the pooled dataset.

2.1 Semiparametric Mixture Model↩︎

To capture intra-client heterogeneity, we introduce a latent component variable \(Z\in[C]\) and model each observed distribution \(P_i^{X,Y}\) through an augmented distribution \(P_i^{X,Y,Z}\). Specifically, we assume \[\label{eq:mixture} (x_{ij},y_{ij})\mid(z_{ij}=c)\sim P_{ic}^{X,Y},~z_{ij}\sim P_i^Z \text{ with } P_i^Z(\{c\})=\pi_{ic},~\forall c\in[C],~j\in[n_i].\tag{1}\] Here, \(\pi_i\mathrel{\vcenter{:}}=(\pi_{i1},\dots,\pi_{iC})\) with \(\pi_{ic}\ge0\) and \(\sum_c\pi_{ic}=1\). Marginalizing over \(Z\) gives \(P_i^{X,Y}=\sum_c\pi_{ic}P_{ic}^{X,Y}\). Notably, we allow both the mixing weights \(\pi_i\) and the component distributions \(P_{ic}^{X,Y}\) to vary across clients, providing a flexible model of intra-client heterogeneity.

While allowing fully client-specific \(P_{ic}^{X,Y}\) is expressive, it prevents effective information sharing across clients. We factor \(P_{ic}^{X,Y}\) into its predictive distribution \(P_{ic}^{Y|X}\) and feature distribution \(P_{ic}^X\), and impose structure on both factors to share information while modeling inter-client heterogeneity:

Concept shift. We model the component predictive distribution of client \(i\) as \[\label{eq:concept95shift} P_{ic}^{Y|X}(\{k\}\mid x)\propto\exp(\alpha_{ikc}+\beta_{ikc}^\top g_{\theta_c}(x)),\tag{2}\] where \(g_{\theta_c}(x)\) is a shared embedding for component \(c\). The client-specific parameters \((\alpha_{ikc}, \beta_{ikc})\) allow the marginal predictive distribution \(P_i^{Y|X}\) to vary across clients, hence capturing concept shift.

Covariate shift. We model the component feature distribution of client \(i\) via a DRM with respect to a shared baseline distribution \(G\): \[\label{eq:covariate95shift} dP_{ic}^{X}/dG(x) = \exp\big(\gamma_{ic} + \xi_{ic}^\top h_{\nu_c}(x)\big),\tag{3}\] where \(dP_{ic}^{X}/dG\) is the Radon–Nikodym derivative of \(P_{ic}^X\) relative to \(G\), and \(h_{\nu_c}\) is a shared basis for component \(c\). The client-specific tilting parameters \((\gamma_{ic}, \xi_{ic})\) induce differences in the marginal feature distribution \(P_i^X\) across clients, capturing covariate shift while preserving a shared structure.

Label shift. Finally, differences in the mixing weights \(\pi_i\) and the component distribution \(P_{ic}^{X,Y}\) jointly induce variation in the marginal label distribution \(P_i^Y\) across clients, capturing label shift.

The baseline distribution \(G\) in 3 remains unspecified. Restricting \(G\) to a parametric family imposes an unjustified distributional assumption, which limits model flexibility and increases the risk of misspecification. Instead, we estimate \(G\) via EL, a nonparametric likelihood framework that assigns unknown probability masses to the observed samples, subject to the DRM constraints (see App. 6 for related work on DRM and EL). Specifically, \(G\) is represented as a discrete distribution supported on the pooled observations, taking the form \(G=\sum_{i,j}r_{ij}\delta_{x_{ij}}\) with unknown \(r_{ij} \ge 0\). Crucially, all samples across all clients jointly determine \(G\), serving as a key mechanism for information sharing. To ensure that \(G\) and each \(P_{ic}^X\) define valid probability distributions, the masses \(\{r_{ij}\}\) must satisfy \(\int G(dx)=1\) and \(\int \exp\left(\gamma_{i'c}+\xi_{i'c}^\top h_{\nu_c}(x)\right) G(dx)=1\) for all \(i',c\), i.e.,  \[\label{eq:constraints} \sum_{i,j} r_{ij} = 1, \quad \sum_{i,j} \exp\left(\gamma_{i'c} + \xi_{i'c}^\top h_{\nu_c}(x_{ij})\right) r_{ij} = 1, \quad \forall i' \in [m], c \in [C].\tag{4}\] In summary, the predictive model 2 , the DRM-EL feature model 3 , and the mixing weights \(\pi_i\) together define our semiparametric mixture model.

2.2 Profile Log-EL and EM Algorithm↩︎

We develop a practical learning procedure for the proposed semiparametric mixture model. The main difficulty is that both the latent component assignments and the nonparametric baseline distribution \(G\) are unknown. We first profile out \(G\) through EL, and then exploit the latent mixture structure to derive an EM algorithm that avoids direct optimization of the resulting intractable objective.

Profile log-EL. Let \(\zeta\mathrel{\vcenter{:}}=(\alpha,\beta,\theta,\gamma,\xi,\nu,\pi)\) be the parameters of interest. The log-EL based on \(\mathcal{D}\) is: \[\ell(\zeta, G) = \sum_{i,j} \log r_{ij} + \sum_{i,j} \log \biggl[ \sum_{c} \pi_{ic} \exp \left( \gamma_{ic} + \xi_{ic}^\top h_{\nu_c}(x_{ij}) \right) P_{ic}^{Y|X}(\{y_{ij}\}\mid x_{ij}) \biggr].\] See App. 8.1 for the derivation. Since \(G\) is a nuisance parameter, we work with the profile log-EL \(p\ell(\zeta)\mathrel{\vcenter{:}}=\sup_G \ell(\zeta, G)\), where the supremum is taken under the constraints in 4 . By the Lagrange multiplier method, the optimal baseline weights are: \[\label{eq:optimal95G} r_{ij}^*(\zeta)=n^{-1}\biggl\{1+\sum_{i',c}\lambda_{i'c}\left[\exp\left(\gamma_{i'c}+\xi_{i'c}^\top h_{\nu_c}(x_{ij})\right)-1\right] \biggr\}^{-1},\tag{5}\] where \(\{\lambda_{ic}\}\) are the solutions to \[\label{eq:Lagrange95multipliers95system} \sum_{i,j}\frac{\exp\left(\gamma_{i'c}+\xi_{i'c}^\top h_{\nu_c}(x_{ij})\right)-1}{1+\sum_{i'',c'}\lambda_{i''c'}\left[\exp\left(\gamma_{i''c'}+\xi_{i''c'}^\top h_{\nu_{c'}}(x_{ij})\right)-1\right]}=0,\quad\forall i'\in [m], c\in [C].\tag{6}\] See App. 8.2 for the derivation. Substituting 5 back into \(\ell(\zeta,G)\) yields \(p\ell(\zeta)\). However, directly maximizing \(p\ell(\zeta)\) remains computationally impractical because each evaluation requires solving the nonlinear system in 6 for the Lagrange multipliers \(\{\lambda_{ic}\}\), making end-to-end optimization with standard automatic-differentiation tools such as PyTorch infeasible.

EM algorithm. To address this challenge, we treat the component assignments \(\{z_{ij}\}\) as missing data and employ the EM algorithm to maximize \(p\ell(\zeta)\). Let \(z_{ijc}\mathrel{\vcenter{:}}=\mathbb{1}(z_{ij}=c)\) denote the indicator that the \(j\)-th sample on the \(i\)-th client belongs to the \(c\)-th subpopulation. If \(\{z_{ij}\}\) were observed, the complete-data profile log-EL would be \[p\ell^{\text{c}}(\zeta)=\sum_{i,j}\log r_{ij}^*(\zeta)+\sum_{i,j,c}z_{ijc}\biggl[\log P_{ic}^{Y|X}(\{y_{ij}\}\mid x_{ij})+\gamma_{ic}+\xi_{ic}^\top h_{\nu_c}(x_{ij})+\log \pi_{ic} \biggr].\] See App. 8.3 for derivation. Since \(\{z_{ij}\}\) are unobserved, EM alternates between estimating their posterior distributions under the current parameters and updating the parameters based on these estimates. Specifically, the two steps are:

E-step. Given \(\zeta^{(t)}\), we compute the posterior responsibilities for each \(z_{ij}\) as \[\label{eq:posterior} w_{ijc}^{(t)} \mathrel{\vcenter{:}}= \mathbb{E}\left[z_{ijc}\middle| \mathcal{D};\zeta^{(t)}\right] \propto P_{ic}^{Y|X}(\{y_{ij}\}\mid x_{ij};\zeta^{(t)})\exp\left(\gamma_{ic}^{(t)}+(\xi_{ic}^{(t)})^\top h_{\nu_c^{(t)}}(x_{ij})\right)\pi_{ic}^{(t)}.\tag{7}\] The resulting profile \(Q\)-function is the conditional expectation of the complete-data profile log-EL: \[\label{eq:Q95decomposition} Q^{(t)}(\zeta)\mathrel{\vcenter{:}}=\mathbb{E}\left[p\ell^{\text{c}}(\zeta)\middle|\mathcal{D};\zeta^{(t)}\right] = Q^{(t)}_1(\alpha, \beta, \theta) + Q^{(t)}_2(\gamma, \xi, \nu) + Q^{(t)}_3(\pi),\tag{8}\] where \[Q_1^{(t)}(\alpha,\beta,\theta) \mathrel{\vcenter{:}}= \sum_{i,j,c}w_{ijc}^{(t)}\log P_{ic}^{Y|X}(\{y_{ij}\}\mid x_{ij}), \quad Q_3^{(t)}(\pi) \mathrel{\vcenter{:}}= \sum_{i,j,c}w_{ijc}^{(t)}\log\pi_{ic}, \text{ and }\] \[{!}{ \displaystyle Q_2^{(t)}(\gamma, \xi, \nu) \mathrel{\vcenter{:}}= \sum_{i,j,c}w_{ijc}^{(t)} \left[\gamma_{ic} + \xi_{ic}^\top h_{\nu_c}(x_{ij})\right] - \sum_{i,j}\log \biggl(1 + \sum_{i',c} \lambda_{i'c} \left[\exp\left(\gamma_{i'c} + \xi_{i'c}^\top h_{\nu_c}(x_{ij})\right) - 1\right]\biggr). }\]

M-step. We update \(\zeta^{(t+1)}=\arg\max Q^{(t)}(\zeta)\). Since \(Q_1^{(t)}\), \(Q_2^{(t)}\), and \(Q_3^{(t)}\) are separable, the three parameter blocks can be optimized independently. For convenience, we refer to \((\alpha,\beta,\theta)\) as the supervised block, and to \((\gamma,\xi,\nu)\) as the DRM block. For the mixing weights, maximizing \(Q^{(t)}_3\) under \(\sum_c \pi_{ic} = 1\) yields \[\label{eq:update95pi} \pi_{ic}^{(t+1)}=\frac{1}{n_i}\sum_{j=1}^{n_i}w_{ijc}^{(t)}.\tag{9}\] See App. 8.4 for derivation details of the EM algorithm. For the supervised block, maximizing \(Q_1^{(t)}\) is equivalent to minimizing a weighted cross-entropy loss, which can be directly solved via standard gradient-based methods: \(\left(\alpha^{(t+1)},\beta^{(t+1)},\theta^{(t+1)}\right)=\arg\max Q^{(t)}_1(\alpha,\beta,\theta)\). For the DRM block, directly maximizing \(Q_2^{(t)}\) faces the same computational bottleneck as \(p\ell(\zeta)\): each evaluation still requires solving the nonlinear system for the Lagrange multipliers \(\{\lambda_{ic}\}\). In contrast, the EM algorithm offers a pivotal resolution to this bottleneck: at any critical point of \(Q^{(t)}_2\), each \(\lambda_{ic}\) admits the analytic form \(n^{-1}\sum_{j=1}^{n_i}w_{ijc}^{(t)}\) (see App. 8.5). Substituting this expression back shows that every critical point of \(Q^{(t)}_2\) is also that of the tractable surrogate \[{!}{ \displaystyle\widetilde{Q}_2^{(t)}(\gamma,\xi,\nu)\mathrel{\vcenter{:}}= \sum_{i,j,c}w_{ijc}^{(t)}\left[\gamma_{ic}+\xi_{ic}^\top h_{\nu_c}(x_{ij})\right]-\sum_{i,j}\log\biggl(\sum_{i',j',c}w_{i'j'c}^{(t)}\exp\left(\gamma_{i'c}+\xi_{i'c}^\top h_{\nu_c}(x_{ij})\right)\biggr).}\] We optimize \(\widetilde{Q}^{(t)}_2\) instead, giving the update \(\left(\gamma^{(t+1)},\xi^{(t+1)},\nu^{(t+1)}\right)=\arg\max\widetilde{Q}_2^{(t)}(\gamma,\xi,\nu)\).

2.3 Federated Training and Inference↩︎

Algorithm [algo:fedspm] in App. 7 presents FedSPM, which extends the centralized EM procedure to the client-server setting. As in standard FL, FedSPM follows a local-update-then-aggregate paradigm. In communication round \(t\), the server broadcasts the latest shared parameters \((\theta^{(t)},\gamma^{(t)},\xi^{(t)},\nu^{(t)})\) to all clients. Each client \(i\) then performs a local E-step followed by a local M-step:

Local E-step. Client \(i\) computes \(\{w_{ijc}^{(t)}\}_{j\in[n_i],c\in[C]}\) via 7 using only its own dataset \(\mathcal{D}_i\).

Local M-step. After the local E-step, client \(i\) updates the three parameter blocks separately. First, since the mixing weights \(\pi_i\) are client-specific, client \(i\) updates them locally via 9 without server aggregation. For the supervised block, client \(i\) updates \((\alpha_i^{(t+1)},\beta_i^{(t+1)},\theta_i^{(t+1)})\) by maximizing \[Q_{1i}^{(t)}(\alpha, \beta, \theta) \mathrel{\vcenter{:}}= \sum_{j,c} w_{ijc}^{(t)} \log P_{ic}^{Y|X}(\{y_{ij}\}\mid x_{ij}),\] which depends only on \(\mathcal{D}_i\) and \(\{w_{ijc}^{(t)}\}_{j\in[n_i],c\in[C]}\) available on client \(i\). For the DRM block, client \(i\) updates \((\gamma_i^{(t+1)}, \xi_i^{(t+1)}, \nu_i^{(t+1)})\) by maximizing \[\widetilde{Q}_{2i}^{(t)}(\gamma,\xi,\nu)\mathrel{\vcenter{:}}= \sum_{j,c}w_{ijc}^{(t)}\left[\gamma_{ic}+\xi_{ic}^\top h_{\nu_c}(x_{ij})\right]-\sum_{j}\log\biggl(\sum_{i',c}\tau_{i'c}^{(t)}\exp\left(\gamma_{i'c}+\xi_{i'c}^\top h_{\nu_c}(x_{ij})\right)\biggr),\] where \(\tau_{i'c}^{(t)}\mathrel{\vcenter{:}}=\sum_{j'=1}^{n_{i'}}w_{i'j'c}^{(t)}\) denotes the total responsibility assigned to component \(c\) on client \(i'\). Unlike \(Q_{1i}^{(t)}\), client \(i\) cannot evaluate \(\widetilde{Q}_{2i}^{(t)}\) using only its local information, because \(\widetilde{Q}_{2i}^{(t)}\) involves responsibilities from other clients. However, this cross-client dependence enters only through the summary statistics \(\{\sum_{j'=1}^{n_{i'}}w_{i'j'c}^{(t)}\}_{i'\in[m],c\in[C]}\). Accordingly, after the local E-step, each client \(i\) transmits \(\{\tau_{ic}^{(t)}\}_{c\in[C]}\) to the server, and the server broadcasts the collection \(\{\tau_{ic}^{(t)}\}_{i\in[m],c\in[C]}\) to all clients. This additional communication consists of only \(m\times C\) scalars, which is negligible relative to model transmission and does not reveal raw data or per-sample responsibilities. Armed with these summary statistics, client \(i\) can maximize both \(Q_{1i}^{(t)}\) and \(\widetilde{Q}_{2i}^{(t)}\) via any gradient-based LocalSolver.

Routing and prediction. At inference time, given an external query \(x\), the server first routes it to the most suitable client using maximum a posteriori estimation, with the sample fraction \(\rho_i\) as the prior and the client feature density relative to the baseline distribution \(G\) as the likelihood: \[\label{eq:client95routing} i^*\mathrel{\vcenter{:}}= \arg\max_{i\in [m]}\left\{\rho_i \cdot dP_i^{X}/dG(x)\right\} = \arg\max_{i\in [m]}\biggl\{\rho_i\sum_{c}\pi_{ic}\exp\big(\gamma_{ic} + \xi_{ic}^\top h_{\nu_c}(x)\big)\biggr\}.\tag{10}\] Let \(P_i^{Z|X}\) denote the conditional distribution of the latent component given the query on client \(i\). The selected client \(i^*\) then makes the final prediction: \[\label{eq:local95prediction} \widehat{y}\mathrel{\vcenter{:}}= \arg\max_{k\in [K]} P_{i^*}^{Y|X}(\{k\}\mid x)= \arg\max_{k\in [K]}\biggl\{\sum_{c} P_{i^*c}^{Y|X}(\{k\}\mid x)P_{i^*}^{Z|X}(\{c\}\mid x)\biggr\}.\tag{11}\]

2.4 Convergence Analysis↩︎

We analyze the convergence of FedSPM when LocalSolver is instantiated as local SGD [32] with momentum. Since the mixing weights and the supervised block follow standard EM/GEM updates and inherit the convergence guarantees established by [14], we focus exclusively on the DRM block \(\phi\mathrel{\vcenter{:}}=(\gamma,\xi,\nu)\). For convenience, we cast the optimization of \(\phi\) as a standard minimization problem: \[F(\zeta)\mathrel{\vcenter{:}}= -p\ell(\zeta),\quad f^{(t)}(\phi)\mathrel{\vcenter{:}}= -Q_2^{(t)}(\phi),\quad \widetilde{f}^{(t)}(\phi)\mathrel{\vcenter{:}}= -\widetilde{Q}_2^{(t)}(\phi),\quad \widetilde{f}_i^{(t)}(\phi)\mathrel{\vcenter{:}}= -\widetilde{Q}_{2i}^{(t)}(\phi).\] Assume that \(F\) is bounded below by \(F^*\). In communication round \(t\), each client \(i\) performs \[\phi_{i,0}^{(t)}\mathrel{\vcenter{:}}= \phi^{(t)},\qquad d_{i,-1}^{(t)}\mathrel{\vcenter{:}}= 0,\qquad d_{i,e}^{(t)}\mathrel{\vcenter{:}}= \mu d_{i,e-1}^{(t)}+\widetilde{g}_{i,e}^{(t)},\qquad \phi_{i,e+1}^{(t)}\mathrel{\vcenter{:}}= \phi_{i,e}^{(t)}-\eta d_{i,e}^{(t)},\] for \(e\in\{0,\ldots,E-1\}\), where \(\mu\in [0,1)\) denotes the momentum, \(E\) denotes the number of local steps, \(\eta\) denotes the learning rate, and \(\widetilde{g}_{i,e}^{(t)}\) denotes the stochastic gradient.

Assumptions 1–3 are standard in federated optimization [14], [22].

Assumption 4 requires that, along the iterates, the gradient norm of the DRM objective \(f^{(t)}\) is bounded by that of its tractable surrogate \(\widetilde{f}^{(t)}\) up to an error \(\varepsilon^{(t)}\), without imposing any alignment between their gradient directions, while Assumption 5 quantifies the one-step mismatch between their function-value changes. We further provide empirical support for these assumptions in Sec. 3.2.

Crucially, although FedSPM optimizes the tractable surrogate \(\widetilde{f}^{(t)}\), Theorem 1 establishes convergence toward a stationary point of the exact profiled objective \(F\). For fixed \(E\), if \(\sum_{t=0}^{T-1}\delta^{(t)}=\mathcal{O}(1)\) and \(\sum_{t=0}^{T-1}\varepsilon^{(t)}=\mathcal{O}(\sqrt{T})\), the surrogate errors do not change the standard \(\mathcal{O}(1/\sqrt{T})\) convergence rate. The bound further characterizes the effect of momentum through the accumulation factors \(a_E\) and \(s_E\). When \(\mu=0\), \(a_E=s_E=E\), recovering the usual local SGD scaling. As \(\mu\to 1\), \(a_E=\mathcal{O}(E^2)\) and \(s_E=\mathcal{O}(E^3)\). Consequently, the initial-gap term improves from order \(1/(E\sqrt{T})\) to \(1/(E^2\sqrt{T})\), whereas the stochastic term grows from order \(1/\sqrt{T}\) to \(E/\sqrt{T}\), revealing a trade-off between faster optimization and amplified stochastic noise. The proof is deferred to App. 8.6.

3 Experiments on Benchmark Datasets↩︎

We conduct experiments on three image classification benchmarks of increasing complexity: FMNIST [33], CIFAR-10 [34], and CIFAR-100 [34], with dataset details deferred to App. 9.1. To study each heterogeneity type in isolation and assess their joint effect, we construct semi-synthetic FL settings by superimposing controlled inter-client and intra-client heterogeneity onto the original data.

3.1 Training and Evaluation Setup↩︎

Baselines. We compare FedSPM with representative FL baselines from seven categories: global-model methods (FedAvg [1] and FedProx [20]), a fine-tuning-based method (FedAvgFT [26]), a regularization-based method (Ditto [11]), a cluster-based method (ClusterFL [12]), a representation-learning-based method (FedBABU [29]), mixture-model-based methods (FedEM [14] and FedGMM [15]), and a routing-based method (FedDRM [10]). Notably, only FedGMM, FedDRM, and FedSPM are capable of client routing by design.

Evaluation metrics. Following [10], we evaluate all methods using system and average accuracies. System accuracy is measured on the pooled test set across all clients, computed via 10 and 11 for routing-capable methods, and via majority voting otherwise. Average accuracy is the mean of local test accuracies weighted by client sample fractions.

Model architecture. We instantiate the classification encoder \(g_{\theta_c}\) with ResNet [35] and the routing encoder \(h_{\nu_c}\) with a lightweight CNN [36]. For a fair comparison, all methods use the same classification encoder architecture, and routing-capable methods use the same routing encoder architecture. Notably, the model capacity of cluster- and mixture-model-based methods scales linearly with their number of components. For FedSPM, we report two variants to separate algorithmic gains from increased model capacity. In FedSPM (\(1\times\)), the encoders are shared across components, i.e., \(g_{\theta_c}\equiv g_{\theta}\) and \(h_{\nu_c}\equiv h_{\nu}\) for all \(c\). In FedSPM (\(C\times\)), each component has its own encoders \(g_{\theta_c}\) and \(h_{\nu_c}\).

Dual heterogeneity settings. Since benchmark datasets do not inherently exhibit statistical heterogeneity, we explicitly construct dual heterogeneity based on the mixture representation \(P_i^{X,Y}=\sum_c \pi_{ic}P_{ic}^{X,Y}\), following common practice [10], [15], [37]. Specifically, we first partition the full dataset across 8 clients via class-wise Dirichlet partitioning [38] with concentration parameter \(\alpha_{\text{inter}}=1.0\), inducing label shift and unequal client dataset sizes. Within each client, we further divide the local dataset into 2 latent components, with the mixing weights \(\pi_i\) sampled from a Dirichlet distribution with \(\alpha_{\text{intra}}=2.0\). To induce component-wise covariate shift in \(P_{ic}^{X}\), we apply transformations at two levels: client-level transformations, which introduce inter-client heterogeneity through combinations of color shifts and spatial intensity biases [10]; and component-level transformations, which introduce intra-client heterogeneity through color-channel inversion [15]. In addition, for each client-component pair, we randomly generate a label permutation [15] to induce component-wise concept shift in \(P_{ic}^{Y|X}\). Since the type and degree of each heterogeneity are specified by construction, the ground-truth data-generating mechanism is fully known, making these benchmarks well-suited for controlled comparisons. See App. 9.2 for illustrative visualizations.

Training details. We employ local SGD with momentum as the LocalSolver. For fine-tuning-based methods, we additionally perform one epoch of local fine-tuning before evaluation. For cluster- and mixture-model-based methods, we set the number of components \(C\) to 3. Additional implementation and hyperparameter details are provided in App. 9.3.

3.2 Experimental Results↩︎

Main results. Tab. 1 yields four main observations. First, FedSPM consistently outperforms all baselines in both metrics across all datasets, regardless of model capacity. Second, routing-free personalization methods (red background) trade system accuracy for average accuracy compared with global-model methods, revealing severe model drift. The substantial recovery in system accuracy achieved by FedDRM and FedSPM further underscores the necessity of client routing for system-level performance. Third, the poor performance of mixture-model-based methods indicates that personalizing only the mixing weights is insufficient to handle severe concept shift without component-level personalization. Fourth, FedSPM (\(1\times\)) slightly outperforms FedSPM (\(C\times\)), suggesting that under fixed data budgets on these relatively simple tasks, the benefit of increased model capacity does not outweigh the sample dilution from separating encoders per component.

Impact of dual heterogeneity intensity. We further evaluate the robustness of FedSPM under different levels of inter-client and intra-client heterogeneity on FMNIST. To isolate their effects, we vary the Dirichlet concentration parameter \(\alpha_{\text{inter}}\in\{0.5,1.0,2.0\}\) while fixing \(\alpha_{\text{intra}}=2.0\), and vary \(\alpha_{\text{intra}}\in\{0.5,1.0,2.0\}\) while fixing \(\alpha_{\text{inter}}=1.0\). As shown in Tab. 2, FedSPM consistently achieves the best performance across all settings, demonstrating its strong adaptability to dual heterogeneity.

Impact of \(\mu\). We study the impact of momentum \(\mu\) under a constant learning rate on FMNIST. To align the empirical evaluation with our theory, we report the running average of \(\|\nabla_\phi F(\zeta^{(t)})\|^2\), together with the normalized cumulative sums of the estimated bridge errors \(\widehat{\delta}^{(t)}\) and \(\widehat{\varepsilon}^{(t)}\). See detailed estimation procedure in App. 9.4. Since computing these quantities requires evaluating the profiled objective \(F\) and the DRM objective \(f^{(t)}\), both of which depend on the Lagrange multipliers \(\{\lambda_{ic}\}\), we numerically solve the nonlinear system in 6 for \(\{\lambda_{ic}\}\) at each communication round. Fig. 2 yields three observations. First, the averaged gradient norm consistently decreases for all tested values of \(\mu\), suggesting that FedSPM empirically converges toward a stationary point of the profiled objective \(F\). Second, as \(\mu\) increases, the convergence first improves and then deteriorates, with the best performance achieved at \(\mu=0.95\), consistent with the momentum-induced trade-off predicted by our theory. Third, the bridge errors are empirically well controlled: their normalized cumulative estimates remain small or decrease during training. This provides empirical support for the asymptotic behaviors \(\bar{\delta}_T=o(1)\) and \(\bar{\varepsilon}_T=o(1)\), thereby validating the soundness of Assumptions 4 and 5.

Impact of \(C\). We study the impact of the number of components \(C\) on FMNIST. To largely isolate the effect of varying \(C\) from model capacity changes, we employ a shared encoder across all components so that increasing \(C\) only adds lightweight component-specific heads. As shown in Fig. 3, increasing \(C\) from 1 to 2 yields the largest gain, supporting the benefit of mixture modeling under dual heterogeneity. Further increasing \(C\) brings smaller but consistent improvements, suggesting that mild over-specification can provide a more flexible approximation to the predictive distribution, consistent with theoretical results on over-specified mixture-of-experts models [39], [40].

4 Experiment on Real Medical Dataset↩︎

We complement the controlled experiments on benchmark datasets with a case study on Fed-ISIC2019 [41], a real-world medical dataset where dual heterogeneity arises naturally from varying imaging protocols, skewed class proportions, and diverse pathological features. Fed-ISIC2019 comprises 23,247 RGB dermoscopic images collected from 6 clinical centers for an 8-class skin-lesion classification task, naturally forming a 6-client FL system where each client corresponds to one clinical center. Since no heterogeneity is artificially imposed, this natural setting serves two purposes: validating the necessity of jointly modeling dual heterogeneity in real-world FL applications and verifying the effectiveness of our semiparametric model against parametric alternatives. The training details are provided in App. 9.5.

Fig. 4 yields four main observations. First, FedSPM achieves the best performance across all three metrics, showing its practical feasibility in real-world medical scenarios. Second, mixture-model-based methods generally outperform methods that do not explicitly model intra-client heterogeneity, confirming the presence of intra-client heterogeneity in real-world medical data and validating the necessity of jointly modeling dual heterogeneity. Third, FedSPM substantially improves routing accuracy over FedGMM, highlighting the limitation of parametric GMMs and validating the effectiveness of our semiparametric model. Finally, in contrast to the benchmark results, FedSPM (\(C\times\)) outperforms FedSPM (\(1\times\)), suggesting that on this medically complex dataset, the benefit of expanded model capacity for capturing diverse pathological features outweighs the sample dilution effect.

5 Conclusion↩︎

This paper presents FedSPM, a semiparametric mixture framework that extends the routing-prediction FL paradigm from inter-client heterogeneity to dual heterogeneity. By combining DRM-based EL with mixture modeling, FedSPM leverages latent intra-client structure to improve both routing and prediction while balancing model flexibility with effective information sharing. A federated EM procedure further enables practical optimization in the client-server setting. Empirical results on controlled benchmarks and real-world medical data show that FedSPM outperforms existing approaches in both routing and prediction performance under complex dual heterogeneity. Overall, this work paves the way for expertise-aware FL systems that transform dual heterogeneity from an obstacle into a source of system-level intelligence.

6 Related Work on DRM and EL↩︎

The DRM, first introduced by [30], provides a statistical framework for modeling multiple related populations by formulating their densities as ratios with respect to a shared baseline distribution. It is highly flexible and subsumes commonly used parametric families such as the binomial, gamma, and normal distributions [42]. As a semiparametric model, the DRM does not impose parametric assumptions on the baseline distribution, which can instead be handled nonparametrically via EL [31]. This integration gives rise to the DRM-based EL approach, which has garnered substantial attention over the past few decades [43]–[46]. Moreover, [47] formally established the equivalence between the maximum DRM-based EL estimators and the corresponding dual EL estimators. Since the dual EL admits an analytical form and is computationally tractable, it significantly facilitates the implementation and application of DRM-based EL methods [48]–[50].

Recently, [10] first introduced DRM-based EL into FL and established a routing-prediction FL paradigm, thereby enabling explicit client routing together with personalized representation learning. However, this paradigm treats each client as internally homogeneous. Since intra-client heterogeneity often arises in real-world FL scenarios, ignoring such latent structure degrades both routing and prediction performance, motivating our semiparametric mixture extension.

7 The Federated EM Algorithm↩︎

8 Mathematical Details↩︎

8.1 Derivation of the Log-EL↩︎

The observed data contain only \((x_{ij},y_{ij})\), while the component assignment \(z_{ij}\) is latent. Therefore, the observed-data log-EL is obtained from the augmented distribution \(P_i^{X,Y,Z}\) by marginalizing over all possible values of \(Z\): \[\begin{align} \ell(\zeta, G) &\mathrel{\vcenter{:}}= \sum_{i,j} \log \sum_c \underbrace{P_i^{X,Y,Z}(dx_{ij},\{y_{ij}\},\{c\})}_{\text{complete-data contribution for }Z=c}\\ &= \sum_{i,j} \log \sum_c \pi_{ic} P_{ic}^{X,Y}(dx_{ij},\{y_{ij}\}) \\ &= \sum_{i,j} \log \sum_c \pi_{ic} P_{ic}^{Y|X}(\{y_{ij}\}\mid x_{ij}) P_{ic}^{X}(dx_{ij}) \\ &= \sum_{i,j} \log \sum_c \pi_{ic} P_{ic}^{Y|X}(\{y_{ij}\}\mid x_{ij}) \frac{dP_{ic}^X}{dG}(x_{ij})G(dx_{ij}) \\ &= \sum_{i,j} \log \sum_c \pi_{ic} P_{ic}^{Y|X}(\{y_{ij}\}\mid x_{ij}) \frac{dP_{ic}^X}{dG}(x_{ij})r_{ij}\\ &=\sum_{i,j} \log r_{ij} + \sum_{i,j} \log \biggl[ \sum_{c} \pi_{ic} \exp \left( \gamma_{ic} + \xi_{ic}^\top h_{\nu_c}(x_{ij}) \right) P_{ic}^{Y|X}(\{y_{ij}\}\mid x_{ij})\biggr]. \end{align}\]

8.2 Derivation of the Optimal Baseline Weights↩︎

For fixed \(\zeta\), maximizing \(\ell\) over \(G\) reduces to maximizing \(\sum_{i,j}\log r_{ij}\) subject to the constraints in 4 . The corresponding Lagrangian is: \[{!}{ \displaystyle \mathcal{L}\mathrel{\vcenter{:}}= \sum_{i,j}\log r_{ij} -n\kappa\biggl[\sum_{i,j}r_{ij}-1\biggr] -n\sum_{i',c}\lambda_{i'c}\sum_{i,j} \left[ \exp\left(\gamma_{i'c}+\xi_{i'c}^{\top}h_{\nu_c}(x_{ij})\right)-1 \right]r_{ij}. }\] Setting \(\partial\mathcal{L}/\partial r_{ij}=0\) gives \[r_{ij}^{-1} -n\kappa -n\sum_{i',c}\lambda_{i'c} \left[ \exp\left(\gamma_{i'c}+\xi_{i'c}^{\top}h_{\nu_c}(x_{ij})\right)-1 \right]=0.\] Multiplying the above equation by \(r_{ij}\), summing over \((i,j)\), and applying the constraints in 4 gives \(\kappa=1\), which yields \[r_{ij}^*(\zeta)=\frac{1}{n} \biggl\{ 1+ \sum_{i',c}\lambda_{i'c} \left[ \exp\left(\gamma_{i'c}+\xi_{i'c}^{\top}h_{\nu_c}(x_{ij})\right)-1 \right] \biggr\}^{-1}.\] Substituting \(r_{ij}^*(\zeta)\) back into 4 gives the following nonlinear system for \(\{\lambda_{ic}\}\): \[\sum_{i,j} \frac{ \exp\left(\gamma_{i'c}+\xi_{i'c}^{\top}h_{\nu_c}(x_{ij})\right)-1 }{ 1+ \sum_{i'',c'}\lambda_{i''c'} \left[ \exp\left(\gamma_{i''c'}+\xi_{i''c'}^{\top}h_{\nu_{c'}}(x_{ij})\right)-1 \right] }=0, \quad \forall i'\in[m],\;c\in[C].\]

8.3 Derivation of the Complete-Data Profile Log-EL↩︎

Unlike the observed-data log-EL, the complete-data log-EL treats the latent component indicators \(\{z_{ijc}\}\) as observed. Thus, it uses the corresponding complete-data contribution from \(P_i^{X,Y,Z}\) directly, without marginalizing over \(Z\): \[\begin{align} \ell^{\text{c}}(\zeta, G) &\mathrel{\vcenter{:}}= \sum_{i,j,c}z_{ijc}\log P_i^{X,Y,Z}(dx_{ij},\{y_{ij}\},\{c\})\\ &= \sum_{i,j,c}z_{ijc}\log\Bigl[\pi_{ic}P_{ic}^{X,Y}(dx_{ij},\{y_{ij}\})\Bigr]\\ &= \sum_{i,j,c}z_{ijc}\log\Bigl[\pi_{ic}P_{ic}^{Y|X}(\{y_{ij}\}\mid x_{ij})P_{ic}^{X}(dx_{ij})\Bigr]\\ &= \sum_{i,j,c}z_{ijc}\log\Bigl[\pi_{ic}P_{ic}^{Y|X}(\{y_{ij}\}\mid x_{ij})\frac{dP_{ic}^X}{dG}(x_{ij})G(dx_{ij})\Bigr]\\ &= \sum_{i,j,c}z_{ijc}\log\Bigl[\pi_{ic}\frac{dP_{ic}^X}{dG}(x_{ij})r_{ij} P_{ic}^{Y|X}(\{y_{ij}\}\mid x_{ij}) \Bigr]\\ &=\sum_{i,j}\log r_{ij}+\sum_{i,j,c}z_{ijc}\Bigl[\log P_{ic}^{Y|X}(\{y_{ij}\}\mid x_{ij})+\gamma_{ic}+\xi_{ic}^\top h_{\nu_c}(x_{ij})+\log \pi_{ic} \Bigr]. \end{align}\] Since maximizing \(\ell^{\text{c}}\) over \(G\) for fixed \(\zeta\) also reduces to maximizing \(\sum_{i,j}\log r_{ij}\) subject to the constraints in 4 , \(\ell^{\text{c}}\) shares the same optimal baseline weights as \(\ell\). Thus, substituting \(r_{ij}^*(\zeta)\) back into \(\ell^{\text{c}}(\zeta,G)\) yields the complete-data profile log-EL: \[\begin{align} p\ell^{\text{c}}(\zeta)&\mathrel{\vcenter{:}}= \sup_G\ell^{\text{c}}(\zeta,G)\\ &=\sum_{i,j}\log r_{ij}^*(\zeta)+\sum_{i,j,c}z_{ijc}\biggl[\log P_{ic}^{Y|X}(\{y_{ij}\}\mid x_{ij})+\gamma_{ic}+\xi_{ic}^\top h_{\nu_c}(x_{ij})+\log \pi_{ic} \biggr]. \end{align}\]

8.4 EM Algorithm Details↩︎

The EM algorithm is an iterative procedure that starts from an initial estimate \(\zeta^{(0)}\). At iteration \(t\), given the current parameter estimate \(\zeta^{(t)}\), the algorithm alternates between the following two steps.

E-step. The latent component assignments \(\{z_{ij}\}\) are estimated by their posterior responsibilities, which serve as soft guesses computed from the observed data under the current parameter estimate: \[w_{ijc}^{(t)} \mathrel{\vcenter{:}}= \mathbb{E}\left[z_{ijc}\middle|\mathcal{D};\zeta^{(t)}\right] =P_i^{Z|X,Y}(\{c\}\mid x_{ij},y_{ij};\zeta^{(t)}).\] The second equality follows because \(z_{ijc}\) is the indicator of the event \(\{Z=c\}\). Applying Bayes’ rule with prior \(\pi_{ic}^{(t)}\) and likelihood \(P_{ic}^{X,Y}(dx_{ij},\{y_{ij}\};\zeta^{(t)})\) gives \[w_{ijc}^{(t)} = \frac{ \pi_{ic}^{(t)}P_{ic}^{X,Y}(dx_{ij},\{y_{ij}\};\zeta^{(t)}) }{ \sum_{c'} \pi_{ic'}^{(t)}P_{ic'}^{X,Y}(dx_{ij},\{y_{ij}\};\zeta^{(t)}) }.\] Using the factorization of \(P_{ic}^{X,Y}\) into \(P_{ic}^{Y|X}\) and \(P_{ic}^{X}\) gives \[w_{ijc}^{(t)} = \frac{ \pi_{ic}^{(t)} P_{ic}^{Y|X}(\{y_{ij}\}\mid x_{ij};\zeta^{(t)}) P_{ic}^{X}(dx_{ij};\zeta^{(t)}) }{ \sum_{c'} \pi_{ic'}^{(t)} P_{ic'}^{Y|X}(\{y_{ij}\}\mid x_{ij};\zeta^{(t)}) P_{ic'}^{X}(dx_{ij};\zeta^{(t)}) }.\] Under the EL representation of \(G\), \(P_{ic}^{X}(dx_{ij})=dP_{ic}^{X}/dG(x_{ij})r_{ij}\). Therefore, \[w_{ijc}^{(t)} = \frac{ \pi_{ic}^{(t)} P_{ic}^{Y|X}(\{y_{ij}\}\mid x_{ij};\zeta^{(t)}) \frac{dP_{ic}^{X}}{dG}(x_{ij};\zeta^{(t)})r_{ij} }{ \sum_{c'} \pi_{ic'}^{(t)} P_{ic'}^{Y|X}(\{y_{ij}\}\mid x_{ij};\zeta^{(t)}) \frac{dP_{ic'}^{X}}{dG}(x_{ij};\zeta^{(t)})r_{ij} }.\] Finally, substituting the DRM density ratio and canceling the common factor \(r_{ij}\) yields \[w_{ijc}^{(t)} = \frac{ P_{ic}^{Y|X}(\{y_{ij}\}\mid x_{ij};\zeta^{(t)}) \exp\left(\gamma_{ic}^{(t)}+(\xi_{ic}^{(t)})^\top h_{\nu_c^{(t)}}(x_{ij})\right) \pi_{ic}^{(t)} }{ \sum_{c'} P_{ic'}^{Y|X}(\{y_{ij}\}\mid x_{ij};\zeta^{(t)}) \exp\left(\gamma_{ic'}^{(t)}+(\xi_{ic'}^{(t)})^\top h_{\nu_{c'}^{(t)}}(x_{ij})\right) \pi_{ic'}^{(t)} }.\] Then, the resulting profile \(Q\)-function is the conditional expectation of the complete-data profile log-EL with respect to \(\{z_{ij}\}\): \[\begin{align} Q^{(t)}(\zeta) &\mathrel{\vcenter{:}}= \mathbb{E}\left[p\ell^{\text{c}}(\zeta)\middle|\mathcal{D};\zeta^{(t)}\right] \\ &= Q_1^{(t)}(\alpha,\beta,\theta) + Q_2^{(t)}(\gamma,\xi,\nu) + Q_3^{(t)}(\pi), \end{align}\] where \[\begin{align} Q_1^{(t)}(\alpha,\beta,\theta) \mathrel{\vcenter{:}}=& \sum_{i,j,c}w_{ijc}^{(t)}\log P_{ic}^{Y|X}(\{y_{ij}\}\mid x_{ij}),\\ Q_2^{(t)}(\gamma, \xi, \nu) \mathrel{\vcenter{:}}=& \sum_{i,j,c}w_{ijc}^{(t)} \left[\gamma_{ic} + \xi_{ic}^\top h_{\nu_c}(x_{ij})\right]\\ &- \sum_{i,j}\log \biggl(1 + \sum_{i',c} \lambda_{i'c} \left[\exp\left(\gamma_{i'c} + \xi_{i'c}^\top h_{\nu_c}(x_{ij})\right) - 1\right]\biggr),\\ Q_3^{(t)}(\pi) \mathrel{\vcenter{:}}=& \sum_{i,j,c}w_{ijc}^{(t)}\log\pi_{ic}. \end{align}\]

M-step. Instead of directly maximizing the profile log-EL, the M-step maximizes \(Q^{(t)}\) with respect to \(\zeta\), i.e.,  \[\zeta^{(t+1)}=\arg\max Q^{(t)}(\zeta).\] Here we focus on the analytic update of the mixing weights \(\pi\), which solves \[\begin{align} \max_{\pi}\quad &\sum_{ijc} w_{ijc}^{(t)}\log \pi_{ic} \\ \mathrm{s.t.}\quad &\sum_{c} \pi_{ic}=1,\qquad i\in [m],\\ &\pi_{ic}\ge 0,\qquad i\in [m],\quad c\in [C]. \end{align}\] Using the method of Lagrange multipliers, the Lagrangian is \[\mathcal{L}_\pi \mathrel{\vcenter{:}}= \sum_{ijc} w_{ijc}^{(t)}\log \pi_{ic} + \sum_{i} \iota_i \left(\sum_{c}\pi_{ic}-1\right).\] At the optimum, we require \[\frac{\partial\mathcal{L}_\pi}{\partial \pi_{ic}}=0,\quad\frac{\partial\mathcal{L}_\pi}{\partial \iota_i}=0, \quad i\in[m],\quad c\in[C].\] Solving these equations gives \[\pi_{ic}^{(t+1)} = \frac{1}{n_i} \sum_{j=1}^{n_i}w_{ijc}^{(t)}.\]

8.5 Derivation of the Analytical Form of the Lagrange Multipliers↩︎

We derive the analytical form of \(\{\lambda_{ic}\}\) from the first-order optimality condition of \(Q_2^{(t)}\). For each \(i_0\in [m]\) and \(c_0\in [C]\), taking the derivative with respect to \(\gamma_{i_0,c_0}\) gives \[\frac{\partial Q_2^{(t)}}{\partial \gamma_{i_0,c_0}} =\sum_{j=1}^{n_{i_0}}w_{i_0,j,c_0}^{(t)} -\lambda_{i_0,c_0}\sum_{i,j} \frac{ \exp\left(\gamma_{i_0,c_0}+\xi_{i_0,c_0}^{\top}h_{\nu_{c_0}}(x_{ij})\right) }{ 1+ \sum_{i',c}\lambda_{i'c} \left[ \exp\left(\gamma_{i'c}+\xi_{i'c}^{\top}h_{\nu_c}(x_{ij})\right)-1 \right] },\] where the term involving \(\partial\lambda_{i_0,c_0}/\partial\gamma_{i_0,c_0}\) vanishes due to 6 . By 4 and 5 , we further have \[\sum_{i,j} \frac{ \exp\left(\gamma_{i_0,c_0}+\xi_{i_0,c_0}^{\top}h_{\nu_{c_0}}(x_{ij})\right) }{ 1+ \sum_{i',c}\lambda_{i'c} \left[ \exp\left(\gamma_{i'c}+\xi_{i'c}^{\top}h_{\nu_c}(x_{ij})\right)-1 \right] } =n\sum_{i,j} \exp\left(\gamma_{i_0,c_0}+\xi_{i_0,c_0}^{\top}h_{\nu_{c_0}}(x_{ij})\right)r_{ij}^* =n.\] Therefore, \[\frac{\partial Q_2^{(t)}}{\partial \gamma_{i_0,c_0}} =\sum_{j=1}^{n_{i_0}}w_{i_0,j,c_0}^{(t)} -n\lambda_{i_0,c_0}.\] Setting the derivative to zero yields \[\lambda_{i_0,c_0}=\frac{1}{n}\sum_{j=1}^{n_{i_0}}w_{i_0,j,c_0}^{(t)}.\]

8.6 Convergence Proofs↩︎

The proof proceeds in four steps. First, we unroll the local momentum recursion and rewrite the aggregated update as a single descent step driven by a weighted effective gradient. The weights are the momentum accumulation factors \(\{\mu_{E,e}\}_{e=0}^{E-1}\), which make the dependence on \(\mu\) explicit through \(a_E\) and \(s_E\). Second, we decompose this effective gradient into the averaged exact gradient, the local-update drift, and the stochastic noise. Lemmas 5 and 6 bound the stochastic noise and drift terms, respectively. Third, Lemma 7 establishes a one-round descent estimate for the tractable surrogate \(\widetilde{f}^{(t)}\). Finally, we combine Assumptions 4 and 5 with the EM bridge arguments in Lemmas 1 and 2 to transfer the descent to the profiled log-EL and telescope over the outer rounds.

8.6.1 Additional Notations↩︎

For the convergence proof, define the stochastic gradient noise as \[\omega_{i,e}^{(t)}\mathrel{\vcenter{:}}= \widetilde{g}_{i,e}^{(t)}- \nabla\widetilde{f}_i^{(t)}(\phi_{i,e}^{(t)}).\] Let \(\mathcal{F}^{(-1)}\) be the trivial \(\sigma\)-algebra. For \(t\ge0\), define \[\mathcal{F}^{(t)}\mathrel{\vcenter{:}}= \sigma\left( \omega_{i,e}^{(s)}: 0\le s\le t, \;i\in[m], \;0\le e\le E-1 \right).\] For each client \(i\) and local step \(e\), define \[\mathcal{F}_{i,-1}^{(t)}\mathrel{\vcenter{:}}=\mathcal{F}^{(t-1)},\quad \mathcal{F}_{i,e}^{(t)}\mathrel{\vcenter{:}}= \mathcal{F}^{(t-1)}\vee \sigma\left( \omega_{i,0}^{(t)},\dots,\omega_{i,e}^{(t)} \right).\]

8.6.2 Key Lemmas↩︎

8.6.3 Proof of Theorem 1↩︎

9 Experimental Details↩︎

9.1 Benchmark Dataset Details↩︎

We conduct experiments on FMNIST [33], CIFAR-10 [34], and CIFAR-100 [34]. FMNIST contains 70,000 \(28\times 28\) grayscale images from 10 fashion classes. CIFAR-10 contains 60,000 \(32\times 32\) RGB images from 10 object classes. CIFAR-100 contains 60,000 \(32\times 32\) RGB images from 100 object classes, grouped into 20 superclasses. We construct three tasks of increasing complexity: (a) 10-class classification on FMNIST, (b) 10-class classification on CIFAR-10, and (c) 20-class classification using the CIFAR-100 superclasses.

9.2 Visualization of Dual Heterogeneity↩︎

We provide visualizations of the dual heterogeneity settings in our benchmark experiments.

To induce covariate shift, images are transformed at both the client and component levels. At the component level, component 0 retains the original image, while component 1 applies grayscale inversion for FMNIST and green-channel inversion for CIFAR-10/100. At the client level, the 8 clients correspond to the \(2^3\) binary combinations of three orthogonal dimensions: color shift (red- vs. blue-dominant), vertical spatial bias (top- vs. bottom-brightened), and horizontal spatial bias (left- vs. right-brightened). As shown in Fig. 5, these hierarchical transformations yield a visually conspicuous covariate shift across clients.

Fig. 6 illustrates the impact of \(\alpha_{\text{inter}}\) on client label distributions. A smaller \(\alpha_{\text{inter}}\) produces increasingly skewed class proportions and uneven dataset sizes across clients, corresponding to intensified inter-client heterogeneity. Similarly, Fig. 7 visualizes how \(\alpha_{\text{intra}}\) modulates the intra-client mixing weights. While a small \(\alpha_{\text{intra}}\) concentrates local data into a single dominant component, a larger \(\alpha_{\text{intra}}\) yields more balanced mixtures, indicating stronger intra-client heterogeneity.

9.3 Training Details on Benchmark Datasets↩︎

We employ local SGD with momentum, a batch size of 128, and an initial learning rate of 0.01 with cosine annealing. Each client performs 10 local steps per round with momentum \(0.9\) for image classification, and 15 local steps with momentum \(0.95\) for client routing. The encoder output dimensions and total communication rounds are set to 32 and 200 for FMNIST, and 64 and 400 for both CIFAR-10 and CIFAR-100. For fine-tuning-based methods, we additionally perform one epoch of local fine-tuning before evaluation. For cluster- and mixture-model-based methods, we set the number of components \(C\) to 3.

9.4 Estimation of the Bridge Errors↩︎

This section details the empirical estimation of the bridge errors \(\varepsilon^{(t)}\) and \(\delta^{(t)}\) reported in Fig. 2. Computing these quantities requires evaluating the DRM objective \(f^{(t)}\) and its gradient \(\nabla f^{(t)}\), both of which depend on the Lagrange multipliers \(\{\lambda_{ic}\}\). Therefore, at each communication round, we numerically solve the nonlinear system in 6 for \(\{\lambda_{ic}\}\) using a Jacobian-based root solver, followed by a trust-region least-squares refinement when necessary. The resulting multipliers are then used to compute \(\{f^{(t)}(\phi^{(t)})\}_{t=0}^{T-1}\) and \(\{\nabla f^{(t)}(\phi^{(t)})\}_{t=0}^{T-1}\) along the training trajectory.

Estimation of the gradient bridge error \(\varepsilon^{(t)}\). Assumption 4 gives \[\mathbb{E}\left[\left\|\nabla f^{(t)}(\phi^{(t)})\right\|^2\right]\leq \Gamma_2\mathbb{E}\left[\left\|\nabla\widetilde{f}^{(t)}(\phi^{(t)})\right\|^2\right]+\varepsilon^{(t)}.\] Motivated by this inequality, we first estimate the coefficient \(\Gamma_2\) from the ratios between the gradient norms of the DRM objective and its tractable surrogate. To prevent numerical instability when the surrogate gradient norm vanishes in the denominator, our evaluation is restricted to the valid set: \[\mathcal{T}_{\text{valid}}\mathrel{\vcenter{:}}= \left\{t:\left\|\nabla\widetilde{f}^{(t)}(\phi^{(t)})\right\|^2>10^{-9}\right\}.\] We then estimate \(\Gamma_2\) by the 95% quantile of the valid ratios to ensure robustness against outliers: \[\widehat{\Gamma}_2\mathrel{\vcenter{:}}= \text{Quantile}_{0.95}\left\{ \frac{ \left\|\nabla f^{(t)}(\phi^{(t)})\right\|^2 }{ \left\|\nabla\widetilde{f}^{(t)}(\phi^{(t)})\right\|^2 } : t\in\mathcal{T}_{\mathrm{valid}} \right\}.\] Finally, given \(\widehat{\Gamma}_2\), \(\varepsilon^{(t)}\) is computed as the remaining positive gap: \[\widehat{\varepsilon}^{(t)}\mathrel{\vcenter{:}}= \max\left(\left\|\nabla f^{(t)}(\phi^{(t)})\right\|^2-\widehat{\Gamma}_2\left\|\nabla \widetilde{f}^{(t)}(\phi^{(t)})\right\|^2,0\right).\]

Estimation of the function-value bridge error \(\delta^{(t)}\). Assumption 5 gives \[\mathbb{E}\left[f^{(t)}(\phi^{(t+1)})-f^{(t)}(\phi^{(t)})\right]\leq \mathbb{E}\left[\widetilde{f}^{(t)}(\phi^{(t+1)})-\widetilde{f}^{(t)}(\phi^{(t)})\right] +\delta^{(t)}.\] Accordingly, we estimate \(\delta^{(t)}\) as \[\widehat{\delta}^{(t)}\mathrel{\vcenter{:}}= \max\left(\left[f^{(t)}(\phi^{(t+1)})-f^{(t)}(\phi^{(t)})\right]-\left[\widetilde{f}^{(t)}(\phi^{(t+1)})-\widetilde{f}^{(t)}(\phi^{(t)})\right],0\right).\]

9.5 Training Details on Real Medical Dataset↩︎

The training setup for Fed-ISIC2019 largely follows that of the benchmark experiments, with several adjustments. Specifically, the encoder output dimension is set to 256 for image classification and 128 for client routing. The number of communication rounds is set to 400, and the batch size is set to 32.

References↩︎