Beyond Aggregation: Guiding Clients in Heterogeneous Federated Learning
Authors: Zijian Wang
Xiaofei Zhang
Xin Zhang
Yukun Liu
Qiong Zhang
Venue: ICLR 2026 Poster (2026)
Keywords: client allocation
density ratio model
empirical likelihood
healthcare
patient routing
server routing
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: In federated learning, not all clients are equally useful. We automatically route each query to the most relevant clients, improving accuracy while preserving data privacy.

Abstract: Federated learning (FL) is increasingly adopted in domains like healthcare, where data privacy is paramount. 
A fundamental challenge in these systems is statistical heterogeneity—the fact that data distributions vary significantly across clients (e.g., different hospitals may treat distinct patient demographics). 
While current FL algorithms focus on aggregating model updates from these heterogeneous clients, the potential of the central server remains under-explored.
This paper is motivated by a healthcare scenario: could a central server not only coordinate model training but also guide a new patient to the hospital best equipped for their specific condition? We generalize this idea to propose a novel paradigm for FL systems where the server actively guides the allocation of new tasks or queries to the most appropriate client.
To enable this, we introduce a density ratio model and empirical likelihood-based framework that simultaneously addresses two goals: (1) learning effective local models on each client, and (2) finding the best matching client for a new query. 
Empirical results demonstrate the framework's effectiveness on benchmark datasets, showing improvements in both model accuracy and the precision of client guidance compared to standard FL approaches. This work opens a new direction for building more intelligent and resource-efficient FL systems that leverage heterogeneity as a feature, not just a bug.

Published as a conference paper at ICLR 2026
BEYOND AGGREGATION: GUIDING CLIENTS IN HET- EROGENEOUS FEDERATED LEARNING
Zijian Wang'
Xiaofei Zhang
Xin Zhang
Yukun Liu
Qiong Zhang!*
'Remmin University of China
Zhongnan University of Economics and Law
8 Meta
*Fast China Normal University
ABSTRACT
Federated learning (FL) is increasingly adopted in domains like healthcare, where data privacy is paramount. A fundamental challenge in these systems is statisti- cal heterogeneity- the fact that data distributions vary significantly across clients (e.g.. different hospitals may treat distinct patient demographics). While cur- rent FL algorithms focus on aggregating model updates from these heterogeneous clients, the potential of the central server remains under-explored. This paper is motivated by a healthcare scenario: could a central server not only coordinate model training but also guide new patient to the hospital best equipped for their specific condition? We generalize this idea to propose a novel paradigm for FL systems where the server actively guides the allocation of new tasks or queries to the most appropriate client. To enable this, we introduce a density ratio model and empirical likelihood-bol
framework that simultaneously addresses two goals: (I) learning effective local models on each client, and (2) finding the best match- ing client for new query. Empirical results demonstrate the framework's effec- tiveness on benchmark datasets, showing improvements in both model accuracy and the precision of client guidance compared to standard FL approaches. This work opens a new direction for building more intelligent and resource- resource-efficient FL. systems that leverage heterogeneity as a feature, not just bug. Code is available at htttss:// com zi
INTRODUCTION
Federated learning (FL) has emerged powerful paradigm for training machine learning models across distributed data sources without sharing raw data. By enabling clients such as hospitals, financial institutions, mobile devices to collaboratively train models under the coordination of central server, FL offers a practical solution for
learning in sensitive domains (Li
al. 2020a
Long al. 2020 Xu al. 2021)
key challenge in applying FL in practice is statistical het- erogeneity: clients often hold data drawn from different, identically distributed populations. may serve distinct patient demographies; may encounter different fraud patterns; and on mobile devices, user behavior varies widely. Such heterogeneity can cause lo- cal models to drift apart, leading to slower convergence (Li al. 2020b), biased updates (Karimireddy et al., 2020), and global models that underperform when applied back to indi- vidual clients TT Dinh et al.) 2020). To address these issues, most FL systems suppress heterogeneity through aggregation corrections, client reweighting, or personalization, leaving the server largely passive as aggregator of local updates. We argue this misses a key opportunity: the server should actively exploit heterogeneity rather than merely mitigate it.
healthcare, hospitals
finance, banks
This way!
FL server as an intelligent router: Leveraging learned data distributions to di- rect queries the most special- ized client, rather than applying a global model for diagnosis.
Figure 1:
"Correspondence Qiong Zhang (qiong. zhang@ruccc.oooo wdw...edd.c.o edu.cn).

---

Published as a conference paper at ICLR 2026
Consider a healthcare scenario: different hospitals may excel at treating different patient groups depending on their location and/or expertise. When a new patient arrives, instead of merely de- ploying a global model for diagnosis, the server could help identify the hospital best equipped to provide care, leveraging local data distributions to capture specialized expertise. A cartoon illustra- tion of this scenario is given in Fig. [ Similar opportunities exist in other domains: in finance, the server could direct a fraud detection query to the bank whose historical data best matches the trans- action profile; in personalized services, it could route query to the client with the most relevant user base. These examples illustrate that statistical heterogeneity across clients-off seen as an obstacle can instead become a valuable resource. They motivate the central insight of our work:
Beyond coordinating training, the server can actively exploit client heterogeneiiy-ttrnsfoooom it from a challenge into resource by guiding new queries the most suitable client.
Much of the existing work in FL has focused on mitigating the challenges of statistical heterogeneity, without using the server for guiding new queries. One major line of research develops aggregation
algorithms to reduce the bias induced by non-identically distributed data. Examples include methods
that modify local updates before aggregation (Gao et al. [2022] Guo et al. 2023 Zhang al.) 2023),
reweight client contributions (Wang et al. 2020; Yin et al. 2024). or introduce regularization terms
to align local objectives with the global one (Li et al. 2020b AAarr et al., 2021 Li et al. 2021b).
These approaches aim to learn single global model that performs reasonably well across all clients, but they do not leverage heterogeneity as an asset. A second line of work explores personalization in FL. Rather than enforcing a universal global model, personalization methods adapt models to each client's local distribution (Li et al. 2021d), often through fine-tuning (T Dinh et al. 2020 Collins iet al., 2021 Tan et al., 2022 Ma et al., 2022). multi-task learning (Smith et al.J [017]] Li et al. 2021c). or meta-learning Fallah et al. 2020). While these approaches improve local performance, they are typically not designed to address the challenge of guiding new queries tasks to the most appropriate client. Another related direction is client clustering (Ghosh et al.) 2020 Li et al., 2021a BBriggs et al. 2020 Kim et al. 202T Long et al., 2023), where clients with similar data distri- butions are grouped and trained jointly within each cluster. This can improve performance under heterogeneity, but still assumes the server's role is limited to coordinating training and distributing models, rather than supporting query routing or task allocation. Overall, while these approaches are effective for their intended goals, they stop short of enabling the server to actively guide new queries to the most suitable client.
Motivated by this gap, we introduce a new paradigm in which the FL. server not only coordinates training but also learns to guide each incoming query to the client best suited to handle it. Achieving this goal requires two capabilities: (i) effective information sharing across clients despite hetero- geneity, and (ii) principled way to quantify how each client's 's data distribution differs from the others so that queries can be meaningfully matched. To achieve this, we develop F FedDRM, uni- fied framework grounded in density ratio model (DRM) (Anderson 1979) and empirical likelihood (EL) (Owen) (001). DRM represents cach client's 's distribution as a multiplicatiie density tilt of a baseline distribution, while EL. facilitates nonparametric model learning, enabling the estimation of this baseline distribution in a data-driven manner without parametric assumptions. After profiling out the baseline distribution, the resulting objective decomposes into two interpretable cross-entropy components: one for predicting class labels and another for identifying sample's ' client of origin. The first supports standard FL training; the second supplies precisely the signal needed for query- to-client routing, enabling the server to exploit -raoher than suppress-staaisticaal heterogeneity.
This formulation leads to three key contributions. First, we propose the first statistically grounded FL framework that jointly learns heterogeneous predictive models and the distributional structure required for query routing within a single principled objective. Second, we develop new algo- rithmic correction for the classification component of the EL. objective. Because each client is associated with only a single class label for client identification, the vanilla loss suffers from an extreme form of label shift; we propose simple yet effective reweighting adjustment that yields more stable classifier. Third, through experiments on benchmark datasets, we demonstrate that our approach consistently improves both predictive accuracy and routing precision compared to stan- dard FL methods, underscoring the benefits of integrating guidance directly into the FL workflow. Togetherr, these developments transform the FL server from a passive aggregator into an intelligent
2

---

