FedSPM: Routing-Enabled Federated Learning under Dual Heterogeneity via Semiparametric Mixture
Zijian Wang1
Pengfei Li2
Guangyu Yang1
Qiong Zhang1∗
1
arXiv:2607.04085v1 [cs.LG] 5 Jul 2026
2
Institute of Statistics and Big Data, Renmin University of China Department of Statistics and Actuarial Science, University of Waterloo
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 O(1/ 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 here.
1
Introduction
Federated learning (FL) [27] 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 [43] and mobile devices [20], making FL particularly promising when data are sensitive, geographically dispersed, or impractical to centralize [21, 14, 44]. 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 [15], thereby slowing convergence [24] and degrading global model performance [50]. In contrast, the recent routing-prediction FL paradigm [40] reframes inter-client heterogeneity as a useful signal of client specialization. Like personalized FL [23, 10, 3], 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. ∗ Correspondence to: Qiong Zhang ([email protected]) and Guangyu Yang ([email protected])
Preprint.
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 [26, 41]. 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 [46, 7], whereas visually similar cases may correspond to different diagnoses [4, 6]. 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 [22, 1], aggregation reweighting [38, 25], or refined optimization [13, 34], but their shared predictor still struggles to adapt to client-specific distributions. Conventional personalized methods, including local fine-tuning [39], regularization [9, 23], client clustering [10, 35], and representation learning [3, 30], improve prediction within each client’s local domain, but provide no mechanism for server-side routing of external queries. Routing-based personalization [40] 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 [26, 41], but its assumption of shared component distributions across clients limits flexibility in modeling client-specific latent structures.
Figure 1: Overview of the FedSPM framework. 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) [2] 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
2
Method
Problem Formulation. Consider an FL system with m clients for K-class classification. Each client i ∈ [m] := {1, . . . , m} has a local dataset Di := {(xij , yij )}j∈[ni ] with samples drawn Pm independently from a client-specific distribution PiX,Y .SLet n := i=1 ni be the total sample size, m ρi := ni /n be the sample fraction of client i, and D := i=1 Di be the pooled dataset. 2.1
Semiparametric Mixture Model
To capture intra-client heterogeneity, we introduce a latent component variable Z ∈ [C] and model each observed distribution PiX,Y through an augmented distribution PiX,Y,Z . Specifically, we assume X,Y (xij , yij ) | (zij = c) ∼ Pic , zij ∼ PiZ with PiZ ({c}) = πic , ∀c ∈ [C], j ∈ [ni ]. (1) P X,Y Here, πi := (πi1 , . . . , πiC ) with πic ≥ 0 and c πic = 1. Marginalizing over Z gives Pi = P X,Y X,Y . Notably, we allow both the mixing weights π and the component distributions P π P i ic ic ic c to vary across clients, providing a flexible model of intra-client heterogeneity. X,Y While allowing fully client-specific Pic is expressive, it prevents effective information sharing Y |X X,Y X across clients. We factor Pic into its predictive distribution Pic and feature distribution Pic , 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 Y |X
⊤ Pic ({k} | x) ∝ exp(αikc + βikc gθc (x)), (2) where gθc (x) is a shared embedding for component c. The client-specific parameters (αikc , βikc ) Y |X allow the marginal predictive distribution Pi 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: X ⊤ dPic /dG(x) = exp γic + ξic hνc (x) , (3) X X where dPic /dG is the Radon–Nikodym derivative of Pic relative to G, and hνc is a shared basis for component c. The client-specific tilting parameters (γic , ξic ) induce differences in the marginal feature distribution PiX across clients, capturing covariate shift while preserving a shared structure. X,Y Label shift. Finally, differences in the mixing weights πi and the component distribution Pic jointly induce variation in the marginal label distribution PiY 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. A for related work on DRM and EL). Specifically, GP is represented as a discrete distribution supported on the pooled observations, taking the form G = i,j rij δxij with unknown rij ≥ 0. Crucially, all samples across all clients jointly determine G, serving as a key mechanism for information sharing. X To valid probability distributions, the masses {rij } must satisfy R ensure that G and R each Pic define G(dx) = 1 and exp γi′ c + ξi⊤′ c hνc (x) G(dx) = 1 for all i′ , c, i.e., X X rij = 1, exp γi′ c + ξi⊤′ c hνc (xij ) rij = 1, ∀i′ ∈ [m], c ∈ [C]. (4) i,j
i,j
In summary, the predictive model (2), the DRM-EL feature model (3), and the mixing weights πi together define our semiparametric mixture model. Remark 1 (Identifiability and interpretability). Our model is intended as a flexible approximation to the client-specific joint distribution PiX,Y , rather than as a tool for recovering identifiable or interpretable latent structures. Its non-identifiability arises from both the neural network implementations of the representation maps gθc and hνc , and the mixture structure itself. For the latter, if X,Y X,Y two components on client i collapse, i.e., Pic = Pic for some c1 ̸= c2 , then redistributing 1 2 P X,Y X,Y mass between πic1 and πic2 leaves Pi = c πic Pic unchanged. Thus, the learned components should be viewed as auxiliary constructs for routing and prediction, not as recovered true subgroups. 3
X,Y X,Y Remark 2 (Comparison to existing work). Prior work such as [26, 41] assume Pic = Pjc for all i, j, c, i.e., identical component distributions across clients. This assumption is often unrealistic. For instance, even data from clinically similar patient groups can still differ substantially across hospitals due to site-specific acquisition protocols, imaging devices, and preprocessing pipelines.
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 ζ := (α, β, θ, γ, ξ, ν, π) be the parameters of interest. The log-EL based on D is: X X X Y |X ⊤ ℓ(ζ, G) = log rij + log πic exp γic + ξic hνc (xij ) Pic ({yij } | xij ) . i,j
c
i,j
See App. C.1 for the derivation. Since G is a nuisance parameter, we work with the profile log-EL pℓ(ζ) := supG ℓ(ζ, G), where the supremum is taken under the constraints in (4). By the Lagrange multiplier method, the optimal baseline weights are: X −1 ∗ rij (ζ) = n−1 1 + λi′ c exp γi′ c + ξi⊤′ c hνc (xij ) − 1 , (5) i′ ,c
where {λic } are the solutions to exp γi′ c + ξi⊤′ c hνc (xij ) − 1 = 0, P 1 + i′′ ,c′ λi′′ c′ exp γi′′ c′ + ξi⊤′′ c′ hνc′ (xij ) − 1 i,j
X
∀i′ ∈ [m], c ∈ [C].
(6)
See App. C.2 for the derivation. Substituting (5) back into ℓ(ζ, G) yields pℓ(ζ). However, directly maximizing pℓ(ζ) remains computationally impractical because each evaluation requires solving the nonlinear system in (6) for the Lagrange multipliers {λ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 {zij } as missing data and employ the EM algorithm to maximize pℓ(ζ). Let zijc := 1(zij = c) denote the indicator that the j-th sample on the i-th client belongs to the c-th subpopulation. If {zij } were observed, the complete-data profile log-EL would be X X Y |X ∗ ⊤ pℓc (ζ) = log rij (ζ) + zijc log Pic ({yij } | xij ) + γic + ξic hνc (xij ) + log πic . i,j
i,j,c
See App. C.3 for derivation. Since {zij } 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 ζ (t) , we compute the posterior responsibilities for each zij as h i (t) Y |X (t) (t) (t) wijc := E zijc D; ζ (t) ∝ Pic ({yij } | xij ; ζ (t) ) exp γic + (ξic )⊤ hν (t) (xij ) πic . c
(7)
The resulting profile Q-function is the conditional expectation of the complete-data profile log-EL: h i (t) (t) (t) Q(t) (ζ) := E pℓc (ζ) D; ζ (t) = Q1 (α, β, θ) + Q2 (γ, ξ, ν) + Q3 (π), (8) where (t)
Q1 (α, β, θ) :=
X
(t)
Y |X
wijc log Pic
({yij } | xij ),
(t)
Q3 (π) :=
i,j,c (t)
Q2 (γ, ξ, ν) :=
X i,j,c
X
(t)
wijc log πic , and
i,j,c
X X (t) ⊤ wijc γic + ξic hνc (xij ) − log 1 + λi′ c exp γi′ c + ξi⊤′ c hνc (xij ) − 1 . i′ ,c
i,j
4
(t)
(t)
(t)
M-step. We update ζ (t+1) = arg max Q(t) (ζ). Since Q1 , Q2 , and Q3 are separable, the three parameter blocks can be optimized independently. For convenience, we refer to (α, β, θ) as the (t) supervised block, and to (γ, ξ, ν) as the DRM block. For the mixing weights, maximizing Q3 under P c πic = 1 yields ni 1 X (t) (t+1) w . (9) πic = ni j=1 ijc See App. C.4 for derivation details of the EM algorithm. For the supervised block, maximizing (t) Q1 is equivalent to minimizing a weighted cross-entropy loss, which can be directly solved via (t) standard gradient-based methods: α(t+1) , β (t+1) , θ(t+1) = arg max Q1 (α, β, θ). For the DRM (t) block, directly maximizing Q2 faces the same computational bottleneck as pℓ(ζ): each evaluation still requires solving the nonlinear system for the Lagrange multipliers {λic }. In contrast, the EM (t) algorithm offers a pivotal resolution to this bottleneck: at any critical point of Q2 , each λic admits P (t) n i the analytic form n−1 j=1 wijc (see App. C.5). Substituting this expression back shows that every (t)
critical point of Q2 is also that of the tractable surrogate X X (t) X (t) (t) ⊤ ⊤ e Q2 (γ, ξ, ν) := wijc γic + ξic hνc (xij ) − log wi′ j ′ c exp γi′ c + ξi′ c hνc (xij ) . i,j,c
i′ ,j ′ ,c
i,j
e (t) instead, giving the update γ (t+1) , ξ (t+1) , ν (t+1) = arg max Q e (t) (γ, ξ, ν). We optimize Q 2 2 2.3
Federated Training and Inference
Algorithm 1 in App. B presents FedSPM, which extends the centralized EM procedure to the clientserver setting. As in standard FL, FedSPM follows a local-update-then-aggregate paradigm. In communication round t, the server broadcasts the latest shared parameters (θ(t) , γ (t) , ξ (t) , ν (t) ) to all clients. Each client i then performs a local E-step followed by a local M-step: (t)
Local E-step. Client i computes {wijc }j∈[ni ],c∈[C] via (7) using only its own dataset Di . Local M-step. After the local E-step, client i updates the three parameter blocks separately. First, since the mixing weights πi are client-specific, client i updates them locally via (9) without server (t+1) (t+1) (t+1) aggregation. For the supervised block, client i updates (αi , βi , θi ) by maximizing X (t) (t) Y |X Q1i (α, β, θ) := wijc log Pic ({yij } | xij ), j,c (t) which depends only on Di and {wijc }j∈[ni ],c∈[C] available on client i. For the DRM block, client i (t+1) (t+1) (t+1) updates (γi , ξi , νi ) by maximizing
e (t) (γ, ξ, ν) := Q 2i
X
(t) ⊤ wijc γic + ξic hνc (xij ) −
j,c (t)
X
log
X
(t) τi′ c exp
γi′ c + ξi⊤′ c hνc (xij )
,
i′ ,c
j
Pni′
(t) ′ j ′ =1 wi′ j ′ c denotes the total responsibility assigned to component c on client i . (t) e (t) using only its local information, because Q e (t) involves Unlike Q1i , client i cannot evaluate Q 2i 2i
where τi′ c :=
responsibilities from other clients. However, this cross-client dependence enters only through the Pn ′ (t) summary statistics { j ′i=1 wi′ j ′ c }i′ ∈[m],c∈[C] . Accordingly, after the local E-step, each client i (t)
(t)
transmits {τic }c∈[C] to the server, and the server broadcasts the collection {τic }i∈[m],c∈[C] to all clients. This additional communication consists of only m × C scalars, which is negligible relative to model transmission and does not reveal raw data or per-sample responsibilities. Armed with these (t) e (t) via any gradient-based LocalSolver. summary statistics, client i can maximize both Q1i and Q 2i 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 ρi as the prior and the client feature density relative to the baseline distribution G as the likelihood: X ∗ X ⊤ i := arg max ρi · dPi /dG(x) = arg max ρi πic exp γic + ξic hνc (x) . (10) i∈[m]
i∈[m]
5
c
Z|X
Let Pi denote the conditional distribution of the latent component given the query on client i. The selected client i∗ then makes the final prediction: X Y |X Y |X Z|X := yb arg max Pi∗ ({k} | x) = arg max Pi∗ c ({k} | x)Pi∗ ({c} | x) . (11) k∈[K]
2.4
k∈[K]
c
Convergence Analysis
We analyze the convergence of FedSPM when LocalSolver is instantiated as local SGD [36] with momentum. Since the mixing weights and the supervised block follow standard EM/GEM updates and inherit the convergence guarantees established by [26], we focus exclusively on the DRM block ϕ := (γ, ξ, ν). For convenience, we cast the optimization of ϕ as a standard minimization problem: F (ζ) := −pℓ(ζ),
(t)
e (t) (ϕ), fe(t) (ϕ) := −Q 2
f (t) (ϕ) := −Q2 (ϕ),
(t) e (t) (ϕ). fei (ϕ) := −Q 2i
Assume that F is bounded below by F ∗ . In communication round t, each client i performs (t)
ϕi,0 := ϕ(t) ,
(t)
di,−1 := 0,
(t)
(t)
(t)
di,e := µdi,e−1 + gei,e ,
(t)
(t)
(t)
ϕi,e+1 := ϕi,e − ηdi,e ,
for e ∈ {0, . . . , E − 1}, where µ ∈ [0, 1) denotes the momentum, E denotes the number of local (t) steps, η denotes the learning rate, and gei,e denotes the stochastic gradient. (t)
Assumption 1 (Smoothness). For all i and t, fei
is L-smooth. (t)
Assumption 2 (Stochastic gradient). For all i, t, and e, gei,e is unbiased with bounded variance σ 2 , (t) (t) (t) (t) (t) (t) (t) 2 (t) i.e., E gei,e ϕi,e = ∇fei (ϕi,e ) and E gei,e − ∇fei (ϕi,e ) ϕi,e ≤ σ 2 . Assumption 3 (Inter-client heterogeneity). There exist constants Γ0 , Γ1 ≥ 0 such that, for all t, Pm e(t) (t) 2 ≤ Γ0 + Γ1 ∇fe(t) (ϕ(t) ) 2 . i=1 ρi ∇fi (ϕ ) Assumptions 1–3 are standard in federated optimization [38, 26]. Assumption 4 (Gradient bridge). There exist a constant Γ2 ≥ 0 and a non-negative sequence 2 2 {ε(t) }t≥0 such that, for all t, E ∇f (t) (ϕ(t) ) ≤ Γ2 E ∇fe(t) (ϕ(t) ) + ε(t) . Assumption 5 (Function-value bridge). There exists a non-negative sequence {δ (t) }t≥0 such that, for all t, E f (t) (ϕ(t+1) ) − f (t) (ϕ(t) ) ≤ E fe(t) (ϕ(t+1) ) − fe(t) (ϕ(t) ) + δ (t) . Assumption 4 requires that, along the iterates, the gradient norm of the DRM objective f (t) is bounded by that of its tractable surrogate fe(t) up to an error ε(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. √ √ Theorem 1. Under Assumptions 1–5, when η = Θ(1/ T ) and ηLaE < min{1/3, 1/ 3 + 8Γ1 }, FedSPM under local SGD with momentum satisfies: ! T −1 m 2 1 X Γ2 Γ2 (F (ζ (0) ) − F ∗ ) Γ2 Lσ 2 sE X 2 Γ2 L2 (a2E Γ0 + σ 2 sE ) (t) √ √ + δ̄T + ρi + + ε̄T , E ∇ϕ F (ζ ) = O T t=0 aE T aE T aE T i=1
where µe,s :=
E−1 E−1 T −1 T −1 X X 1 X (t) 1 − µe−s 1 X (t) , aE := δ , ε̄T := µE,e , sE := µ2E,e , δ̄T := √ ε . 1−µ T t=0 T t=0 e=0 e=0
Crucially, although FedSPM optimizes the tractable surrogate fe(t) , Theorem 1 establishes convergence PT −1 toward a stationary point of the exact profiled objective F . For fixed E, if t=0 δ (t) = O(1) and √ √ PT −1 (t) = O( T ), the surrogate errors do not change the standard O(1/ T ) convergence rate. t=0 ε The bound further characterizes the effect of momentum through the accumulation factors aE and 2 sE . When µ = 0, aE = sE = E, recovering the usual local SGD scaling. As µ√→ 1, aE = O(E √ ) 3 2 and sE = O(E ). Consequently, the initial-gap term √ improves √ from order 1/(E T ) to 1/(E T ), whereas the stochastic term grows from order 1/ T to E/ T , revealing a trade-off between faster optimization and amplified stochastic noise. The proof is deferred to App. C.6. 6
3
Experiments on Benchmark Datasets
We conduct experiments on three image classification benchmarks of increasing complexity: FMNIST [42], CIFAR-10 [18], and CIFAR-100 [18], with dataset details deferred to App. D.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: globalmodel methods (FedAvg [27] and FedProx [22]), a fine-tuning-based method (FedAvgFT [39]), a regularization-based method (Ditto [23]), a cluster-based method (ClusterFL [10]), a representationlearning-based method (FedBABU [30]), mixture-model-based methods (FedEM [26] and FedGMM [41]), and a routing-based method (FedDRM [40]). Notably, only FedGMM, FedDRM, and FedSPM are capable of client routing by design. Evaluation metrics. Following [40], 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θc with ResNet [11] and the routing encoder hνc with a lightweight CNN [19]. 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×), the encoders are shared across components, i.e., gθc ≡ gθ and hνc ≡ hν for all c. In FedSPM (C×), each component has its own encoders gθc and hνc . Dual heterogeneity settings. Since benchmark datasets do not inherently exhibit statistical heterogeneity, we explicitly construct dual heterogeneity based on the mixture representation P X,Y PiX,Y = , following common practice [37, 41, 40]. Specifically, we first partition c πic Pic the full dataset across 8 clients via class-wise Dirichlet partitioning [45] with concentration parameter α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 πi sampled from a Dirichlet X distribution with αintra = 2.0. To induce component-wise covariate shift in Pic , we apply transformations at two levels: client-level transformations, which introduce inter-client heterogeneity through combinations of color shifts and spatial intensity biases [40]; and component-level transformations, which introduce intra-client heterogeneity through color-channel inversion [41]. In addition, for each client-component pair, we randomly generate a label permutation [41] to induce component-wise Y |X concept shift in Pic . 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. D.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. D.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 systemlevel 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 componentlevel personalization. Fourth, FedSPM (1×) slightly outperforms FedSPM (C×), 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. 7
Table 1: System and average accuracies on benchmark datasets. Results are reported as mean ± standard deviation over 3 random seeds. Superscripts (1), (2), and (3) denote the first-, second-, and third-best results, respectively. 1× and C× denote the model capacity. System Accuracy
Method
Average Accuracy
FMNIST
CIFAR-10
CIFAR-100
FMNIST
CIFAR-10
CIFAR-100
FedAvg (1×) FedProx (1×) FedAvgFT (1×) Ditto (1×) ClusterFL (C×) FedBABU (1×) FedEM (C×) FedGMM (C×) FedDRM (1×)
39.15 ± 1.31 39.14 ± 1.32 35.42 ± 1.90 34.30 ± 2.42 31.36 ± 2.20 34.27 ± 1.19 29.73 ± 2.68 31.98 ± 1.99 53.68 ± 1.01(3)
35.02 ± 0.32 34.94 ± 0.36 32.71 ± 0.22 33.12 ± 0.45 28.38 ± 2.48 33.39 ± 0.55 26.71 ± 2.07 25.36 ± 0.48 47.47 ± 0.09(3)
18.17 ± 0.53 18.15 ± 0.49 16.29 ± 0.69 16.02 ± 0.69 14.68 ± 1.59 16.64 ± 0.25 13.69 ± 0.62 11.97 ± 0.96 27.15 ± 0.37(3)
39.15 ± 1.31 39.14 ± 1.32 50.55 ± 1.04 54.87 ± 1.18 55.77 ± 2.32(3) 38.36 ± 1.76 38.79 ± 0.19 38.34 ± 0.58 53.69 ± 1.01
35.02 ± 0.32 34.94 ± 0.36 41.40 ± 0.91 43.91 ± 0.26 45.33 ± 1.20 36.90 ± 0.47 38.50 ± 0.42 38.86 ± 0.42 47.62 ± 0.11(3)
18.17 ± 0.53 18.15 ± 0.49 23.51 ± 0.54 25.06 ± 0.54 27.31 ± 1.26 19.09 ± 0.62 19.98 ± 0.44 20.70 ± 0.25 27.55 ± 0.37(3)
FedSPM (1×) FedSPM (C×)
58.82 ± 0.61(1) 48.39 ± 0.21(1) 28.33 ± 0.37(1) 58.86 ± 0.62(1) 48.59 ± 0.22(1) 28.84 ± 0.40(1) 58.23 ± 0.17(2) 48.03 ± 0.58(2) 27.54 ± 0.45(2) 58.32 ± 0.18(2) 48.32 ± 0.61(2) 28.07 ± 0.46(2)
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 αinter ∈ {0.5, 1.0, 2.0} while fixing αintra = 2.0, and vary αintra ∈ {0.5, 1.0, 2.0} while fixing αinter = 1.0. As shown in Tab. 2, FedSPM consistently achieves the best performance across all settings, demonstrating its strong adaptability to dual heterogeneity. Table 2: System and average accuracies under varying dual heterogeneity intensities. Inter-Client Heterogeneity (αinter ) Method
System Accuracy 0.5
1.0
2.0
0.5
1.0
2.0
System Accuracy 0.5
1.0
Average Accuracy
2.0
0.5
1.0
2.0
FedAvg (1×) FedProx (1×) FedAvgFT (1×) Ditto (1×) ClusterFL (C×) FedBABU (1×) FedEM (C×) FedGMM (C×) FedDRM (1×)
49.85 39.15 31.37 49.85 39.15 31.37 42.15 37.59 39.15 42.15 37.59 39.15 49.88 39.14 31.34 49.88 39.14 31.34 42.10 37.58 39.14 42.10 37.58 39.14 45.87 35.42 27.67 62.45 50.55 45.01 36.72 33.42 35.42 54.11 50.99 50.55 45.28 34.30 25.56 65.87 54.87 49.36 34.26 31.42 34.30 58.22 55.44 54.87 40.77 31.36 23.96 68.80 55.77(3) 52.41 34.79 31.22 31.36 61.48(3) 57.66(3) 55.77(3) 42.93 34.27 27.82 50.14 38.36 31.45 37.16 35.09 34.27 40.70 39.21 38.36 37.08 29.73 25.70 49.68 38.79 31.57 33.43 28.88 29.73 41.78 39.47 38.79 34.33 31.98 24.57 49.22 38.34 31.83 28.92 27.90 31.98 44.08 38.59 38.34 (3) (3) (3) (3) (3) (3) (3) (3) 69.03 53.68 52.51 69.05 53.69 52.51 60.25 55.52 53.68 60.28 55.53 53.69
FedSPM (1×) FedSPM (C×)
74.20(1) 58.82(1) 55.68(2) 74.24(1) 58.86(1) 55.70(2) 62.83(1) 63.10(1) 58.82(1) 62.89(1) 63.20(1) 58.86(1) 73.15(2) 58.23(2) 56.57(1) 73.21(2) 58.32(2) 56.65(1) 62.48(2) 61.33(2) 58.23(2) 62.58(2) 61.40(2) 58.32(2)
101
50
̂ 1 ∑ δ(s) t s=1
t
10−1
t
25
√
1 ∑ ε(s) ̂ t s=1
102
t
1 ∑ ‖∇ F(ζ (s))‖2 t s=1 ϕ
Intra-Client Heterogeneity (αintra )
Average Accuracy
101 0
0
200
400
600
800
1000
10−3
Communication Round t μ = 0.99
0
200
400
600
800
1000
Communication Round t μ = 0.95
μ = 0.9
μ = 0.6
0
200
400
600
800
1000
Communication Round t μ = 0.3
μ=0
Figure 2: Impact of momentum µ. From left to right, the panels show the running average of ∥∇ϕ F (ζ (t) )∥2 and the normalized cumulative sums of εb(t) and δb(t) . Impact of µ. We study the impact of momentum µ under a constant learning rate on FMNIST. To align the empirical evaluation with our theory, we report the running average of ∥∇ϕ F (ζ (t) )∥2 , together with the normalized cumulative sums of the estimated bridge errors δb(t) and εb(t) . See detailed estimation procedure in App. D.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 {λic }, we numerically solve the nonlinear system in (6) for {λic } at each communication round. Fig. 2 yields three observations. First, the averaged gradient norm consistently decreases for all tested values of µ, suggesting that FedSPM empirically converges toward a stationary point of the profiled objective F . Second, as µ increases, the convergence first improves and then deteriorates, with the best performance achieved at µ = 0.95, consistent with the momentum-induced trade-off predicted 8
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 δ̄T = o(1) and ε̄T = o(1), thereby validating the soundness of Assumptions 4 and 5.
4
0.60
System Acc
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 [12, 28].
0.58 0.56 0.54 1
2
3
4
Number of Components C
5
Figure 3: Impact of C.
Experiment on Real Medical Dataset
We complement the controlled experiments on benchmark datasets with a case study on FedISIC2019 [29], 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. D.5.
64 62 60
FedAvgFT FedProx
FedBABU Ditto
Routing Acc
64
System Acc
Average Acc
66
61 58 55
85 75 65 55
FedEM ClusterFL
FedGMM FedDRM
FedSPM (C × ) FedSPM (1 × )
FedAvg
Figure 4: Average, system, and routing accuracies on Fed-ISIC2019. 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-modelbased 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×) outperforms FedSPM (1×), 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. 9
References [1] D. A. E. Acar, Y. Zhao, R. Matas, M. Mattina, P. Whatmough, and V. Saligrama. Federated learning based on dynamic regularization. In International Conference on Learning Representations, 2021. [2] J. A. Anderson. Multivariate logistic compounds. Biometrika, 66(1):17–26, 1979. [3] M. G. Arivazhagan, V. Aggarwal, A. K. Singh, and S. Choudhary. Federated learning with personalization layers. arXiv preprint arXiv:1912.00818, 2019. [4] J. C. T. Braga, A. Scope, I. Klaz, P. Mecca, P. Spencer, and A. A. Marghoob. Melanoma mimicking seborrheic keratosis: an error of perception precluding correct dermoscopic diagnosis. Journal of the American Academy of Dermatology, 58(5):875–880, 2008. [5] S. Cai, J. Chen, and J. V. Zidek. Hypothesis testing in the presence of multiple samples under density ratio models. Statistica Sinica, 27(2):761–783, 2017. [6] C. Carrera, S. Segura, P. Aguilera, M. Scalvenzi, C. Longo, A. Barreiro, P. Broganelli, S. Cavicchini, A. Llambrich, P. Zaballos, L. Thomas, J. Malvehy, S. Puig, and I. Zalaudek. Dermoscopic clues for diagnosing melanomas that resemble seborrheic keratosis. JAMA Dermatology, 153 (6):544–551, 2017. [7] L. Changchien, S. W. Dusza, A. L. C. Agero, A. J. Korzenko, R. P. Braun, D. Sachs, M. H. U. Usman, A. C. Halpern, and A. A. Marghoob. Age- and site-specific variation in the dermoscopic patterns of congenital melanocytic nevi: an aid to accurate classification and assessment of melanocytic nevi. Archives of Dermatology, 143(8):1007–1014, 2007. [8] J. Chen and Y. Liu. Quantile and quantile-function estimations under density ratio model. The Annals of Statistics, 41(3):1669–1692, 2013. [9] C. T. Dinh, N. H. Tran, and T. D. Nguyen. Personalized federated learning with moreau envelopes. In Advances in Neural Information Processing Systems, 2020. [10] A. Ghosh, J. Chung, D. Yin, and K. Ramchandran. An efficient framework for clustered federated learning. In Advances in Neural Information Processing Systems, 2020. [11] K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In IEEE Conference on Computer Vision and Pattern Recognition, 2016. [12] N. Ho, C.-Y. Yang, and M. I. Jordan. Convergence rates for gaussian mixtures of experts. Journal of Machine Learning Research, 23(323):1–81, 2022. [13] H. Hsu, H. Qi, and M. Brown. Measuring the effects of non-identical data distribution for federated visual classification. arXiv preprint arXiv:1909.06335, 2019. [14] P. Kairouz, H. B. McMahan, B. Avent, A. Bellet, M. Bennis, A. N. Bhagoji, K. Bonawitz, Z. Charles, G. Cormode, R. Cummings, R. G. L. D’Oliveira, H. Eichner, S. E. Rouayheb, D. Evans, J. Gardner, Z. Garrett, A. Gascón, B. Ghazi, P. B. Gibbons, M. Gruteser, Z. Harchaoui, C. He, L. He, Z. Huo, B. Hutchinson, J. Hsu, M. Jaggi, T. Javidi, G. Joshi, M. Khodak, J. Konečný, A. Korolova, F. Koushanfar, S. Koyejo, T. Lepoint, Y. Liu, P. Mittal, M. Mohri, R. Nock, A. Özgür, R. Pagh, H. Qi, D. Ramage, R. Raskar, M. Raykova, D. Song, W. Song, S. U. Stich, Z. Sun, A. T. Suresh, F. Tramèr, P. Vepakomma, J. Wang, L. Xiong, Z. Xu, Q. Yang, F. X. Yu, H. Yu, and S. Zhao. Advances and open problems in federated learning. Foundations and Trends in Machine Learning, 14(1–2):1–210, 2021. [15] S. P. Karimireddy, S. Kale, M. Mohri, S. J. Reddi, S. U. Stich, and A. T. Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International Conference on Machine Learning, 2020. [16] R. Kay and S. Little. Transformations of the explanatory variables in the logistic regression model for binary data. Biometrika, 74(3):495–501, 1987. [17] A. Keziou and S. Leoni-Aubin. On empirical likelihood for semiparametric two-sample density ratio models. Journal of Statistical Planning and Inference, 138(4):915–928, 2008. 10
[18] A. Krizhevsky. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009. [19] Y. LeCun, L. Bottou, Y. Bengio, and P. Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278–2324, 1998. [20] L. Li, Y. Fan, M. Tse, and K.-Y. Lin. A review of applications in federated learning. Computers & Industrial Engineering, 149:106854, 2020. [21] T. Li, A. K. Sahu, A. Talwalkar, and V. Smith. Federated learning: Challenges, methods, and future directions. IEEE Signal Processing Magazine, 37(3):50–60, 2020. [22] T. Li, A. K. Sahu, M. Zaheer, M. Sanjabi, A. Talwalkar, and V. Smith. Federated optimization in heterogeneous networks. In Proceedings of Machine Learning and Systems, 2020. [23] T. Li, S. Hu, A. Beirami, and V. Smith. Ditto: Fair and robust federated learning through personalization. In International Conference on Machine Learning, 2021. [24] X. Li, K. Huang, W. Yang, S. Wang, and Z. Zhang. On the convergence of fedavg on non-iid data. In International Conference on Learning Representations, 2020. [25] Z. Li, T. Lin, X. Shang, and C. Wu. Revisiting weighted aggregation in federated learning with neural networks. In International Conference on Machine Learning, 2023. [26] O. Marfoq, G. Neglia, A. Bellet, L. Kameni, and R. Vidal. Federated multi-task learning under a mixture of distributions. In Advances in Neural Information Processing Systems, 2021. [27] B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. y Arcas. Communication-efficient learning of deep networks from decentralized data. In International Conference on Artificial Intelligence and Statistics, 2017. [28] H. Nguyen, T. Nguyen, and N. Ho. Demystifying Softmax gating function in Gaussian mixture of experts. In Advances in Neural Information Processing Systems, 2023. [29] J. Ogier du Terrail, S.-S. Ayed, E. Cyffers, F. Grimberg, C. He, R. Loeb, P. Mangold, T. Marchand, O. Marfoq, E. Mushtaq, B. Muzellec, C. Philippenko, S. Silva, M. Teleńczuk, S. Albarqouni, S. Avestimehr, A. Bellet, A. Dieuleveut, M. Jaggi, S. P. Karimireddy, M. Lorenzi, G. Neglia, M. Tommasi, and M. Andreux. Flamby: Datasets and benchmarks for cross-silo federated learning in realistic healthcare settings. In Advances in Neural Information Processing Systems, 2022. [30] J. Oh, S. Kim, and S.-Y. Yun. Fedbabu: Toward enhanced representation for federated image classification. In International Conference on Learning Representations, 2022. [31] A. B. Owen. Empirical likelihood ratio confidence regions. The Annals of Statistics, 18(1): 90–120, 1990. [32] J. Qin. Empirical likelihood in biased sample problems. The Annals of Statistics, 21(3): 1182–1196, 1993. [33] J. Qin. Empirical likelihood ratio based confidence intervals for mixture proportions. The Annals of Statistics, 27(4):1368–1384, 1999. [34] S. J. Reddi, Z. Charles, M. Zaheer, Z. Garrett, K. Rush, J. Konečný, S. Kumar, and H. B. McMahan. Adaptive federated optimization. In International Conference on Learning Representations, 2021. [35] F. Sattler, K.-R. Müller, and W. Samek. Clustered federated learning: Model-agnostic distributed multi-task optimization under privacy constraints. IEEE Transactions on Neural Networks and Learning Systems, 32(8):3710–3722, 2021. [36] S. U. Stich. Local SGD converges fast and communicates little. In International Conference on Learning Representations, 2019. 11
[37] Y. Tan, C. Chen, W. Zhuang, X. Dong, L. Lyu, and G. Long. Is heterogeneity notorious? taming heterogeneity to handle test-time shift in federated learning. In Advances in Neural Information Processing Systems, 2023. [38] J. Wang, Q. Liu, H. Liang, G. Joshi, and H. V. Poor. Tackling the objective inconsistency problem in heterogeneous federated optimization. In Advances in Neural Information Processing Systems, 2020. [39] K. Wang, R. Mathews, C. Kiddon, H. Eichner, F. Beaufays, and D. Ramage. Federated Evaluation of On-device Personalization. arXiv preprint arXiv:1910.10252, 2019. [40] Z. Wang, X. Zhang, X. Zhang, Y. Liu, and Q. Zhang. Beyond aggregation: Guiding clients in heterogeneous federated learning. In International Conference on Learning Representations, 2026. [41] Y. Wu, S. Zhang, W. Yu, Y. Liu, Q. Gu, D. Zhou, H. Chen, and W. Cheng. Personalized federated learning under mixture of distributions. In International Conference on Machine Learning, 2023. [42] H. Xiao, K. Rasul, and R. Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv:1708.07747, 2017. [43] J. Xu, B. S. Glicksberg, C. Su, P. Walker, J. Bian, and F. Wang. Federated learning for healthcare informatics. Journal of Healthcare Informatics Research, 5(1):1–19, 2021. [44] Q. Yang, Y. Liu, T. Chen, and Y. Tong. Federated machine learning: Concept and applications. ACM Transactions on Intelligent Systems and Technology, 10(2):1–19, 2019. [45] M. Yurochkin, M. Agarwal, S. Ghosh, K. Greenewald, N. Hoang, and Y. Khazaeni. Bayesian nonparametric federated learning of neural networks. In International Conference on Machine Learning, 2019. [46] I. Zalaudek, S. Grinschgl, G. Argenziano, A. A. Marghoob, A. Blum, E. Richtig, I. H. Wolf, R. Fink-Puches, H. Kerl, H. P. Soyer, and R. Hofmann-Wellenhof. Age-related prevalence of dermoscopy patterns in acquired melanocytic naevi. British Journal of Dermatology, 154(2): 299–304, 2006. [47] B. Zhang. Quantile estimation under a two-sample semi-parametric model. Bernoulli, 6(3): 491–511, 2000. [48] B. Zhang. Assessing goodness-of-fit of generalized logit models based on case-control data. Journal of Multivariate Analysis, 82(1):17–38, 2002. [49] Q. Zhang, Q. Tian, and P. Li. Neyman-pearson multiclass classification under label noise via empirical likelihood. arXiv preprint arXiv:2603.21623, 2026. [50] Y. Zhao, M. Li, L. Lai, N. Suda, D. Civin, and V. Chandra. Federated learning with Non-IID data. arXiv preprint arXiv:1806.00582, 2018.
A
Related Work on DRM and EL
The DRM, first introduced by Anderson [2], 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 [16]. 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 [32, 33, 47, 48]. Moreover, Keziou and Leoni-Aubin [17] 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 [8, 5, 49]. 12
Recently, Wang et al. [40] 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.
B
The Federated EM Algorithm
Algorithm 1: FedSPM for t = 0, 1, . . . , T − 1 do 2 Server broadcasts (θ(t) , γ (t) , ξ (t) , ν (t) ) to all clients 3 for client i ∈ [m] in parallel do (t) 4 Compute {wijc }j∈[ni ],c∈[C] via (7) Pni (t) wijc }c∈[C] to the server 5 Send {τic = j=1 6 end 7 Server broadcasts {τic }i∈[m],c∈[C] to all clients 8 for client i ∈ [m] in parallel do (t+1) 9 Update {πic }c∈[C] via (9)
1
(t+1)
(t+1)
(t+1)
(t)
(t)
(αi , βi , θi ) ← LocalSolver(αi , βi , θ(t) ) (t+1) (t+1) (t+1) 11 (γi , ξi , νi ) ← LocalSolver(γ (t) , ξ (t) , ν (t) ) (t+1) (t+1) (t+1) (t+1) 12 Client i sends (θi , γi , ξi , νi ) to the server 13 end Pm (t+1) (t+1) (t+1) (t+1) 14 Server updates (θ(t+1) , γ (t+1) , ξ (t+1) , ν (t+1) ) ← i=1 ρi (θi , γi , ξi , νi ) 15 end 10
C
Mathematical Details
C.1
Derivation of the Log-EL
The observed data contain only (xij , yij ), while the component assignment zij is latent. Therefore, the observed-data log-EL is obtained from the augmented distribution PiX,Y,Z by marginalizing over all possible values of Z: ℓ(ζ, G) :=
X i,j
c
=
X
log
X
i,j
c
=
X
log
X
i,j
c
=
X
log
X
i,j
c
=
X
log
X
i,j
c
=
X
log rij +
i,j
log
X
PiX,Y,Z (dxij , {yij }, {c}) | {z }
complete-data contribution for Z=c X,Y πic Pic (dxij , {yij }) Y |X
X ({yij } | xij )Pic (dxij )
Y |X
({yij } | xij )
X dPic (xij )G(dxij ) dG
Y |X
({yij } | xij )
X dPic (xij )rij dG
πic Pic
πic Pic πic Pic X i,j
log
X
Y |X ⊤ πic exp γic + ξic hνc (xij ) Pic ({yij } | xij ) .
c
13
C.2
Derivation of the Optimal Baseline Weights
For fixed ζ, maximizing ℓ over G reduces to maximizing The corresponding Lagrangian is:
L :=
X
log rij − nκ
X
i,j
rij − 1 − n
X
λi′ c
i′ ,c
i,j
P
i,j log rij subject to the constraints in (4).
X exp γi′ c + ξi⊤′ c hνc (xij ) − 1 rij . i,j
Setting ∂L/∂rij = 0 gives −1 rij − nκ − n
X
λi′ c exp γi′ c + ξi⊤′ c hνc (xij ) − 1 = 0.
i′ ,c
Multiplying the above equation by rij , summing over (i, j), and applying the constraints in (4) gives κ = 1, which yields X −1 1 ∗ ⊤ rij (ζ) = 1+ λi′ c exp γi′ c + ξi′ c hνc (xij ) − 1 . n ′ i ,c
∗ Substituting rij (ζ) back into (4) gives the following nonlinear system for {λic }:
exp γi′ c + ξi⊤′ c hνc (xij ) − 1 = 0, ⊤ ′′ ′ ′′ ′ i′′ ,c′ λi c exp γi c + ξi′′ c′ hνc′ (xij ) − 1
X
P
1+ i,j C.3
∀i′ ∈ [m], c ∈ [C].
Derivation of the Complete-Data Profile Log-EL
Unlike the observed-data log-EL, the complete-data log-EL treats the latent component indicators {zijc } as observed. Thus, it uses the corresponding complete-data contribution from PiX,Y,Z directly, without marginalizing over Z: X ℓc (ζ, G) := zijc log PiX,Y,Z (dxij , {yij }, {c}) i,j,c
=
X
=
X
=
X
h i X,Y zijc log πic Pic (dxij , {yij })
i,j,c
h i Y |X X zijc log πic Pic ({yij } | xij )Pic (dxij )
i,j,c
h i dP X Y |X zijc log πic Pic ({yij } | xij ) ic (xij )G(dxij ) dG i,j,c
h dP X i Y |X zijc log πic ic (xij )rij Pic ({yij } | xij ) dG i,j,c h i X X Y |X ⊤ = log rij + zijc log Pic ({yij } | xij ) + γic + ξic hνc (xij ) + log πic . =
X
i,j
i,j,c
P Since maximizing ℓc over G for fixed ζ also reduces to maximizing i,j log rij subject to the ∗ constraints in (4), ℓc shares the same optimal baseline weights as ℓ. Thus, substituting rij (ζ) back c into ℓ (ζ, G) yields the complete-data profile log-EL: pℓc (ζ) := sup ℓc (ζ, G) G
=
X
∗ log rij (ζ) +
i,j
C.4
X
Y |X ⊤ zijc log Pic ({yij } | xij ) + γic + ξic hνc (xij ) + log πic .
i,j,c
EM Algorithm Details
The EM algorithm is an iterative procedure that starts from an initial estimate ζ (0) . At iteration t, given the current parameter estimate ζ (t) , the algorithm alternates between the following two steps. 14
E-step. The latent component assignments {zij } are estimated by their posterior responsibilities, which serve as soft guesses computed from the observed data under the current parameter estimate: h i (t) Z|X,Y wijc := E zijc D; ζ (t) = Pi ({c} | xij , yij ; ζ (t) ). The second equality follows because zijc is the indicator of the event {Z = c}. Applying Bayes’ rule (t) X,Y with prior πic and likelihood Pic (dxij , {yij }; ζ (t) ) gives (t)
π P X,Y (dxij , {yij }; ζ (t) ) (t) wijc = P ic (t)ic X,Y . (t) ) c′ πic′ Pic′ (dxij , {yij }; ζ Y |X
X,Y into Pic Using the factorization of Pic
X and Pic gives
Y |X
(t)
X π P ({yij } | xij ; ζ (t) )Pic (dxij ; ζ (t) ) (t) wijc = P ic (t)ic Y |X . (t) )P X (dx ; ζ (t) ) ij ic′ c′ πic′ Pic′ ({yij } | xij ; ζ X X Under the EL representation of G, Pic (dxij ) = dPic /dG(xij )rij . Therefore, (t) Y |X dP X πic Pic ({yij } | xij ; ζ (t) ) dGic (xij ; ζ (t) )rij (t) . wijc = P X (t) Y |X (t) ) dPic′ (x ; ζ (t) )r ij ij c′ πic′ Pic′ ({yij } | xij ; ζ dG
Finally, substituting the DRM density ratio and canceling the common factor rij yields Y |X (t) (t) (t) Pic ({yij } | xij ; ζ (t) ) exp γic + (ξic )⊤ hν (t) (xij ) πic c (t) . wijc = P Y |X (t) ) exp γ (t) + (ξ (t) )⊤ h (t) (x ) π (t) ({y } | x ; ζ P ′ ′ ′ ′ ′ ij ij ij ic ic ic ic c ν c′
Then, the resulting profile Q-function is the conditional expectation of the complete-data profile log-EL with respect to {zij }: h i Q(t) (ζ) := E pℓc (ζ) D; ζ (t) (t)
(t)
(t)
= Q1 (α, β, θ) + Q2 (γ, ξ, ν) + Q3 (π), where (t)
Q1 (α, β, θ) :=
X
Y |X
(t)
wijc log Pic
({yij } | xij ),
i,j,c (t)
Q2 (γ, ξ, ν) :=
X
(t) ⊤ wijc γic + ξic hνc (xij )
i,j,c
−
X
X ⊤ ′ ′ log 1 + λi c exp γi c + ξi′ c hνc (xij ) − 1 , i′ ,c
i,j (t) Q3 (π) :=
X
(t) wijc log πic .
i,j,c
M-step. Instead of directly maximizing the profile log-EL, the M-step maximizes Q(t) with respect to ζ, i.e., ζ (t+1) = arg max Q(t) (ζ). Here we focus on the analytic update of the mixing weights π, which solves X (t) max wijc log πic π
s.t.
ijc
X
i ∈ [m],
πic = 1,
c
πic ≥ 0,
i ∈ [m], 15
c ∈ [C].
Using the method of Lagrange multipliers, the Lagrangian is ! Lπ :=
X
(t) wijc log πic +
X
ijc
ιi
X
πic − 1 .
c
i
At the optimum, we require ∂Lπ = 0, ∂ιi
∂Lπ = 0, ∂πic Solving these equations gives
i ∈ [m],
c ∈ [C].
n
(t+1)
πic C.5
=
i 1 X (t) w . ni j=1 ijc
Derivation of the Analytical Form of the Lagrange Multipliers (t)
We derive the analytical form of {λic } from the first-order optimality condition of Q2 . For each i0 ∈ [m] and c0 ∈ [C], taking the derivative with respect to γi0 ,c0 gives ni0 (t) X X exp γi0 ,c0 + ξi⊤0 ,c0 hνc0 (xij ) ∂Q2 (t) , P wi0 ,j,c0 − λi0 ,c0 = ∂γi0 ,c0 1 + i′ ,c λi′ c exp γi′ c + ξi⊤′ c hνc (xij ) − 1 j=1 i,j where the term involving ∂λi0 ,c0 /∂γi0 ,c0 vanishes due to (6). By (4) and (5), we further have X X ∗ exp γi0 ,c0 + ξi⊤0 ,c0 hνc0 (xij ) =n P exp γi0 ,c0 + ξi⊤0 ,c0 hνc0 (xij ) rij = n. ⊤ 1 + i′ ,c λi′ c exp γi′ c + ξi′ c hνc (xij ) − 1 i,j i,j Therefore,
ni0 (t) X ∂Q2 (t) = wi0 ,j,c0 − nλi0 ,c0 . ∂γi0 ,c0 j=1
Setting the derivative to zero yields ni
0 1X (t) . λi0 ,c0 = w n j=1 i0 ,j,c0
C.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 {µE,e }E−1 e=0 , which make the dependence on µ explicit through aE and sE . 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 fe(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. C.6.1
Additional Notations
For the convergence proof, define the stochastic gradient noise as (t) (t) (t) (t) ωi,e := gei,e − ∇fei (ϕi,e ).
Let F (−1) be the trivial σ-algebra. For t ≥ 0, define (s) F (t) := σ ωi,e : 0 ≤ s ≤ t, i ∈ [m], 0 ≤ e ≤ E − 1 . For each client i and local step e, define (t)
Fi,−1 := F (t−1) ,
(t) (t) (t) Fi,e := F (t−1) ∨ σ ωi,0 , . . . , ωi,e . 16
C.6.2
Key Lemmas
Lemma 1 (EM function-value bridge). At communication round t, for any ζ, we have F (ζ) − F (ζ (t) ) ≤ f (t) (ϕ) − f (t) (ϕ(t) ). Proof. We first prove the evidence lower bound (ELBO) for the profiled log-EL: X X X Y |X ∗ ⊤ pℓ(ζ) = log rij (ζ) + log πic exp γic + ξic hνc (xij ) Pic ({yij } | xij ) . i,j
(t) c wijc = 1, Jensen’s inequality gives
For each sample (i, j), using X
X
pℓ(ζ) =
∗ log rij (ζ) +
i,j
≥
X
c
i,j
P
log
X c
i,j ∗ log rij (ζ) +
i,j
X
(t)
wijc log
Y |X ⊤ γic + ξic hνc (xij ) Pic ({yij } | xij )
(t) πic exp
wijc
(t)
wijc Y |X ⊤ πic exp γic + ξic hνc (xij ) Pic ({yij } | xij ) (t)
wijc
i,j,c
= Q(t) (ζ) −
X
(t)
(t)
wijc log wijc .
i,j,c
By (7), when ζ = ζ
(t)
, the Jensen inequality above becomes tight. Hence, X (t) (t) pℓ(ζ (t) ) = Q(t) (ζ (t) ) − wijc log wijc . i,j,c
Subtracting this equality from the previous lower bound gives pℓ(ζ) − pℓ(ζ (t) ) ≥ Q(t) (ζ) − Q(t) (ζ (t) ). (t)
(t)
By the monotone ascent property of the standard EM/GEM updates for Q1 and Q3 , we have (t)
(t)
Q(t) (ζ) − Q(t) (ζ (t) ) ≥ Q2 (ϕ) − Q2 (ϕ(t) ). (t)
Therefore, since F = −pℓ and f (t) = −Q2 , we obtain F (ζ) − F (ζ (t) ) ≤ f (t) (ϕ) − f (t) (ϕ(t) ).
Lemma 2 (EM gradient bridge). At communication round t, we have ∇ϕ F (ζ (t) ) = ∇f (t) (ϕ(t) ). Proof. By the ELBO for the profiled log-EL established in the proof of Lemma 1, for any ζ, pℓ(ζ) − pℓ(ζ (t) ) ≥ Q(t) (ζ) − Q(t) (ζ (t) ). Equivalently, pℓ(ζ) − Q(t) (ζ) ≥ pℓ(ζ (t) ) − Q(t) (ζ (t) ). Therefore, pℓ(ζ) − Q(t) (ζ) attains a local minimum at ζ (t) . Since the involved functions are differentiable, its ϕ-gradient vanishes at ζ (t) , namely ∇ϕ pℓ(ζ (t) ) = ∇ϕ Q(t) (ζ (t) ). (t)
By (8), only Q2 depends on ϕ. Thus, (t)
∇ϕ Q(t) (ζ (t) ) = ∇Q2 (ϕ(t) ). (t)
Since F = −pℓ and f (t) = −Q2 , we obtain (t)
∇ϕ F (ζ (t) ) = −∇ϕ pℓ(ζ (t) ) = −∇Q2 (ϕ(t) ) = ∇f (t) (ϕ(t) ).
17
Lemma 3 (Momentum unrolling identity). At communication round t, we have ϕ(t+1) = ϕ(t) − ηe g (t) , where ge(t) :=
m X
ρi
E−1 X
(t)
µE,e gei,e .
e=0
i=1
Proof. By unrolling the momentum recursion, for every e = 0, . . . , E − 1, (t) (t) (t) di,e = µdi,e−1 + gei,e =
e X
(t)
µe−s gei,s .
s=0
Therefore, (t)
ϕi,E − ϕ(t) = −η
E−1 X
(t)
di,e
e=0
= −η
E−1 e XX
(t)
µe−s gei,s
e=0 s=0
= −η
= −η
E−1 X
E−1 X
s=0
e=s
E−1 X
! e−s
µ
(t)
gei,s
(t)
µE,s gei,s .
s=0
Using the aggregation rule and
Pm
i=1 ρi = 1 gives
ϕ(t+1) − ϕ(t) =
m X
(t) ρi ϕi,E − ϕ(t)
i=1
= −η
m X
ρi
E−1 X
(t)
µE,e gei,e
e=0
i=1 (t)
= −ηe g .
Lemma 4 (Gradient decomposition). At communication round t, we have ge(t) = aE ∇fe(t) (ϕ(t) ) + D(t) + Ξ(t) , where D(t) :=
m X i=1
ρi
E−1 X
h i (t) (t) (t) µE,e ∇fei (ϕi,e ) − ∇fei (ϕ(t) ) ,
e=0
and Ξ(t) :=
m X
ρi
E−1 X
(t)
µE,e ωi,e .
i=1
e=0
E−1 X
(t)
Proof. By definition, ge(t) = =
m X
ρi
i=1
e=0
m X
E−1 X
i=1
ρi
µE,e gei,e
(t) (t) µE,e ∇fei (ϕi,e ) + Ξ(t) .
e=0
18
(t) Adding and subtracting ∇fei (ϕ(t) ) gives
ge(t) =
m X
ρi
(t)
µE,e ∇fei (ϕ(t) ) + D(t) + Ξ(t)
e=0
i=1
=
E−1 X
E−1 X
µE,e
! m X
e=0
(t)
ρi ∇fei (ϕ(t) ) + D(t) + Ξ(t)
i=1
= aE ∇fe(t) (ϕ(t) ) + D(t) + Ξ(t) .
Lemma 5 (Stochastic noise bound). Under Assumption 2, at communication round t, we have m X 2 F (t−1) ≤ σ 2 sE ρ2i . E Ξ(t) i=1 (t)
Proof. For each client i, {ωi,e }E−1 e=0 is a martingale difference sequence with h i 2 (t) (t) (t) (t) Fi,e−1 ≤ σ 2 . E ωi,e Fi,e−1 = 0, E ωi,e Thus the within-client cross terms vanish, and 2 E−1 E−1 X X 2 (t) (t) F (t−1) E µE,e ωi,e F (t−1) = µ2E,e E ωi,e e=0
e=0
≤ σ 2 sE . Since the accumulated noises are conditionally independent across clients given F (t−1) , the crossclient terms also vanish. Hence, 2 X m E−1 X 2 (t) E Ξ(t) F (t−1) = ρ2i E µE,e ωi,e F (t−1) e=0
i=1 m X
≤ σ 2 sE
ρ2i .
i=1
Lemma 6 (Gradient drift bound). Under Assumptions 1–3, at communication round t, if ηLaE < √ 1/ 3, we have 2 2 2a2E L2 η 2 2 2 e(t) (ϕ(t) ) E D(t) a Γ + Γ F (t−1) ≤ ∇ f + s σ . 0 1 E 1 − 3a2E L2 η 2 E Proof. By Jensen’s inequality and L-smoothness, D(t)
2
≤ aE
m X
ρi
i=1
≤ aE L2
m X
E−1 X
Me(t) :=
m X
ρi E
(t)
(t)
2
e=0
ρi
i=1
Set
(t)
µE,e ∇fei (ϕi,e ) − ∇fei (ϕ(t) )
(t) ϕi,e − ϕ(t)
E−1 X
(t)
µE,e ϕi,e − ϕ(t)
2
.
e=0
2
F
(t−1)
i=1
19
,
(t) := Mmax
max
0≤e≤E−1
Me(t) .
Then
D(t)
E
2
(t) F (t−1) ≤ a2E L2 Mmax .
(t)
It remains to bound Mmax . For e ≥ 1, (t)
ϕi,e − ϕ(t) = −η
e−1 X
(t)
µe,s gei,s .
s=0
Using h i (t) (t) (t) (t) (t) (t) gei,s = ∇fei (ϕ(t) ) + ∇fei (ϕi,s ) − ∇fei (ϕ(t) ) + ωi,s , martingale cancellation, Young’s inequality with equal weights, and L-smoothness imply !2 m e−1 X X 2 (t) (t) 2 M ≤ 2η µe,s ρi ∇fe (ϕ(t) ) e
i
s=0 2
2
+ 3η L
i=1 e−1 X
µe,s
s=0
! e−1 X
µe,s Ms(t) + 2η 2 σ 2
s=0
e−1 X
µ2e,s .
s=0
By Assumption 3, m X
(t)
ρi ∇fei (ϕ(t) )
2
≤ Γ0 + Γ1 ∇fe(t) (ϕ(t) )
2
.
i=1
Since e−1 X
e−1 X
µe,s ≤ aE ,
s=0
µ2e,s ≤ sE ,
(t) Ms(t) ≤ Mmax ,
s=0
we get Me(t) ≤ 2η 2 a2E
Γ0 + Γ1
∇fe(t) (ϕ(t) )
2
(t) + 2η 2 sE σ 2 + 3η 2 L2 a2E Mmax .
Taking the maximum over e gives 2 2η 2 a2E Γ0 + Γ1 ∇fe(t) (ϕ(t) ) + sE σ 2 (t) Mmax ≤
.
1 − 3a2E L2 η 2
Substituting this into the bound for D(t) proves the result. Lemma 7 (One-round descent). Under Assumptions 1–3, at communication round t, if ηLaE < √ min{1/3, 1/ 3 + 8Γ1 }, we have h i E fe(t) (ϕ(t+1) ) − fe(t) (ϕ(t) ) F (t−1) m
≤−
X 2 aE η ∇fe(t) (ϕ(t) ) + sE Lη 2 σ 2 ρ2i + 3aE L2 η 3 a2E Γ0 + sE σ 2 . 4 i=1
Proof. By Lemma 3 and L-smoothness, D E Lη 2 2 fe(t) (ϕ(t+1) ) − fe(t) (ϕ(t) ) ≤ −η ∇fe(t) (ϕ(t) ), ge(t) + ge(t) . 2 Using Lemma 4, expanding the square, and taking conditional expectation, the terms involving ⟨∇fe(t) (ϕ(t) ), Ξ(t) ⟩ vanish. Since 2⟨D(t) , Ξ(t) ⟩ ≤ D(t) 20
2
+ Ξ(t)
2
,
we obtain
i h E fe(t) (ϕ(t+1) ) − fe(t) (ϕ(t) ) F (t−1) 2 2 La2E η 2 ≤ −aE η ∇fe(t) (ϕ(t) ) + ∇fe(t) (ϕ(t) ) 2 i E hD (t) (t) 2 e + −η + LaE η E ∇f (ϕ ), D(t) F (t−1) 2 2 F (t−1) . F (t−1) + Lη 2 E Ξ(t) + Lη 2 E D(t)
Since aE Lη < 1/3, −η + LaE η 2 = −η(1 − aE Lη) ≤ 0. Using
we have
E a D 2 2 1 E − ∇fe(t) (ϕ(t) ), D(t) ≤ ∇fe(t) (ϕ(t) ) + D(t) , 2 2aE i h E fe(t) (ϕ(t+1) ) − fe(t) (ϕ(t) ) F (t−1) 2 2 aE η η Lη 2 ≤− ∇fe(t) (ϕ(t) ) + F (t−1) + E D(t) 2 2aE 2 2 F (t−1) . + Lη 2 E Ξ(t)
By Lemmas 5 and 6, 2 η Lη 2 F (t−1) E D(t) + 2aE 2 2 aE L2 η 3 (1 + aE Lη) 2 (t) (t) 2 e ≤ a Γ + Γ ∇ f (ϕ ) + s σ , 0 1 E E 1 − 3a2E L2 η 2 and 2
Lη E
(t)
Ξ
2
F
(t−1)
2 2
≤ sE Lη σ
m X
ρ2i .
i=1
The drift contribution to the gradient coefficient equals aE η
2 Γ1 (aE Lη)2 (1 + aE Lη) ∇fe(t) (ϕ(t) ) . 2 2 2 1 − 3aE L η
The learning rate condition implies Γ1 (aE Lη)2 (1 + aE Lη) 1 ≤ , 2 2 2 1 − 3aE L η 4
1 + aE Lη ≤ 3. 1 − 3a2E L2 η 2
Combining the above inequalities proves the lemma. C.6.3
Proof of Theorem 1
Proof. By Lemma 7 and the tower property, h i E fe(t) (ϕ(t+1) ) − fe(t) (ϕ(t) ) m X 2 aE η (t) (t) e E ∇f (ϕ ) ≤− + sE Lη 2 σ 2 ρ2i + 3aE L2 η 3 a2E Γ0 + sE σ 2 . 4 i=1 Combining the preceding inequality with Lemma 1 and Assumption 5, we obtain h i 2 aE η E ∇fe(t) (ϕ(t) ) ≤ E F (ζ (t) ) − F (ζ (t+1) ) + δ (t) 4 m X 2 2 + sE Lη σ ρ2i + 3aE L2 η 3 a2E Γ0 + sE σ 2 . i=1
21
Summing over t = 0, . . . , T − 1 and using F (ζ (T ) ) ≥ F ∗ , T −1 T −1 X 2 aE η X (t) (t) e ≤ F (ζ (0) ) − F ∗ + E ∇f (ϕ ) δ (t) 4 t=0 t=0 + T sE Lη 2 σ 2
m X
ρ2i + 3T aE L2 η 3 a2E Γ0 + sE σ 2 .
i=1
Thus, T −1 X
1 E T t=0
∇fe(t) (ϕ(t) )
2
≤
PT −1 4 F (ζ (0) ) − F ∗ + t=0 δ (t) aE ηT m
+
4sE Lησ 2 X 2 ρi + 12L2 η 2 a2E Γ0 + sE σ 2 . aE i=1
By Lemma 2 and Assumption 4, 2 2 + ε(t) . ≤ Γ2 E ∇fe(t) (ϕ(t) ) E ∇ϕ F (ζ (t) ) Averaging over t and substituting the preceding bound gives 4Γ F (ζ (0) ) − F ∗ + PT −1 δ (t) T −1 2 2 t=0 1 X E ∇ϕ F (ζ (t) ) ≤ T t=0 aE ηT m
+
4Γ2 sE Lησ 2 X 2 ρi aE i=1
−1 1 TX ε(t) . + 12Γ2 L2 η 2 a2E Γ0 + sE σ 2 + T t=0
√ Taking η = Θ(1/ T ) yields m T −1 2 1 X Γ2 (F (ζ (0) ) − F ∗ ) Γ2 Γ2 Lσ 2 sE X 2 √ √ + ρi E ∇ϕ F (ζ (t) ) = O δ̄T + T t=0 aE aE T aE T i=1 ! Γ2 L2 (a2E Γ0 + σ 2 sE ) + ε̄T , + T
which completes the proof.
D
Experimental Details
D.1
Benchmark Dataset Details
We conduct experiments on FMNIST [42], CIFAR-10 [18], and CIFAR-100 [18]. FMNIST contains 70,000 28 × 28 grayscale images from 10 fashion classes. CIFAR-10 contains 60,000 32 × 32 RGB images from 10 object classes. CIFAR-100 contains 60,000 32 × 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. D.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 22
inversion for FMNIST and green-channel inversion for CIFAR-10/100. At the client level, the 8 clients correspond to the 23 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 αinter on client label distributions. A smaller αinter produces increasingly skewed class proportions and uneven dataset sizes across clients, corresponding to intensified interclient heterogeneity. Similarly, Fig. 7 visualizes how αintra modulates the intra-client mixing weights. While a small αintra concentrates local data into a single dominant component, a larger αintra yields more balanced mixtures, indicating stronger intra-client heterogeneity.
client:0 client:1 client:2 client:3 client:4 client:5 client:6 client:7 component:0 component:0 component:0 component:0 component:0 component:0 component:0 component:0
client:0 client:1 client:2 client:3 client:4 client:5 client:6 client:7 component:1 component:1 component:1 component:1 component:1 component:1 component:1 component:1
inter = 0.5
5000
10000
15000
7 6 5 4 3 2 1 0 0
class0
class1
class2
class3
Sample Num
inter = 1.0
5000
10000
Sample Num
class4
5000
10000
15000
class6
class7
class8
class9
15000
class5
inter = 2.0
7 6 5 4 3 2 1 0 0
Client
7 6 5 4 3 2 1 0 0
Client
Client
Figure 5: Visualization of an FMNIST sample under client- and component-level transformations.
Sample Num
intra = 0.5
7 6 5 4 3 2 1 0 15000 0
intra = 1.0
5000
10000
Sample Num
7 6 5 4 3 2 1 0 15000 0
intra = 2.0
Client
7 6 5 4 3 2 1 0 0
Client
Client
Figure 6: Visualization of client label distributions under varying αinter .
5000
10000
Sample Num
component0
component1
5000
10000
Sample Num
15000
Figure 7: Visualization of client mixing weights under varying αintra . D.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 23
of local fine-tuning before evaluation. For cluster- and mixture-model-based methods, we set the number of components C to 3. D.4
Estimation of the Bridge Errors
This section details the empirical estimation of the bridge errors ε(t) and δ (t) reported in Fig. 2. Computing these quantities requires evaluating the DRM objective f (t) and its gradient ∇f (t) , both of which depend on the Lagrange multipliers {λic }. Therefore, at each communication round, we numerically solve the nonlinear system in (6) for {λic } using a Jacobian-based root solver, followed by a trust-region least-squares refinement when necessary. The resulting multipliers are then used to −1 −1 compute {f (t) (ϕ(t) )}Tt=0 and {∇f (t) (ϕ(t) )}Tt=0 along the training trajectory. Estimation of the gradient bridge error ε(t) . Assumption 4 gives 2 2 (t) (t) (t) (t) e ≤ Γ2 E ∇f (ϕ ) + ε(t) . E ∇f (ϕ ) Motivated by this inequality, we first estimate the coefficient Γ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: 2 Tvalid := t : ∇fe(t) (ϕ(t) ) > 10−9 . We then estimate Γ2 by the 95% quantile of the valid ratios to ensure robustness against outliers: ∇f (t) (ϕ(t) ) 2 b Γ2 := Quantile0.95 : t ∈ T . valid 2 ∇fe(t) (ϕ(t) ) b 2 , ε(t) is computed as the remaining positive gap: Finally, given Γ 2 2 b 2 ∇fe(t) (ϕ(t) ) , 0 . εb(t) := max ∇f (t) (ϕ(t) ) − Γ Estimation of the function-value bridge error δ (t) . Assumption 5 gives h i h i E f (t) (ϕ(t+1) ) − f (t) (ϕ(t) ) ≤ E fe(t) (ϕ(t+1) ) − fe(t) (ϕ(t) ) + δ (t) . Accordingly, we estimate δ (t) as h i h i δb(t) := max f (t) (ϕ(t+1) ) − f (t) (ϕ(t) ) − fe(t) (ϕ(t+1) ) − fe(t) (ϕ(t) ) , 0 . D.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.
24