Published as a conference paper at ICLR 2026
router capable of directing queries to the most suitable client, opening the door FL systems that are not only
but also adaptive and expertise-awail
FEDDRM: GUIDING CLIENTS IN HETEROGENEOUS FL
2.1 PROBABILISTIC DESCRIPTION OF STATISTICAL HETEROGENEIT!
Consider FL system with m clients. Let D; ((Xij, Yoy))== denote the training the i-th client, where each sample is drawn independently from Px,-- We consider the multi-class classification case where Yij E [K] = {1..... K} with marginal distribution P((a) k) Tssk for k E [K], and features conditioned on the labels are distributed as Xiyl(Y, k) pk We
denote the marginal distribution of the features on client as p and the conditional distribution of Y given x as (T = x|X = >))=r Different types statistical heterogeneity can be described in terms of the family of distributions {P (PxY m
Covariate shift: Clients differ their marginal feature distributions while sharing the same con- ditional label distribution. In our notation, this corresponds to
p(() for i + i', but p (Y = = = z) = py k| = x) for all x and k. X
Label shift: Clients have different label marginals but share the same conditional feature distri- butions given the label. Equivalently,
T As1 .... T(K) * nे == (TP1. Tyk for some i + v, but p(() -p(') for all k.
In practice, real-world FL systems often exhibit combinations of these shifts, which leads to the full distributional shift where both n, and ( ( K
1==1 may vary across clients.
2.2 SEMIPARAMETRI DENSITY RATIO MODEL
For clarity, we begin with the special case of covariate shift across clients. Extensions to other types of heterogeneity will then follow naturally. Let go (xx) represent a feature embedding (e.g., an embedding from DNN parameterized by 0) s.t. the conditional distribution of Y|X is given by:
explox go()))
P(Y = k| X =
Σ, expla,, go()))
We drop the superscript (i) since this conditional distribution remains the same across all clients under covariate shift. Applying Bayes' rule to (T) we derive that the class-conditional distributions connected by exponential function:
!P! /ap( (x) = exp(a) + + go(x))
(2)
where dP(() (d denotes the Radon-Nikodym derivative Og + / for i [m].
deriviaati of dP with respect to dp(1) and at
To facilitate knowledge transfer across clients in FL, we assume their datasets share some common
underlying statistical structure. Specifically, we relate the client distributions (P(t) pp)mm fi=1 through
hypothetical reference measure p(()) at the server, using DRM Anderson 1979):
dP( /dp (x) exp(7: + E (go(x))
where (+) is parametric function with parameters 7. since the server may not have data directly, although the formulation also applies when server-side
We refer p(0) hypothetical reference
DRM provides a framework for modeling the relationship between two or more populations that share similar characteristics. It is highly flexible and encompasses several commonly used parametric distribution families- such as the binomial, exponential, and normal familics-an special cases (Kay & Litle 1987.
3

---

Published as a conference paper at ICLR 2026
data are available. The DRM captures differences in the conditional distributions of X|Y across clients via density ratios, with log-ratios modeled linearly in the embeddings. This avoids estimating each distribution separately, focusing instead on relative differences. When the covariate shift is not too severe, the marginal distributions of different clients are connected through this parametric form, making FL effective by leveraging shared structure across clients. On the other hand, if the distributions differ too drastically, combining data from different clients is unlikely to improve performance; in this case, even if the DRM assumption does not hold, it is not limitation of the formulation but a consequence of the inherent nature of the problem. Thus, the assumption is reasonable in practice. When % and & = 0. (3) reduces the IID case. Under this assumption, we obtain the following relationship between the marginal feature distributions:
Theorem 2.1. With 2 and (3). the marginal distributions of also satisfy the DRM:
dP /dP (I)) exptn! + 5. (go(((=))
(4)
where of Y + log(#aa/ma) for for all E [m]. and p an unspecifted reference measure.
See proof in App. This theorem relates each client's marginal distribution to the reference distribution through a parametric tilt, which directly facilitates construction of the likelihood. the reference measure XX was fully specified, all {PX, x,y | would also be fully determined, and one could estimate the unknown model parameters using a standard maximum likelihood approach. In practice, however, p(o) is unknown, and assuming follows a parametric family risks model mis-specification and potentially biased inference.
To address this challenge, we adopt flexible, nonparametric approach based on EL Owen 2001 that integrate data across heterogencous populations via DRM (Qin & Zhang 1997 Fokianos et al, 2001 Chen Liu| 2013 Li al. 2017 Liu et al., 2017 2025). EL constructs likelihood functions directly from the observed data without requiring a parametric form. Instead of specifying a prob- ability model, it assigns probabilities to the observed samples and maximizes the nonparametric likelihood subject to constraints, such as moment conditions. Unlike classical parametric likeli- hood, EL adapts flexibly to the data, making particularly suitable when the underlying distribution is unknown or complex, but valid structural or moment conditions are available. Specifically, we let
Pij P ((X.))) 2 O, vi [ j E [ms],
treating the Pij as parameters. In this way, the reference measure P is represented as an atomic measure without any parametric assumptions, and most importantly all samples across clients are leveraged for information sharing. To ensure that p(() and (P(() 1m are valid probability measures, the following constraints are imposed:
m
as
m
1,
Σ>
Psy exp (ol + & (x00 11, VIE [m].
(5)
2.3 A SURPRISINGLY SIMPLE DUAL LOSS
With the semiparametrie DRM for heterogeneous FL., we propose a maximum likelihood approach for model learning. Let = {pu), a = {ax]. 3 = (3x). y = {n & = {E.) and Ç = (a, 6, r),
the log empirical likelihood of the model based on datasets across clients is
EN (p. = Σ log p(o) X.Y (1X,, Yo)3)= = Σ I(Yg = k)) log = k|X,,)) + Σ log pX, ((X1+})
8.j.k
-110 = log P(Y klx,,) Σ ho + hy (go (X+;)) + Pty).
i.j.k
where the last equality makes use of Theorem 2.T Since our goal is to learn [ on each client, the weight becomes nuisance parameter, which we profile out to learn the parameters that are connected to the conditional distribution of YIX = x. The profile log-EL of Ç is defined as (C) sup, LN (p, () where the supremum is under constraints (5). By the method of Lagrange multiplier, we show in App. C.2 that an analytical form of the profile log-EL is
pCN(() = Σ 1Y k)) log P(,, k| k())) + Σ (od < (oo(()))) log (gg))((()
(6)
4
4

---

Published as a conference paper at ICLR 2026
where PL)(C) -
(explil + 5/ h (00(+16))--1 }
and the Lagrange multi-
pliers {pt) are the solution to
explni + & h (go })} - -
Σ 1. Ev Pr exp( e, (go
0.
Although the profile log-EL in 6 has a closed analytical form, computing it typically requires solving a system of m equations for the Lagrange multipliers, which can be computationally de- manding. Interestingly, at the optimal solution these multipliers admit closed-form expression, yielding a surprisingly simple dual formulation of the profile log-EL presented below.
Theorem 2.2 (Dual form). For any fixed 9, the Lagrange multipliers Pt = =n/N and the profile log-EL EE-ELKO in (6) becomes
exp(1 + E, hr (((o(((sg))) E exp(nt + El h (g0(zz))
exp(ann + BTT go(s)))
p = Σ log
Σ log
Σn explax k + B gol (Zig))
1
1,
at optimality, up to some constant where of = log(ns/n) + 7
See App. C.3 for proof. The theorem allows us to define the overall loss function as the negative
profile log-EL:
((() -p(x(() = Σ (cali, (.
Σ
ferlvig (oeliig-90 (z1;): B),
where (CE(y, x; a, -aa, + By x) ) + log(E> exp(ax + x[x)} is the cross-entropy loss.
D 1
22
g
00
Remark 2.3 (Beyond covariate shift). Our method is de-
scribed under covariate shift. The derivations key steps ( and () do not require the marginal distribution of Y to be iden- tical across clients, which allows us also accommodate la-
Clant chweo
TSrgel
Client dass. Hink
TArgel idaas head
block
Tongetibbt
Tangetiabol GooT 12u
bel shift. Importantly. we show that our approach extends to the more general setting where both Y X and X differ across clienty in App. D In this case, after a detailed derivation, we find that the overall loss simplifies to a minor adjustment in the target-class classification head. Concretely, the target- class classification loss is equipped with a client-specific lin- head, resulting in the final architecture shown in Fig. 2 Interestingly, this architecture closely resembles those in personalized FL methods such as (Collins) al. 202T:: y target-class classification is performed with client-specific heads, while our new client classification component relies on a single shared head across all clients.
CE: Client
GE: Target dasssiicutunn
CE Client Sisttllaataam
GE: Target dassel
clsoaa Restaor
Client
Figure 2: Network architecture. Gray blocks are shared among all clients, while colored blocks specific to each client.
Remark 2.4 (Guiding new queries). Although the derivation is mathematically involved, the re- sulting loss function is remarkably simple: it consists of two cross-entropy terms, each associated with a distinct classification task. The first term identifies the client from which sample originates, while the second predicts its target class. The additional client-classificttion head produces, for each query, a predictive distribution over clients according to (4). By routing a query to the client with the highest predicted probability, we obtain a principled mechanism for assigning new data to the client best equipped to handle it.
2.4 OPTIMIZATION ALGORITIMM
The overall loss E(C) 7 is defined as if all datasets were pooled together. Since optimizing {(C) with vanilla SGD and weight decay is equivalent to minimizing loss function with an explicit L penalty, we denote the loss as (Ç) {({) (p/2)][<1l, with minimizer ČN. The subscript N is used to indicate that this weight is based on N samples. In the FL setting, the global loss decomposes naturally into client-specific contributions: (*(S) Σm (M./NN)&,(S) where
4(() (,(,g)) + (a, (p/2)]1<113.
4(7,8)= n. ",'E, loE(i, (90(())))) and (( B) = n. Σ, (CE a,,
5

---

Published as a conference paper at ICLR 2026
A key difference arises between these two terms. For the client-clasificattin loss le(), 5), the i-th client only observes samples labeled with its own client index i. In contrast, the target-class loss (a, 9) typically spans multiple target labels per client (though with varying proportions). This asymmetry leads to more pronounced gradient drifß in E)) To il- lustrate, consider the gradient of the client-llasiffication loss with respect to Ykk
5.25
5.00
4.00
3.50
00./0rk = ni Σ may pof( k) -p- (10(()9)1))))
200
Communication Round
400 600 800
Figure 3: Relative gradient drift.
exp(1x =T)x). Since 1(: R-) O for all k * the gradient contributed by client provides no meaningful information about other clients' parameters. As result, local updates the client-olassffication head are inherently biased, which in turn am- plifies gradient drift relative to target-class head. Fig. 3 shows this effect on a 10-class classification task with 3 clients and a randomly generated embedding using FedAvg: the gradient drift for client classification is markedly more severe than that for target classification.
where (:) exp(1>>
Reweighting strategy. To address this, we draw on reweighting principles (Chen et al. 2018 Liu et aL. 202T) to propose a simple yet effective method with theoretical guarantees. Our approach down-weight client classification loss, whose gradient exhibits larger drift, resulting in the per- client loss:
(C(C) = x))())) + Ali(a, 3),
for A > 0.5 and the reweighted global loss is ε(C) = Em (m:/N)))))) see Algorithm
Algorithm 1: FedDRM
Input: Clients m. communication rounds T, local steps E, learning rate 1j. trade-off X
Initialize backbone ((0), target head {(a,)) (0)
(0( )m and client head
for to 0. - do
Server broadcasts (0(*) ,-()), ((()) to all clients
for client i E [m]]] in parallel do
(0(h-0) , (1.0) (((.)) (1,0) 7. c
(1.0)
(1) (()) -(0, , -.(10)) - (0
a
((1)
for k:
k= 1,.... 1 do
Get target loss Es(a (+kt ,((
ofc.8) and client loss (i(T(ti))) on ef cle.)) of llt
(0.)))
7
&(<(tt))) +- Ala(a, (t.k) ofth0) oltck) + (1- x)(((T(" (t.)) (ea) cltck) o((c))
c/ + c((i8) -ni(((
10
end
(0(-2) (411), ,((t2)) Tt (((11), ((442) 5y g((tt)) (o((:)) ((( ,(L A(1,8) vo (11,), c((,))
Client sends (0(1+1)) .i (1t1) (e11) m . (((+2) back to the server
12
13
end
Server updates (g(c+),, f(t-1), , yethl, (o0t),g(tt)) g(tti)y + Em " (0((t)) s. .h e e((+1))
(1+1) (++1)
14
15 end
To accelerate convergence, larger value of À is desirable. However, as X + 1, the target- class classification begins to dominate, which hinders effective training of the client classifica- tion and ultimately weakens the model's ability to guide clients. To illustrate the trade-off be- tween accuracy and convergence, we consider simplified setting where the embedding is fixed (i.e., 0 and r are known) and the true data-generating mechanism follows a multinomial logis- girue, a.e ,,r,).. We define the heterogeneity mea- Ve(C)) which admits the decomposition (²(() X (a. (c, 3). Let G², Go client and class denote the corresponding max-
tic model with parameters girue = -t _ sure
(1 - X) ET, C imum values across updating rounds = 0, 1,1,..... T 7-1. 1. Then, G < (1 - G) + A G2G class With this notation place, we state the following result:
The gradient drift of the client loss is 20 := Em (n./N)) (T, 2)12, that of the target-class loss is Go m (ns/N)][V(((aa))) (M

---

Published as a conference paper at ICLR 2026
Theorem 2.5. Assume 1o is p-stronngy convex and L-smooth. Suppose nj S L and furthermore
nLE & 1/4 Let ((t) be the output after t communication rounds. Then T, T.N - 00 have
[(1 ((1-))] Il + p}-- {AlIZal min + +p} n1 E² { X) Pinou + X² x2G Gduss
IK ST - O₂
N
1 THp)
where || - >mmn(()) and Letiens and Teloss denote the Fisher information matrices with respect to () and (a. 3). respectively.
The proof and the detailed definition of Fisher information matrix is deferred to App. The first term in the bound capture the statistical accuracy. while the last term reflects the convergence rate. For faster convergence, larger À is preferred, while for higher accuracy, À must be chosen to balance {(1 X) ml + p) - and [>] + p}- Together, these terms reveal the trade-off role of X. In practice, since the Fisher information matrices and gradient drifts are unknown, À can be tuned using validation set. We empirically demonstrate the trade-off in Fig. 4
3 EXPERIMENTS ON BENCHMARK DATASETS
3.1 EXPERIMENT SETTINGS
Datasets. We conduct experiments on CIFAR-10 and CIFAR- (Krizhessky 2009). each contain- ing 60, 000 x2 RGB images. CIFAR-10 has 10 classes with 6, 000 000 per class. CIFAR-100 has 100 classes, with 600 images per class, grouped into 20 superclasses. Based on these datasets, we construct three tasks of increasing complexity: (a) 10-class classification on CIFAR-10, (b) 20-class classification using the CIFAR-100 superclasses, and (c) 100-class classification using the fine-grained CIFAR-100 labels.
Non-IID settings. Since standard benchmark datasets do not inherently exhibit statistical hetero- geneity, we simulate non-IID scenarios following common practice (Wu et al. 2023 Tan et al. 2023
Lu et al. 2024). We introduce both label and covariate shifts. For label shift, we construct client datasets using two partitioning strategies: (1) Dirichlet partition with a 0.3 (Dir-0.3): Follow- ing (Yurochkin et al., (201)). we draw class proportions for each client from a Dirichlet distribution with concentration parameter a = 0.3, leading to heterogeneous label marginals and unequal dataset sizes across clients. (2) S shards per client (S-SPC): Following (McMahan et al. 2017), we sort the data by class, split it into equal-sized, random to each client. This yields equal dataset sizes while restricting each client's 's label support to at most S classes. Fach dataset first partitioned across clients using one of the partitioning strate- gies, and within each client, the local dataset is further split 70/30 into training and test sets. For covariate shift, all three nonlinear transformations are applied to each client's dataset: (1) gamma correction: brightness adjustment with client-specific gamma factor 7. (2) hue adjustment: color rotation with client-specific hue factor Ah. (3) saturation scaling: color vividness adjustment with client-specific saturation factor We set y {0.6, 1.4), Ah (-0.1, 0.1). and x. E (0.5, 1.5) in the main experiment, resulting in an 8-client setting. See examples in App. F.I
homogenosss shards, and assign S shards uniformly at
Baselines. We compare our FedDRM against variety of state-of-the-art personalized FL tech- niques, which learn a local model on each client. Ditto (Li et al. 2021c) encourages local models to stay close via global regularization. FedRep (Collins al. [22 learns a global backbone with local linear heads. FedBABU (Oh et al. (022) freezes local classifiers while training a global back- bone, then fine-tunes classifiers per client. FedPAC (Xu et al. 2023) personalizes through feature alignment to a global backbone. FedALA (Zhang et al. 2023) learns client-wise mixing weights that adaptively interpolate between the local and global models. FedAS (Yang et al.) (024) aligns local weights to the global model, followed by client-specific updates. ConFREE (Zheng et al.) (025) resolves conflicts among client updates before server aggregation. We also compare with other stan- dard FL algorithms- FedAvg (McMahan et al. 20177), FedProx (Li et al. (2200b), and FedSAM (Qu) ctt al. 2022h-which aim to achieve a single global model under data heterogeneity. To ensure fair comparison, we fine-tune their global models locally each client, yielding personalized variants denoted FedAvgFT, FedProxFT, and FedSAMFT.
Network architecture. We ResNet- 18 (He et al. 2016) as the feature extractor (backbone), which encodes each input image into a 512-dimensional embedding. For the baselines, this embed- ding is projected to 256 dimensions via linear layer and fed into the image classifier. FedDRM
7

---

Published as a conference paper at ICLR 2026
extends this design by adding separate client- classiftissiiiiiionnn head: the 512 -iz-dimnnsional embedding is projected to 256 dimensions and fed into the client classifier. Importantly, FedDRM uses the same image classification architecture as all baselines.
Training details. To ensure fair comparison, all methods are trained for 800 communication rounds with 10 local steps per round and batch size of 128. For fine-tuning-bn based methods, allocate 700 rounds for global training and 100 rounds for local fine-tuning. We use SGD with momentum 0.9, an initial learning rate of 0.01 with cosine annealing, and weight decay 5 10-1, Method-specific hyperparameters are tuned to achieve their best performance.
3.2 EVALUATION PROTOCOL
To assess the effectiveness of proposed method in guiding clients under heterogeneous FL., introduce a new performance metric, termed system accuracy. This metric is designed to evaluate the server's ability to guide clients effectively. Concretely, we construct pooled test set from all clients. For F FMdDRM, we first use the client classification head to identify the most likely client for each test sample by maximizing the client classification probability. The local model of the selected client is then used to predict the image class label. For baseline methods, which lack this client- guidance mechanism, we instead apply a majority-voting strategy: each client's personalized model makes prediction for every sample in the pooled test set, and the majority label is taken as the final prediction. The overall classification accuracy on the pooled test set is reported as the system accuracy. We also report the widely used average accuracy in personalized FL, which measures each local model's classification accuracy on its own test set. The final value is computed as the weighted average across all clients, with weights proportional to the size of each client's training set. In all experiments, we report the mean and standard deviation of both average accuracy and system accuracy over the final 50 communication rounds.
3.3
MAIN RESULTS
We present the system accuracy and average accuracy in Tab. and Tab.
respectively. Across
Table 1: System accuracy on CIFAR-10/20/100 under Dir-0.3 and 5/25-SPC settingss
CI
CIFAR-20
CIFAR-100 100
Method
Dir-0.3
5-SPC
Dir-0.3
25-SPC
Dir-0.3
25-SPC
Ditto
47.64 0.25
46,99 1 0.23
29.56 1 0.18
31.87 1 10.1
15.97 1
19.51 1 0.16
FedRep
24.96 = 0.19
33.19 1119-=00 0.22
23.83 + 0.15
24.82 + 0.20
11.11
12.02 0.12
FedBABU
57.43 0.17
57.17 0.24
36.96 0.11
40.78
0.13 0.13
22.92
0.17
97.27+00.1
FedPAC
25.14 0.21
33.24
23.83
24.83
1..17
11.99 0.12
PedALA
61.33 0.17
53.20 0.20
32.7811.01
35.79
20.10
25.8011..
FedAS
28.76 U.19
39.71=0.20
27.16 0.16
0.15
13.87
13.51 0.14
COFFREE
25.66 0.22
34.06 0.22
24.08 1: 0.17
2..18
11.32
12.12 0.13
FodAvgFT
54.90 0.22
56.19 0.17
37.53
0.16
25.21
27.96 0.17
FedPtoxFT
55.01 0.20
56.27 0.21
37161
0..5
20.13
27.82 0.18
FedSAMFT
55.83 0.21
51.73 = 0.19
$1.2300.14
36.1
22.97
26.89 0.15
FedDRM
62.78 0.20
58.50 0.23
37.63 1 0.21
4.14
26.03 0.11
31.24 0.17
all settings, FedDRM consistently outperforms the baselines on both metries, demonstrating its abil- ity to leverage statistical heterogeneity for system-level intelligence while also providing effective client-level personalization. In contrast, the baselines primarily focus on addressing data hetero- geneity, resulting in lower system accuracy due to disagreements among their personalized models. Additionally, when using a majority-vote approach as an intelligence router, baseline methods must evaluate all local models, whereas FedDRM requires evaluating only a single model. The shared backbone in FedDRM can also be efficiently repurposed for image prediction by feeding it into the corresponding client-specific classification head. We also compare the influence of label shift in this experiment beyond covariate shift, the results align with our expectation that the less severe label shift Dir-0.3 case has higher accuracy than S-SPC for all methods.
8

---

Published as a conference paper at ICLR 2026
Table 2: Average accuracy on CIFAR-10/20100 under Dir-0.3 and 5/25-SPC settings.
CIFAR-10
CIFAR-20
CIFAR-100
Method
D.r-0.3
5-SPC
Dir-0.3
25-SPC
Dir-0.3
25-SPC
Ditto
76.34 = 0.11
65.17=0.11
40.36 + 0.18
44.83 + 0.19
29.25
0.16
36.58
0.18
FedRep
76.49 0.15
64.96
41.54 0.16
46.57 + 0.19
31.22
0.15
39.11
FedBABU
78.22 0.14
00.22
11.18
32.91 0.14
0.14
FedPAC
76.53 0.13
10.19
41.60 0.16
46.53 1 0.19
31.20 1.
39.13 0.22
FedALA
64.35 = 2.40
33.30
= 0.47
21.83 0.94
27.83 + 1.39
FedAS
78.69 0.17
69.82 + 0.16
45.65
0.18
51.73
36.06 + 0.13
44.26 +: 0.19
CoFFREE
76.73 = 0.16
65.59 0.17
41.91 = 0.16
47.04 77.04:0.2 0.21
31.57 0.15
39.63 0.17
FedAvgFT
79.08 0.11
=
46.55 + 0.15
0.11
36.83 0.11
43.94 + 0.20
FedProxFT
79.07 = 0.12
72.07 0.18
46.58 = 0.19
52.49 0.18
36.87 0.18
43.96 : 0.19
FedSAMFT
75.53 0.11
66.30 1 0.16
41.11 1.166
44.89 0.16
32.53 0.14
40.25 0.16
FedDRM
80.26 + 0.16
72.50 0.16
47.91 + 0.18
53.72 + 0.20
37.91 0.15
46.73 0.15
3.4 SENSITIVITY ANALYSIS
We evaluate the sensitivity of our method to several key factors. Experimental details are reported
in App. F2
Impact of weight À on system accuracy. The reweighting
parameter is crucial for deploying the EL-based framework in the FL setting. As shown in Fig. 4 we observe the expected trade-off between two objectives: increasing À places more emphasis image classification and less client classifica- tion. This shift improves overall accuracy but reduces client accuracy, consistent with Thm. 2.5 The best balance between the two is achieved at À = 0.8, where system accuracy peaks, marking the optimal trade-off for the task of guiding queries in the FL system.
60
SystmAzz Chertho
Acuuray (%)
S
8
Trade-off Parameter
Figure Client & image accu- racy trade-off on CIFAR-10 under the Dir-0.3 setting.
Covariate shift intensity. We have already demonstrated in the main results that label shift is detrimental to all methods, with more severe shifts causing greater harm. To further examine the impact of covariate shift, we fix the degree label shift and vary covariate shift at three intensity levels-lww, mid, and high-by adjusting the parameters of the nonlinear color transformations. As shown in Fig- 5 the results reveal a clear trade-off: higher co- variate shift intensifies differences between client data distributions, which facilitates client routing but simultaneously weakens information sharing across clients, thereby making image classification more difficult. Additional results examining the sensitivity of our method to the severity of label shift are provided in App. F2.4
Backbone sharing strategy. In formulation, the target- class classification task uses the embedding ge(x) for an in- put feature x, while the client-classifcation task uses the embedding h (g) for the same feature. Since both ge and hy are parameterized functions, the optimal shar- ing strategy between the two is not obvious. To explore this, evaluate four cases: sharing, shallow sharing, mid sharing, and deep sharing. As shown in Fig. 6 all strategies perform similarly, with shallow sharing slightly ahead. However, given the substantial increase in parameters for shallow sharing, deep sharing offers more parameter- efficient alternative while maintaining strong performance.
80
5025
70
Accuraay
8
80
40
Intensity
AvorageAAc
ClontAcs
System Aco
Figure 5: Influence of covariate shift intensity on CIFAR-10 under the Dir-0.3 setting.
Number of clients. To check scalability, we set the number of clients from 8 to 32 and compare FedDRM against the top-2 baselines from the main experiments. As shown in Tab. 3 while all methods exhibit a moderate performance decline as the client pool expands (a common challenge in FL), FedDRM consistently maintains significant performance advantage across both system and average accuracy. This demonstrates that our method scales effectively, preserving its superiority even the system grows.

---

Published as a conference paper at ICLR 2026
Table 3: Sensitivity analysis on the number of clients m.
Mothod
System Accuracy
Average Accuracy
16
82
= 16
Jh 24
82
FedAS
32.58
0.19
36.55 0.21
38.12
34.50
78.86 0.12
73.28 0,15
70.90 0.16
13.45
FedAvgiF
0.15
1..1
FediuR
69.59 0.20
51.61
0.22
50.18 + 0.20
46.62
0.17
80.47 + 0.12
74.25 0.15
76.04 + 0.14
74.45 0.18
4 EXPERIMENT ON REAL MEDICAL DATASET
To further demonstrate FedDRM's effectiveness in healthcare, we evaluate it on the real medical dataset RETINA, follow- ing (Huang et al. 2025). RETINA comprises fundus im- ages from three clinical centers ACRIMA (Diaz-Pinto al. 2019), Rim (Fumero Batista et al. 2020), and Refuge (Orlando al. 2020). We exclude Drishti, which has only 82 images, while the others provide at least 385. Each 96 x 96 RGB im- age is labeled as Glaucomatous Normal, creating binary classification task.
37.21
70-05
Accurany (%6
1164 71.4
7244
70
Share
Awerage
System
Figure Impact of the sharing strategy on CIFAR-10 under the Dir-0.3 setting using LeNet (Lecun) et al. 1998).
This dataset naturally fits a 3-client system, with each client representing one center. The different image sources cause a covariate shift in RETINA. Furthermore, the class ratios (pos- itive vs. negative) across the three datasets are 1.34, 1.94, and 0.46, introducing a realistic label shift. Our experimental setup largely follows the CIFAR experiments, with several adjustments: the network embedding dimension is set to 4608 and the projection dimension 512. All methods train for 100 communication rounds with batch size of 32. For fin-tuuning-ras -bann methods, we allocate 90 rounds to global training and 10 to local fine-tuning.
95
Average Ace
85.0
Acc
Train Loss
92.5
System
80
87.5-
0
25
50
75
100
Communication Round
Dito
FedBABU
FedALA
ConFREE
FedProxF
FedDRM
FedRep
FedPAC
FedAS
FedAugFT
FedSAMFT
Figure 7: Average accuracy, system accuracy, and train loss on RETINA.
Fig. shows that FedDRM consistently outperforms all baselines on RETINA. Measured in ab- solute accuracy points, FedDRM exceeds the competing methods by 0.83-3.77 points in average accuracy and by 1.41-7.67 points in system ceuuuacc--uusstanne margins given the small size and pronounced heterogeneity of this dataset. These results underscore the robustness and practi- cal relevance of FedDRM the presence simultaneous covariate and label shifts. Furthermore, F FedDRM achieves the lowest training loss and the most stable convergence trajectory, demonstrating its effectiveness in capturing heterogeneous structure in real multi-center medical data.
CONCLUSION
This paper presents FedDRM, novel FL. paradigm that transforms statistical heterogeneity from challenge into resource. By introducing a unified EL based framework, FedDRM simultaneously learns accurate local models and a client-selection policy, enabling a central server to intelligently route new queries to the most appropriate client. Empirical results demonstrate that our method out- performs existing approaches in both client-level personalization and system-level utility, paving the way for more adaptive and eeffurce-effcceent FL systems that actively leverage statistical diversity. We believe that this work marks a meaningful step toward more adaptive, resource- essuurceefficieet, and intelligent FL systems.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENT
Zijian Wang & Qiong Zhang are supported by the National Key R&D Program of China Grant 2024YFA 1015000 and the National Natural Science Foundation of China Grant 12301391. Xiaofei Zhang is supported by National Natural Science Foundation of China Grant 12501394 and the Nat- ural Science Foundation of Hubei Province of China Grant 2025AFC035. Yukun Liu is supported by National Natural Science Foundation of China Grant 12571283. The contributions of Xin Zhang were made in a personal capacity and outside the scope of his employment at Meta. This research did not use any Meta data, tools, or internal resources. The views expressed in this paper are those of the authors and do not necessarily reflect the views of Meta.
REFERENCES
Durmus Alp Emre Acar, Yue Zhao, Ramon Matas, Matthew Mattina, Paul Whatmough, and Venkatesh Saligrama. Federated learning based on dynamic regularization. In International Con- ference on Learning Representations, 2021.
J. A. Anderson. Multivariate logistic compounds. Biometrika, 66(1):17-28) 1979.
Christopher Briggs, Zhong Fan, and Peter Andras. Federated learning with hierarchical clustering of local updates to improve training on non-iid data. In International Joint Conference on Neural Networks, 2020.
Jiahua Chen and Yukun Liu. Quantile and quantile-function estimations under density ratio model. The Annals of Statistics, 116))11669-169,, 2013.
Zhao Chen, Vijay Badrinarayanan, Chen- Yu Lee, and Andrew Rabinovich. Gradnorm: Gradient normalization for adaptive loss balancing in deep multitask networks. In International Conference on Machine Learning, 2018
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared repre- sentations for personalized federated learning. In International Conference Machine Learning, 2021.
Andres Diaz-Pinto, Sandra Morales, Valery Naranjo, Thomas Köhler, Jose M Mossi, and Amparo Navea. Cnns for automatic glaucoma assessment using fundus images: extensive validation. Biomedical Engineering Online, 18(1):1-19, 2019.
Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning with theo- retical guarantees: A model-agnostic meta-learning approach. In Advances Neural Information Processing Systems, 2020.
Konstantinos Fokianos, Benjamin Kedem, Jing Qin, and David Short. A semiparametric approach to the one-way layout. Technometrics, 43(1):56-65 2001.
Francisco José Fumero Batista, Tinguaro Diaz-Aleman, Jose Sigut, Silvia Alayon, Rafacl Arnay, and Denisse Angel-Pereira. RIM-ONE DL: A unified retinal image database for assessing glaucoma using deep learning. Image Analysis & Stereology, 39(3):161-167, 2020.
Liang Gao, Huazhu Fu, Li Li, Yingwen Chen, Ming Xu, and Cheng -Zong Xu. FedDC: Feder- ated learning with non-lll data via local drift decoupling and correction. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022.
Avishck Ghosh, Jichan Chung, Dong Yin, and Kannan Ramchandran. An efficient framework for clustered federated learning. In Advances in Neural Information Processing Systems, 2020.
Yongxin Guo, Xiaoying Tang, and Tao Lin. FedBR: Improving federated learning on heterogeneous data via local learning bias reduction. In International Conference on Machine Learning. 2023.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2016.

---

Published as a conference paper at ICLR 2026
Chun-Yine Huang, Ruinan Jin, Can Zhao, Daguang Xu, and Xiaoxiao Li. Federated learning vir- tual heterogeneous data with local-glober dataset distillation. Transactions on Machine Learning Research, 2025.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International Conference on Machine Learning, 2020.
Richard Kay and Sarah Little. Transformations the explanatory variables the logistic regression model for binary data. Biometrika, 74(3):495-501 1987.
Yeongwoo Kim, Ezeddin Al Hakim, Johan Haraldson, Henrik Eriksson, José Mairton B da Silva, and Carlo Fischione. Dynamic clustering in federated learning. In International Conference on Communicationss 2021.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009.
Y. Lecun, L. Bottou, Y. Bengio, and P. Haffner. Gradient-based learning applied to document recog- nition. Proceedings of the IEEE, 86(11) 36(11))228822244 1998.
Chengxi Li, Gang Li, and Pramod K Varshney. Federated learning with soft clustering. IEEE Internet of Things Journal, 9(10)7773-7782, 2021a.
Li Li, Yuxi Fan, Mike Tse, and Kuo-Yi Yi Lin. review applications in federated learning. Com- puters & Industrial Engineering, 149:106854, 2020a.
Pengfei Li, Yukun Liu, and Jing Qin. Semiparametric inference in genetic mixture model. Journal of the American Statistical Association, 112(519) 2017.
Qinbin Li, Bingsheng He, and Dawn Song. Model-contrastive federated learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2021b.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. In Proceedings of Machine Learning and Systems, 2020b.
Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and robust federated learning through personalization. In International Conference on Machine Learning, 2021c.
Xiaoxiao Li, Meirui Jiang, Xiaofei Zhang, Michael Kamp, and Qi Dou. FedBN: Federated learn- ing on non-lID features via local batch normalization. In International Conference on Learning Representations, 2021d.
Bo Liu, Xingchao Liu, Xiaojie Jin, Peter Stone, and Qiang Liu. Conflict-averse gradient descent for multi-task learning. In Advances in Neural Information Processing Systems, 2021.
Siyan Liu, Chi-Kuang Yeh, Xin Zhang, Qinglong Tian, and Pengfei Li. Positive and unlabeled data:
Model, estimation, inference, and classification. Journal of the American Statistical Association, pp. 1-12, 2025.
Yukun Liu, Pengfei Li, and Jing Qin. Maximum empirical likelihood estimation for abundance in closed population from capture-recapture data. Biometrika, 104(3):527-543 2017.
Guodong Long, Yue Tan, Jing Jiang, and Chengqi Zhang. Federated learning for open banking. In Federated learning: privacy and incentive, pp. 240-254. 2020.
Guodong Long, Ming Xie, Tao Shen, Tianyi Zhou, Xianzhi Wang, and Jing Jiang. Multi-center federated learning: clients clustering for better personalization. World Wide Web, 26(1):481- 2023.
Yang Lu, Lin Chen, Yonggang Zhang, Yiliang Zhang, Bo Han, Yiu-ming Cheung, and Hanzi Wang. Federated learning with extremely noisy clients via negative distillation. In Proceedings of the AAAI Conference Artificill Intelligence, 2024.
12

---

Published as a conference paper at ICLR 2026
Xiaosong Ma, Jie Zhang, Song Guo, and Wenchao Xu. Layer-wised model aggregation for person- alized federated learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication learning of deep networks from decentralized data. In International Conference on Artificial Intelligence and Statisties, 2017.
Jachoon Oh, Sangmook Kim, and Se-Young Yun. Fedbabu: Toward enhanced representation for federated image classification. In International Conference Learning Representations, 2022.
José Ignacio Orlando, Huazhu Fu, João Barbosa Breda, Karel van Keer, Deepti R Bathula, Andrés Diaz-Pinto, Ruogu Fang, Pheng-Ann Heng, Jeyoung Kim, JoonHo Lee, et al. REFUGE challenge: A unified framework for evaluating automated methods for glaucoma assessment from fundus photographs. Medical Image Analysis, 19:10,570, 2020.
Art B Owen. Empirical likelihood. Chapman and Hal//CRC, 2001.
Jing Qin and Biao Zhang. A goodness-of of-tit test for logistic regression models based case-control data. Biometrika, 84(3):609-618, 1997.
Zhe Qu, Xingyu Li, Rui Duan, Yao Liu, Bo Tang, and Zhuo Lu. Generalized federated learning via sharpness aware minimization. In International Conference on Machine Learning. 2022.
Virginia Smith, Chao-Kai Chiang, Maziar Sanjabi, and Ameet Talwalkar. Federated multi-task learning. In Advances in Neural Information Processing Systems, 2017.
Canh T Dinh, Nguyen Tran, and Josh Nguyen. Personalized federated learning with moreau velopes. In Advances in Neural Information Processing Systems, 2020.
Alysa Ziying Tan, Han Yu, Lizhen Cui, and Qiang Yang. Towards personalized federated learning. IEEE Transactions on Neural Networks and Learning Systems, 34( (4(11))9587-9603 2022.
Yue Tan, Chen Chen, Weiming Zhuang, Xin Dong, Lingjuan Lyu, and Guodong Long. Is het- erogeneity notorious? taming heterogeneity to handle test-time shift in federated learning. In Advances in Neural Information Processing Systems, 2023.
Aad W Van der Vaart. Asymptotic statistics. Cambridge university press, 2000.
Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and H Vincent Poor. Tackling the objective in- consistency problem in heterogeneous federated optimization. In Advances Neural Information Processing Systems, 2020.
Yue Wu, Shuaicheng Zhang, Wenchao Yu, Yanchi Liu, Quanquan Gu, Dawei Zhou, Haifeng Chen, and Wei Cheng. Personalized federated learning under mixture of distributions. In International Conference Machine Learning, 2023.
Jian Xu, Xinyi Tong, and Shao-Lun Huang. Personalized federated learning with feature alignment and classifier collaboration. In International Conference on Learning Representations, 2023.
Jie Xu, Benjamin S Glicksberg, Chang Su, Peter Walker, Jiang Bian, and Fei Wang- Federated learning for healthcare informatics. Journal of Healthcare Informatics Research, 5(1):1-19, 2021.
Xiyuan Yang, Wenke Huang, and Mang Ye. Fedas: Bridging inconsistency in personalized fed- erated learning. In Proceedings of the IEEF/VVF Conference Computer Vision and Pattern Recognition, 2024.
Qiaoyun Yin, Zhiyong Feng, Xiaohong Li, Shizhan Chen, Hongyue Wu, and Gaoyong Han. Tackling lata-heterogeneitt variations in federated learning via adaptive aggregate weights. Knowledge-Bel Based Systems, 304: 12484, 2024.
Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, Nghia Hoang, and Yasaman Khazaeni. Bayesian nonparametric federated learning of neural networks. In Interna- fional Conference on Machine Learning, 2019.
13

---

Published as a conference paper at ICLR 2026
Jianqing Zhang, Yang Hua, Hao Wang, Tao Song, Zhengui Xue, Ruhui Ma, and Haibing Guan. Fedala: Adaptive local aggregation for personalized federated learning. In Proceedings of the AAAI Conference Artificial Intelligence, 2023.
Hao Zheng, Zhigang Hu, Liu Yang, Meiguang Zheng, Aikun Xu, and Boyu Wang. Confree: Conflict-free client update aggregation for personalized federated learning. In Proceedings of the AAAI Conference on Artificial Intelligence, 2025.
14

---

Published as a conference paper at ICLR 2026
A THE USE OF LARGE LANGUAGE MODELS (LLMS)
Large language models (LLMs) were used solely as assistive tools for language editing and pol- ishing of the manuscript. The authors take full responsibility for the accuracy and integrity of the manuscript.
B DENSITY RATIO MODEL EXAMPLES
Many parametric distribution families including normal and Gamma special cases of the DRM.
Example B.1 (Normal distribution). For normal distribution @(x; pl, σ²) with mean μ and variance
a². We have (og( ,1, (2; o Op + 01æ: + Ozæ2 where Op log 02 /1 - (1///ff /2/o3))2, 01 M1/of 1 /2/ 82 (0² 2/2 and the basis function is g(x) (1, 2. 2² Example B.2 (Gamma distribution). For gamma distribution with shape parameter a > and rate parameter 3 0. We have log (f; a11, B1)) f(x; a2, (2)} 0o + 0fz + 02 log where 0o = log [()) - log ['()) + an log 31 0g log $2. 01 = B2 B1. 02 = ai - 02 and the basis function is g(x) = (1, (1,2,logı:) lo
C MATHEMATICAL DETAILS BEHIND FEDDRM
C.I DERIVATION OF
Proof. By the total law of probability, the marginal density of x d (x)
Σ RudPll (xll = k)
Σ Ti exp( (a/4 + 3x go (z) )dp()) {
(xly = 1)
Σ Tik exp(atx + 3% go (x) exp(n + &/ hn (x)) (x))) = =)
Σ Ti exp(ax + 3T (x)) explm + & h (gs
(xly 1)
Σ Tn Toi
Tox expla 3o go (x) expln + & ( (xly = 1)
Tn
exp(n + & ho (x)
To1
Let Let = 71 log(Fu/To)) and divide dP (0) x (x:) on both sides completes the proof.
C.2 DERIVATION OF THE PROFILE LOG-LLKELIHOOD
Let (pij.j j E [n.]) Given Ç = (ax, 3,,,,,,,,, v), the empirical log-likelihood function as function of p becomes
m7 n4
(N(P)
ΣΣ
log + constant
i=
where the constant depends only on ( and does not depend on p. We now maximize the empirical log- lieeiiillhood function with respect to p under the constraint (5) method.
using the Lagrange multiplier
Let
m
C
Σ logpij - N2 Σ
Psj
N
Pt
Σ 1
ply (explil + & hn
4
1.j
1=1
Setting
3c
m N Σ
0
Np
p lcxplll + (o0[ - 1].
0p:j
Pij
1=1
15

---

Published as a conference paper at ICLR 2026
Then multiply both sides by Pij and sum over i and j, we have that
ac
0
Pu
Opis
Σ
Npis
Σ
p + & (n 11
V
bd
1=1
N Nu
this gives μ 1. Hence, we get
Púj (1 1 plcppfe + El hy (20( }}}
where prs are solutions
Σ explai + E hy (x0(())) 1.j
0
Pr
pr (exptt + 5, h (go (xx;))))) -
by plugin the expression for Pij into the second constraints
C.3 DERIVATION OF THE VALUE OF LAGRANGE MULTIPLIER AT OPTIMAL
Recall that the profile log-EI. has the following form
(Sx(5)= = Σ (vi 1-) log P(uc = k((r;)) +
Σ (7l + & (gn(((1)))) + log pig(())))
Σ
(1yy = k:) Pog = k((,,))) + Σ
7 + E, (1o(()))))))
1.j
(exptal + s. hn(o0(1)5))) - ]1} }
Σ log 1
Σ 1=1
pl
where prs are solutions
explil + &, h (go(211))) -1 = 0.
Σ J I+ pr /explat + EJ (x((0((
Taking the partial derivative with respect to ?/ we have
0 ÖplN
Pi exp exptr (n) + & h(x,))
Σ 1 En vr pv (exp (// + E, (xx(5)) 1]
0nl
i.j
Er (0pr / (epp() + E hn (xx;)))
1 + Er (exp(1) E, (g0 (x:;)) })} -
bd
Np Σ Piy exp(1) + el ((:;)) + N Σ Opy
Σ
Piy jexp(1) + + ({]h(()))))) - 1]
Tht
on! 1
n Np.
The last inequality is based on the constraint in 5 Hence, we have Pr = n//N which completes the proof.
16

---

Published as a conference paper at ICLR 2026
C.4 DUAL FORM OF THE PROFILE LOG EL
At the optimal value, we have m//N with N = log-EL, we then get
m
n4. Plugin this value into the profile
exp(ares + B go (zi;))
expl1! + E, hn
pCx(()= = Σ log
Σ log
+
Σ, expla; + {] go (Iis))
Ei=l N exp(1) + E7 (n
exp(amm By gol (2i3z)
(n1/m4) exp + 5, My (go(a)))) 1= ( exp(ni + E (n(o((m)))))
Σ log
+
Σ log
ΣS exp(ax 83 go (xin))
1,
exp(a + 3 gol (Ii)))
(expla) El (g0 1=1 expl2t + E (1)((()))))
Σ log
+
log
Ej exp(a; + 3 go (zii))
74
log
N
The last term is an additive constant; the maximization does not depend on value, which com- pletes the proof.
D GENERALIZATION TO OTHER TYPES OF DATA HETEROGENEITY
In this section, we detail how our method generalizes to the setting where both Y|X and X differ across clients. Recall that the log empirical likelihood function is
(N = ΣΣ log X.Y ((Xij. Yi²1)
2=1
Σ
1(y k) log p()) = k(x,,)) + Elog log P ((X,))
k.j.k
We assume that each client has its own linear head for the conditional distribution:
explas + (z))
plo (Y k|X x) explaas + (1)))(x)))
The marginal distributions p( are linked in Theorem 2.1
dP
(x) explat + & hy (go())))
dp x
where p(() is an unspecified reference measure. Using non non- parametric reference distribution, we set
1(0 ((Xx)) Vi [m], j E [n.],
Poj
subject to the constraints
m
27 n
1
ΣΣ i=1j=1
Psy exp iií + + (X)(()))) vle [m].
Then, the log empirical likelihood across clients is
(n(p.C) = Σ 1(y = log ( - k|Xis) + Etst
hT + log p..)
The profile log-EL of ( defined as
plx(C) = E C)
17

---

Published as a conference paper at ICLR 2026
where the supremum is taken under the constraints above. Applying the method of Lagrange multi- pliers, we obtain the analytical form
(CN(() = Σ 1(Y k) log p0() (Ye, = k|Xis) + Ethl + + hr (g0(()6)) + log (t)
where
Piy(C) = N N -
Σ p (eppll + + ( -1]3 }
m
and the Lagrange multipliers {pr)m= solves
explai + & hor (go (xx;)))
Σ
0.
JJ Ev PV exp(1) E, b (go
Using the dual argument from Appendix C.3
the profile log-El. can be rewritten as
exp(ai.uu + PL., 00 (miz))
expl7] + E, (y
pCx(()= = Σ
log
log
Σ exp(aii + B (riz))
Ee N expl + El
1,
exp(os.vs + PT., 90 (Tij)) Ej exp + B 90 (11)
(n1/nd) expl + c, (gg
Σ
log
log
(m jexp + (e
1,
exp(as.v1 + PT. 90 (x(r)))
(expl1) + E, h (go E Dis explri + El (90
log
log
Eg exp(aai 31 (xir))
1,7
Σ log
74
As result, the loss function remains additive in two cross-entropy terms corresponding different tasks. The key difference from the covariate shift case is that, for the client-llassifictton task, each client now has its own linear head.
E SYSTEM ACCURACY CONVERGENCE RATE TRADE-OFF
We show the proof of Theorem 2.5 in this section. To simplify the notation, we consider the follow- ing loss: To assure strong convexity, we minimize the following objective function:
17113 (C)) 2
1
A
log exp(7% 2ig) exp(1,) 2gj)
P 113311
A
log
exp(3 mzg)
Miy
T2
exp(3 xzj)
1=1
=: elentt (7)
(9)
Then, where - (7, 3) stacks the parameters for client-classificttion y and task-classification 3. the global objective (C( m- Σm (C))
Problem setting: Assume the data is generated according to the true multinomial logistic model with parameters <irue (nt game). i.e.
exp(>;T2> 2)
exp(337 (3%T x) Pr(Yalsss = k|z) = ΣK exp(8;:T
Pr(Yetient = =|2)== = Em exp(1;T2) 22)
Let ČN denote the minimizer of (C) with N = Em n, as the total number of samples.
Total error. Let S" be the output of the algorithm after T steps. We decompose the error using the triangle inequality as follows:
- Srw < || 1l + 16 1
(8)
optimization error
statistical ervor
We know bound these two terms respectively.
18

---

Published as a conference paper at ICLR 2026
Lemma E.1 (Asymptotic normality). As N cc, the estimator ÇN satisfies
VN 5 d N10,, I((trae))))
where the Fisher information is block diagonal:
I(c = - ( + pT
0 XIs pT]
with
I, = E (<)) - @ (zz IB BB= ( - Psp(())))))))))>)) p=(())p(())) where p=(x)) (exp( x E, exp(3] .),.... (x)( exp().. E, exp(8]x)) and is the identity matrix.
@ (xxT))
Proof. This result follows from the well-established asymptotic properties of maximum likelihood estimators (Van der Vaart 2000 Section 5.5).
Statistical error. By Lemma B.1 we have
Ni 102 Op (1 A) llmin + +p): N11 1 07( A min + )): (9)
where ||A|lmin Amann (A)) the operator norm, d and p are dimensions of 2 and x respectively.
Optimization error, For communication round & = 0, 1, ,.. T client i sets Gh 5² and performs E' local gradient steps:
L, the server holds Ç² and each
Si (Ct.))
O....,, 1.
After E steps each client returns f and the server aggregates g++1 = m EiTh C.
Define (c) m- E= (C((() (() It can be decomposed nicely as G2(() = (1
G) elient (7) + X class (3). Let G2 G2 clieus and G class denote the corresponding maximum values
across updating rounds = 0, 2,,,...,?7
Then, G² s (1- cliem cliem + X²C² class
In the convergence proof below, omit the subscript p since it does not influence the convergence rate. Because is u-strongly convex and L-smooth, single full-gradient step satisfies
xx - V)(x)) - EN 12 =|| 2v111 2n(Ve(<>) x (w) + (V 1/
<||i 2x -2mull= 2w + 2w
<(1 //]l| 22.
For client at local step r::
Ilct - = - nVV((t)))))) D2(
(1-mm)llsi, 2x113 + ((IVe(()) - Ve((E))))
Iterating over E local steps gives
IC - Ex < (1- 11) 222 + 7 E - ve(
Averaging over 1..... m and using convexity of squared norm:
Æ-1
- < (1 1|c E. l1 + 1 Σ
1
m
r=0
19

---

Published as a conference paper at ICLR 2026
Using L-smoothness and nLE I 1/4, one show (via induction and triangle inequalities)
1
where 6 is the heterogeneity measure. Summing over r = =0,... E --1 gives
E- rm0
C..
i=1
Combine the above:
15 1 12 112 A (1 115 E.v 0P.
Let 8t = 1155 - -īx 11 and a == (1 - (1) E B = E²² Then
T-1
B - a T)
B
81+1 S 08 + B
8T S a² 80
B
Σ
a² = a +
1 Sa 80 +
1-a
j=0
This yields the desired bound
E2
IK<T -2vl12 s
(1 (1
-iNl11 ET) - 1 (1 mp)
(10)
Since 1 1- (1 mn1B 2 1 E > min nuE). the steady enror is of order
On12 F.62 (n) O(y i.e., FedAvg converges linearly to neighhorhood of radius proportional to /nEG2/H /n.
Combining and 1T0p with K gives the final result that
115² - ttr Op ((1 ((1- II] min + p}- +
[AllZollmio + +p-
E2
N
(1 mu)
as both T, N - co, This along with @² S (1 A) G) client + X2 class completes the proof of the theorem.
F EXPERIMENT DETAILS
F.1 VISUALIZATION OF COVARIATE SHIFT AND LABEL SHIFT
In the main experiment, we simulate covariate shift by applying three distinct nonlinear transfor- mations to each client's is dataset. Specifically, we use gamma correction with 7 {0.6, 1.4). hue adjustment with Ah { (00.1, 0.1}. and saturation scaling with x E {0.3, 1.5). This creates 2 = unique combinations of transformations, corresponding to 8-client setting where each client pos- sesses a visually distinct data distribution. A visualization of single image sampled from CIFAR- 10 after applying these transformations is shown in Fig. g As can be clearly seen, the resulting differences in feature distributions across clients are visually striking, highlighting the significant covariate shift simulated in our experiments. We visualize the two types of label shift used in the main experiment in Fig. D The figures show the number of samples from each class across 8 clients. As observed, the 5-SPC setting assigns at most 5 classes each client, whereas the Dir-0.3 setting distributes more classes per client. Thus, the label shift under Dir-0.3 is less severe than under 5- SPC. Our experimental results confirm this observation: all methods achieve higher performance under the less severe Dir-0.3 case.
F.2 SENSITIVITY ANALYSIS DETAILS
For all subsequent sensitivity analyses, unless otherwise specified, we use CIFAR-10 under the Dir- 0.3 setting. The details of each experiment are provided below.
20

---

Published as a conference paper at ICLR 2026
Covariate Shift Main Experiment
1D:01 y.O.G, 06:01, KOS
10:1 4. 0.1 K005
ID:2 y: A6011 K00
ID:3 v114, 4h..1. v:0.5
DDdd y.u.s. 10-0.1 K1S
0::5 v14. 0h:0.11 Ki5
DD y.u.b, 0.011 k1s
1D7/ 11144 0h0.1. KL
Figure 8: Visualization of a sample from CIFAR-10 under various nonlinear transformations.
Label Shift
Dir-0.3
Label Shift 5-SPC
Client
Client
02
4000
Sample Num
Sample Num
Figure 9: Visualization of client data distribution on CIFAR-10 under Dir-0.3 and 5-SPC settings.
F.2.1 NUMBER OF CLIENTS
To investigate the impact of the number of clients, we adopt more fine-graines strategy for simulat- ing covariate shift. We expand the parameter space for each nonlinear transformation to three distinct values: gamma correction with and saturation scaling with & E {0.4, 1.0, 1.6]. Furthermore, we introduce an additional binary transformation, posterization, which reduces the number of bits for each color channel to create flattening effect on the image's color palette. A visualization of these transformations is presented in Fig. 10 In the n-client setting, we apply the first n transformations from this pool.
{0.6. 1.0, 1.4). hue adjustment with Ah E {-0.15, 0.0, 0.15),
In our experiments, we set the maximum number of clients to 32. This is due to two primary challenges. First, the number of clients increases, the amount of data partitioned to each client diminishes significantly. This data scarcity creates scenario where fie-ttnnig-baase methods gain an inherent advantage, as each client's local train and distributions are identical. Second, is hard to design simulation strategy for covariate shift that is both sufficiently distinct and aligned with the model's inductive bias when the number of clients becomes very large.
F.2.2 COVARIATE SHIFT INTENSITY
To evaluate the robustness of our method under varying degrees covariate shift, we construct three intensity levels- mid, and high-by adjusting the parameter ranges of the nonlin-
ear transformations. The specific value ranges for each level are detailed as follows: (1) Low:
[0.9, 1.1), Ah E 0.0.0,, k (0.9, 1.1). (2) Mid: E {0.75, 1.25), Ah
(-0.05, 0.05), x E 1.3.. (3) High: E {0.6, 1.4). Ah E (-0.1.0.1), & E {0.5, 1.5). Visualizations corresponding to these levels are presented in Fig. II Fig. [2 and Fig. 8 It can be seen that the induced covariate shift is nearly imperceptible at the low level and escalates to a stark distinction at the high level, clearly illustrating the progressive intensity of the shift.
21

---

Published as a conference paper at ICLR 2026
Covariate Shift Sensitivity to Number of Clients
013 pohie
840 lus
91d eushd yohbe
Mait wi
Iold tn
1932 st
0211 22
(81) aLAnd
Cill
00
Figure 10: Visualization of CIFAR-10 sample under covariate shift with larger number of clients.
Covariate Shift Low Level
1b:0 I y.u.9. M--0.1, ad..
ID:1 p33.1, Mh-0011, .0.9
ID v:l.9, Mr0.01, 000
ID:3 y111 Mo..0, wa.
y1.1.
y11
Figure 11: Visualization of a sample from CIFAR-10 under low covariate shift intensity.
F.2.3 BACKBONE SHARING STRATEGY
In our formulation, the target-class classification task uses the embedding go(x) for an input feature while the client-classffication task uses the embedding h (z)) for the same feature. Since both go and h are parameterized functions, the optimal sharing strategy between the two not obvious. To explore this, we investigate four nackbone-sharing strategies based on LeNet: no sharing. shallow sharing, mid sharing, and deep sharing. The network architectures are illustrated in Fig. 13 From (a) to (c), the discrepancy between the embeddings for the two tasks decreases, while the number of learnable parameters also reduces. Our empirical results in Fig. show that all strategies perform similarly, with shallow sharing slightly ahead. However, given the substantial increase in parameters for shallow sharing, deep sharing offers a more parameter-efficient alternative while maintaining strong performance.
22

---

Published as a conference paper at ICLR 2026
Covariate Shift Mid Level
1D:0 ydu.5, 6h-00005 kdd.
iD:1 y:1.2% 6h--0.05, xt0.7
ID:2 yy0.75, dho x:07
ID:3 y 06.0.55 402
40.005
Figure 12: Visualization of sample from CIFAR-10 under mid covariate shift intensity,
Cunva
FC1
Target
FC1
Target
Ipplt
Inpul
Convi
Convt
Convz
FC1
FC2
Client Head
Conv2
FC1
FC2
Cleen
(a) No Sharing
(b) Shallow Sharing
FCT
Target Head
Eorget
Inpul
Convt
Carso
Inpul
Convi
Canvo
FC1
FCD
FC1
FC2
Cient Head
Cilent Head
(c) Mid Sharing
(d) Deep Sharing
Figure 13: Visualization of the four different parameter sharing strategies.
Labal Shift Dis0.3
Label Shift Dir113
Labal Shift Di-2.3
2000
4000
5000
5030
2000
4000
0000
5500
1000
2000
30E0
4000
5500
5000 7010
Sample Num
Sample Num
Sample
Figure 14: Visualization of client data distribution under Dir-0.3/1.3/2.3 settings.
Table System accuracy and average accuracy under different Dirichlet parameter values.
System Accuracy
Average Accuracy
Method
0.3
1.3
2.3
0.3
1.3
2.3
Ditto
47.64
0.25
43.78 + 0.23
44.04 0.20
76.34 + 0.11
55.56 + 0.18
52.76
0.18
FedRep
24.96
0.19
35.45 1 0.22
37.87 0.22
76.49
1 0.15
55.09
1 0.17
52.95
0.16
FedBABU
57.43 0.17
54.60 + 0.19
54.25 = 0.20
+
61.61
+ 0.16
59.83 = 0.20
FedPAC.
2.21
35.53 + 0.21
37.84 + 0.20
76.53+0.33
55.08 + 0.15
52.95 0.16
FedALA
61.33
48.47 0.18
46.68 = 0.20
64.35
49.14 0.62
47.29 0.50
FedAS
28.76 0.119
41.43 0.20
46.08 1 0.22
0..19
59.83 1 0.19
58.17 0.18
ConFREE
25.66
96.12
38.78 + 0.23
76.73
55.56
+ 0.18
53.58 = 0.15
FedAvgITT
54.90 0.22
$4.00000
55.61 + 0.19
79.08 0.11
62.87
1 0.16
61.94 0.17
FedProxFT
55.01 0.20
0.15
55,62 - 0.18
79.07 0.12
0.11
61.86 0.19
FedSAMFT
55.83 0.21
49.15 1 0.17
47.71 0.16
75.53 0.111
55.46 1 0.19
53.54 0.15
FedDRM
63.85 0.18
56.83 0.18
56.00 + 0.23
80.25 0.14
64.04 + 0.16
62.30 + 0.15
23

---

Published as a conference paper at ICLR 2026
F.2.4 LABEL SHIFT INTENSITY
To evaluate the robustness of our method under varying degrees of label shift, we compare our method with the baselines across a range of Dirichlet parameters a E {0.3, 1.3, 2.3}. Visualiza- tions corresponding to these settings are presented in Fig. 14 The corresponding results for system accuracy and average accuracy are presented in Tab,
Consistent with prior work (Xu et al., 2023). we can see that smaller a values corresponding to higher data heterogeneity- lead to higher average accuracy for all methods. This occurs because each client's training and testing data are drawn from the same distribution. As a decreases, the local label distributions become increasingly skewed, with some classes receiving negligible probability mass. This effectively reduces the number classes present on each client, thereby simplifying the local classification problem relative to the balanced case. For system accuracy, we find that this trend persists for our method, as the only additional component is the client-routing step. which does not alter the underlying behavior of local classification. In contrast, methods such FedRep exhibit increasing system accuracy as a grows (i.e., as the label distributions become more homogeneous). When is small, the local models become highly personalized and fail to reach consistent consen- sus across clients, causing majority voting to misroute queries and thus lowering system accuracy. As a increases, this inconsistency diminishes, and the aggregated routing accuracy improves. These results further confirm that our method remains robust across varying degrees of label shift.
24


Reviews:
[{"review_id": "z5wJSmKq0A", "reviewer": "Reviewer_avg5", "summary": "This paper proposes FedDRM, a federated learning (FL) framework that extends beyond traditional aggregation to enable the central server to actively route new queries (e.g., patient data) to the most suitable client in heterogeneous environments. Drawing from a healthcare motivation—directing patients to specialized hospitals—the authors model client data distributions using a semiparametric density ratio model (DRM) anchored to a nonparametric empirical likelihood (EL) baseline at the server. Under assumptions like covariate shift, they derive exponential tilts linking client distributions (Eqs. 2-3) and extend this to marginals (Theorem 2.1). The profiled EL loss decomposes into cross-entropy terms for label prediction and client identification, jointly optimizing local accuracy and routing precision. Experiments on CIFAR-10/100 with Dirichlet heterogeneity (α=0.3) compare against baselines like FedAvg, reporting gains in \"system acc\" (routed performance) and guidance metrics, though results appear inconsistent across tables.", "questions": "1. The authors only report the results when Dirichlet (dir) α=0.3 (low heterogeneity; near-IID), but FL's real pain is high heterogeneity. How does FedDRM scale here? Could the authors add analysis across more α values or extend to more datasets to show robustness?\n2. Could you please explain the experimental details? For example, what is the number of clients in the two main tables? I found that the results seem best fit with m=8; however, the results are not well-aligned. \n3. Could the author please provide privacy guarantees for the additional uploaded parameters?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The core idea of reframing heterogeneity as an asset—via server-guided routing— is a genuine advance over passive FL methods. By turning the server into an \"intelligent router\" (Fig. 1), it addresses a practical gap: e.g., in finance, routing fraud queries to distribution-matched banks, or in healthcare, assigning patients to expert clients. This not only boosts efficiency but also highlights FL's untapped potential for adaptive, privacy-preserving systems.\n\n2. The DRM-EL framework elegantly unifies heterogeneous training with query matching. The use of multiplicative tilts (Eq. 3) to capture shifts without raw data sharing is nonparametric yet efficient, and the loss decomposition (label + client CE) provides a clean, interpretable objective. Theorem 2.1's extension to marginals under covariate shift adds rigor, making the method theoretically grounded while extensible to label/full shifts.\n\n3. Despite limitations, the reported \"system acc\" improvements (e.g., leveraging routing for end-to-end gains) demonstrate the framework's promise over aggregation-only baselines. This dual focus—accuracy + precision—positions FedDRM as a step toward resource-efficient FL, with broad applicability in multi-client domains.", "weaknesses": "1. The introduction builds excitement around the routing idea but omits any preview of the DRM or EL machinery, leaving readers disoriented when section 2 dives straight into FedDRM. What is DRM? A brief 1-2 sentence teaser in the intro would bridge this.\n\n2. While results show promise, the scope is too narrow to substantiate claims under severe heterogeneity. You only report Dirichlet (dir) α=0.3 (low heterogeneity; near-IID), but FL's real pain is high heterogeneity. How does FedDRM scale here? Add analysis across α values or extend to more datasets to show robustness.\n\n3. What is the number of clients in the two main tables? I found that the results seem best fit with m=8; however, the results are not well-aligned. The system acc for FedDRM is 59.59% under dir=0.3, which is significantly lower than the result in the first main table.\n\n4. FedDRM requires uploading additional density ratio parameters (e.g., tilts γ_i, ξ_i in Eq. (3), lines 144-147) alongside model updates, potentially undermining FL's privacy guarantees (e.g., vs. FedAvg's aggregated-only sharing, lines 32-35). These tilts encode distributional shifts (e.g., π_i margins), risking inference attacks on sensitive traits like demographics (Fig. 1, lines 45-60).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "P33jhMZLr4", "reviewer": "Reviewer_Yf1W", "summary": "This paper studies the problem of heterogeneity among clients in federated learning. The paper classifies the heterogeneity into two categories: i) Covariate shift ii) label shift. Then the paper investigates the problems of data heterogeneity in federated learning which are combination of these types heterogeneity.To address this problem, the paper defines new local loss functions for clients. Then using some adding heads to a base shared or global model, the federated learning is performed. Experiments on image classification is conducted to show the effectiveness of the proposed algorithms.", "questions": "- While reading the paper, I found it unclear why Theorem 2.1 is important. I suggest adding more explanations to clarify its significance, or alternatively, moving it to the appendix. Could you elaborate on this point?\n- Could you also elaborate on the relevance of data modeling in the experiments to the problems discussed in the introduction?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The paper proposes more responsibilities for the server in federated learning such that the server guides the clients during training. This idea is interesting and can leads to improvements.\n- The convergence rate of the proposed algorithm is obtained for strongly convex loss functions.\n- The proposed algorithm is compared against variety of federated learning baselines.", "weaknesses": "- The most important weakness of the paper is presentation. The importance and novelty of the paper’s contributions are not clearly conveyed. I recommend that the authors explicitly and concisely list the main contributions of the work, such as the specific problems that the proposed algorithm aims to solve. At present, there is no clear evidence or justification demonstrating that the proposed algorithm effectively addresses the key issues claimed in the paper.\n- The paper focuses solely on the multi-class classification problem, which is a relatively minor limitation.\n- The discussion following Theorem 2.5 does not discuss if the proposed algorithm improves the convergence rate compared to existing methods.\n- The data heterogeneity modeled in the experiments is synthetic. While this is common in federated learning research, the introduction suggests that the paper targets particular types of heterogeneity. Because of this, the practical applicability of the studied problem is not clear.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VX2nKiT19y", "reviewer": "Reviewer_C2JH", "summary": "This paper proposed a FL paradigm named FedDRM, that treats client heterogeneity as an asset rather than a limitation. \nBy leveraging an empirical likelihood-based framework, FedDRM simultaneously trains personalized local models and enables the central server to intelligently route new queries to the most suitable client. \nHowever, the experiments are limited to CIFAR benchmarks, which are small scale and fewer than existing methods. They may be insufficent to demonstrate the superity of the proposed method.", "questions": "1. Could the client-classification mechanism be distilled into a lighter-weight model to reduce inference overhead?\n\n2. How could FedDRM adapt to dynamic client populations? For example, the data distributions evolve over time.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The idea of leveraging heterogeneity for query routing instead of merely mitigating it, seems interesting and offer a novel perspective in FL.\n\n2. The empirical likelihood framework provides a principled foundation for joint model training and client matching, with a tractable dual loss formulation.", "weaknesses": "1. The only dataset is CIFAR (including cifar 10 and cifar 100), which is significient different from the data in domains such as healthcare and finance. This limits the support for claims about the method’s potential in such areas.\n\n2. The client-classification head may face challenges when scaling to very large client pools (let's say, >100 clients), as gradient drift could become worsen.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)