Addressing Label Shift in Distributed Learning via Entropy Regularization​
Authors: Zhiyuan Wu
Changkyu Choi
Xiangcheng Cao
Volkan Cevher
Ali Ramezani-Kebrya
Venue: ICLR 2025 Poster (2025)
Keywords: label shifts; test-to-train importance ratio estimation for label distributions; entropy-based regularization; distributed learning
Primary area: other topics in machine learning (i.e., none of the above)

Abstract: We address the challenge of minimizing "true risk" in multi-node distributed learning.\footnote{We use the term node to refer to a client, FPGA, APU, CPU, GPU, or worker.} These systems are frequently exposed to both inter-node and intra-node "label shifts", which present a critical obstacle to effectively optimizing model performance while ensuring that data remains confined to each node.
To tackle this, we propose the Versatile Robust Label Shift (VRLS) method, which enhances the maximum likelihood estimation of the test-to-train label importance ratio. VRLS incorporates Shannon entropy-based regularization and adjusts the importance ratio during training  to better handle label shifts at the test time.
In multi-node learning environments, VRLS further extends its capabilities by learning and adapting importance ratios across nodes, effectively mitigating label shifts and improving overall model performance. Experiments conducted on MNIST, Fashion MNIST, and CIFAR-10 demonstrate the effectiveness of VRLS, outperforming baselines by up to 20\% in imbalanced settings. These results highlight the significant improvements VRLS offers in addressing label shifts. Our theoretical analysis further supports this by establishing high-probability bounds on estimation errors.

Published as a conference paper at ICLR 2025
ADDRESSING LABEL SHIFT IN DISTRIBUTED LEARN- ING VIA ENTROPY REGULARIZATION
Zhiyuan Wu"
Changkyu Choi'
University of Oslo zhi zhiyuanw@ wiraaw@@ffi.m uio.
UIT The Arctic University of Norway
changkyu.choie
Xiangcheng Cao EPFL xi xiangcheng cao. cao.eppflgmaiii
Volkan Cevher
LIONS, EPFL
volkan. cewhereep..eeeeep...... ch
Ali Ramezani-Kebrya
Department of Informatics, University of Oslo
Norwegian Centre for Knowledge-driven Machine Learning (Integreat)
aliłuio.no
ABSTRACT
We address the challenge of minimizing true risk in multi-node distributed learn- ing These systems are frequently exposed to both inter-node and intra-node label shifts, which present critical obstacle to effectively optimizing model performance while ensuring that data remains confined to each node. To tackle this, we propose the Versatile Robust Label Shift (VRLS) method, which enhances the maximum likelihood estimation of the test-to-train label importance ratio. VRLS incorporates Shannon untopy-bassed regularization and adjusts the importance ratio during training to better handle label shifts at the test time. In multi-node learning environ- ments, VRIS further extends its capabilities by learning and adapting importance ratios across nodes, effectively mitigating label shifts and improving overall model performance. Experiments conducted on MNIST, Fashion MNIST, and CIFAR- demonstrate the effectiveness of VRLS, outperforming baselines by up 20% in imbalanced settings. These results highlight the significant improvements VRLS offers in addressing label shifts. Our theoretical analysis further supports this by establishing high-probability bounds on estimation errors. The code available at http:/hhiyuunnie httpps:/ggt thub. com/ zhi yuan- 1uan-11/VRI..... http///SSssii/titeeeeeeeeee
INTRODUCTION
The classical learning theory relies on the assumption that data samples, during training and testing, independently and identically distributed (i.i.d.) drawn from unknown distribution. However, this i..dd assumption is often overly idealistic in real-world settings, where the distributions of training and testing samples can differ significantly and change dynamically as the operational environment evolves. In distributed learning (Kim et al. 2022: Wen et al. [2033] Ye et al., 2023 Luo et al. 2023)
where nodes retain their own data without sharing, these discrepancies actoss nodes become more pronounced, further intensifying the learning challenge (Rahman et al. 2023 Wang et al. 2023
Label shiftts (Lipton et al.) 2018 Garg al. 2022. Mani al. 2022 Zhou et 2023) represent form of distributional discrepancy that arises when the marginal instribbtion of labels in the training set differs from that the test set, i.e., nd (y) 7 (g)) while the conditional distribution of features given labels, p( y). remains largely stable across both datasets. Label shifts commonly manifest both inter- node and intra-node, complicating the learning process in real-world distributed learning scenarios. However, a commonly used learning principle in this distributed setting, empirical risk minimization
"These authors contributed equally to this work.
Wee use the node refer client, FPGA, APU, CPU, GPU, or worker.

---

Published as a conference paper at ICLR 2025
(ERM) (Kur al. 2024), operates under the assumption that the training and test distributions identical on each node and across nodes. This overlooks these shifts, failing to account for the statistical heterogeneity across decentralized data sources. While the current literature (Yin et al.) 2024) addresses statistical heterogeneity across nodes, it often neglects distribution shifts at test or operation time, which has been significant challenge in the entire data science over decades.
The primary technical challenge in addressing label shifts lies in the efficient and accurate estimation of the test-to-train importance ratios, expressed as p" (y) (y) for all labels. A widely popular solution is Maximum Likelihood Label Shift Estimattio (MLLS) (Garg et al., 2020), which frames this estimation as a convex optimization problem, akin the Expectatio-Maxiiization algorithm (Saerens et al.) 2002). Model calibration refers to the process of ensuring that predicted probabilities reflect the true likelihood of correctness, which crucial for improving the accuracy of importance ratio estimation (Guo et al., 2017a; Garg et al. 2020). Bias-Correcter Calibration (BCT) (Alexandari et al., 2020) serves as an efficient calibration method that enhances the EM algorithm within MLLS.
Maximization (EM)
While BCT and other post-hoc calibration techniques (Guo et al. 2017c; Kull et al. 2019 Wang et al. 2021 [Sun et al., 2024) contribute to improved calibration and may potentially improve model performance, their primary focus remains on refining classification outcomes rather than on accurately approximating the true conditional distribution p"(y/ z). The "predictor" in these literature captures the relationship between the input features x: and the corresponding output probabilities across the labels in the discrete label space y, with [>] m, which should approximate the true distribution of p" ("(())))) Despite this goal, training with conventional cross-entropy loss often leads models that produce predictions that are either highly over-confident or under-confident, resulting in poorly calibrated outputs (Guo et al. (017a). Consequently, the predictor fails to capture the
underlying uncertainty inherent in p" [ x which limits its effectiveness in estimating importance
ratios (Alexandari et al. 2020 Garg et al. 2020
Guo et al., 2020: 2017b) Pereyra et al., 2017
McMahan al.) 2017).
To address this limitation, we propose novel Versatile Robust Label Shift (VRLS) method, specifi- cally designed improve importance ratio estimation for tackling the label shift problem. A key idea of our VRLS method is to approximate y" (z((x) in way that accounts for the inherent uncertainty over the label space y for each input æ. Accordingly, we propose a new objective function incorpo- rating regularization to penalize predictions that lack proper uncertainty calibration. We show that training the predictor in this manner significantly reduces estimation eтror under various label shift conditions.
Building upon our VRLS method, we extend its application to multi-node settings by proposing an Importance Weighted-ERN (IW-ERM) framework (Fang al., (023). Within the multi-node distributed environment, our IW-ERM aims to find an unbiased estimate of the overall true risk minimizer across multiple nodes with varying label distributions. By effectively addressing both intra- node and inter-node label shifts with generalization guarantees, our framework handles the statistical heterogeneity inherent in decentralized data sources. Our extensive experiments demonstrate that the IW-ERM framework, which trains predictors exclusively on local node data, significantly improves overall test error. Moreover, maintains convergence rates and privacy levels comparable to standard ERM methods while achieving minimal communication and computational overhead compared to existing baselines. Our main contributions are as follows:
We propose VRLS, which enhances the approximation of the probability distribution p"(y) by incorporating a novel regularization term based on Shannon entropy (Neo et al.) 2024). This regularization leads to more accurate estimation of the test-to-train label importance ratio, resulting improved predictive performance under various label shift conditions.
By integrating our VRLS ratio estimation into multi-node distributed learning environment, we achieve performance close to an upper bound that uses true ratios on Fashion MNIST and CIFAR-10 10 datasets with 5, 100, and 200 nodes. Our IW-ERM framework effectively manages both inter-node and indra-node label shifts while remaining data confined within each node, resulting in up to 20% improvements in average test error over current baselines.
We establish high bighabrobability estimation error bounds for VRLS, as well as brobbbbbbybilityy conver- gence bounds for IW-ERM with VRLS in nonconvex optimization settings (Section Appendices
2

---

Published as a conference paper at ICLR 2025
Table I: Details of the label shift scenarios. Their IW-ERM formulas are presented in Appendix C
Scenario
#Nodes
Assumptions on Distributions
Ratio Node Needs
No-LS in Equation 17
2
p((v) - =( v5(()) p5(() (F)
p((v)/pH(u)
LS on single in Equation 18
pl(()) #
LS on both Equation
pf(v) ( pi(y) + (T) pl()))p( Pftl) pk(v) # (1) for all l
LS on multi in Equation 199
K
((lv)/H(()) for all
Additionally, we demonstrate that incorporating importance weighting does not negatively impact convergence rates or communication guarantees across various optimization settings.
IMPORTANCE RATIO ESTIMATION AND IMPORTANCE WEIGHTED-ERM
Importance ratio estimation Importance ratio estimation for label shifts has been addressed by methods such as solving linear systems (Lipton et al., 2018 Azizzadenesheli et al., (019) and mini- mizing distribution divergences (Garg et al. 2020). primarily in the context of a single node. Lipton et al. 2018); Azizzadeneshelii et al. (2019); Garg et al. (2020) assumed the conditional distribution p(a y remains fixed between the training and test datasets, while the label distribution p(y) changes.
Black Box Shift Estimation (BBSE) (Lipton et al., 2018 Rabanser et al. 2019) and Regularized Learning under Label Shift (RLLS) (Azizzadenesheli et al. (2019) are confusion matrix-based meth- ods for estimating importance ratios label shift problems. While BBSE has been shown consistent even when the predictor is not calibrated, its subpar performance is attributed to information loss inherent in using confusion matrices (Garg et al., (020). To overcome this, Garg al. (2020) has introduced the MLLS, resulting in significant improvements in estimation performance, especially when combined with post-hoc calibration methods like BCT (Shrikumar et al., 2019). This EM algorithm based MLLS method Saerren et al. 2002 Garg et al. 2020) is concave and can be solved efficiently.
Importance Weighted-ERM, Classical ERM seeks to minimize the expected loss over the training distribution using finite samples. However, when there a distribution shift between the training and test data, the objective of ERM is not to minimize the expected loss over the test distribution, regardless of the number of training samples. To address this, IW-ERM is developed (Shimodaira) 2000 Sugiyama et al. 2006 Byrd & C. Lipton| 2019; Fang et al. 2020), which adjusts the training loss by weighting samples according to the importance ratio, i.e., the ratio of test-to-train density, Shhmodaira (2000) has shown that the IW-ERM estimator is asymptotically unbiased under certain conditions. Building on this, Ramezani-KKoa which incorporates importance ratio estimation handle covariate shifts in distributed learning. However, this approach has limitations, as does not address label shifts and the importance ratio estimation method poses potential privacy risks.
et al.] (023b) have recently introduced Federated IW-ERM,
In this work, we focus on label shifts and propose an IW-ERM framework enhanced by our VRLS method. We show that our IW-ERM with VRLS performs comparably to an upper bound that utilizes true importance ratios, all while preserving data privacy across distributed data sources. This approach effectively addresses both intra-node and inter-node label shifts while ensuring convergence in probability to the overall true risk minimizer.
3 VERSATILE ROBUST LABEL SHIFT: REGULARIZED RATIO ESTIMATION
In this section, we introduce the Versatile Robust Label Shift (VRLS) method for importance ratio estimation in single-node setting, which forms the basis of the IW-ERM framework. To solve the optimization problem of IW-ERM, each node l requires an accurate estimate of the ratio:
(k(v) = Σ/ (1)
(1)
where p)J (z) and (3) represent the test and training label densities, respectively. To improve clarity and avoid over-complicating notations, first consider the scenario where we have only one
3

---

Published as a conference paper at ICLR 2025
Algorithm 1 VRLS Importance Ratio Estimation Algorithm
Require:
Labeled training data {(x,, v:))
Require:
Unlabeled test data {æ;}}=ı-
Require:
Initial predictor fo- Optimized predictor fe- and estimated importance ratio ry.
Ensure:
I:
Training:
2:
Optimize fe using Equation 4 via SGD.
3:
Continue until the training loss drops below a threshold or the maximum epochs are reached.
4: Obtain the optimized predictor fo-.
5:
Importance Ratio Estimation:
With optimized predictor fo-, estimate importance ratio ry- using equation Equation (
node under label shifts and then extend to multiple nodes. Table presents various scenarios. In single-node label shift scenario, the goal is to estimate the ratio r(y) (y)) / (y). Following the seminal work of Garg al. (2020), we formulate importance ratio estimation as a Maximum Likelihood Estimation (MLE) problem by constructing an optimization problem based on Kullback- Leibler (KL) divergence to directly estimate r(y). We train a predictor fe to approximate p"(y)æ), where 0 denotes the parameters of neural network. After training, we apply the predictor fo- to finite set of unlabeled samples drawn from the test distribution to obtain predicted label probabilities. These predictions are then used to estimate the ratio r Further details are provided in Algorithm
One of the novelties of VRLS is its ability to better calibrate the predictor, enabling to better approximate the true conditional distribution y ( az). This approximation faces two main chal- lenges, as highlighted in Theorem 3 of (Garg et al.) 2020): finite-sample error and miscalibration error. Entropy-basee regularization can directly tackle miscalibration, which occurs when predicted probabilities systematically deviate from true likelihoods. Building on these insights, we introduce an explicit entropy regularizer into the training objective, which is based on Shannon's entropy (Pereyra et al. 2017
Neo et al. 2024). The regularization term l(fe) is defined as:
N(fe) = = @ff((z))) log (o(folz)).)
m
where & denotes the softmax function, and c represents the t element of the softmax output vector. With this regularization the softmax outputs, VRLS encourages smoother and more reliable predictions that account for inherent uncertainty in the data, leading to more accurate importance ratio estimates and improving the SotA in practice. These improvements are empirically demonstrated in Section 6 Our proposed VRLS objective is formulated as follows:
Tf arg max Ene [log[fo- [r)] FER
where
0- = arg min Ea (CC((((())))))))) (R((o)]
(4)
The vector r in Equation representing the importance ratios for all m classes, belongs to the non-negative real space R This constraint set defined similarly to MLLS (Garg et al. (022). and we use the expected value Fe for estimation, denoting the optimal importance ratio ry. To train the predictor 0, we minimize the cross-entropy loss fCE together with a scaled regularization term (S((fe), where Ç > 0 is coefficient controlling the regularization strength. Incorporating the regularizer (2(fe) improves the model calibration under the influence of fe loss.
4 VRLS FOR MULTI-NODE ENVIRONMENT
We now extend VRLS to the multi-node environment, taking into account the privacy and commu- nication requirements. This extension naturally aligns with the concept of IW-ERM, effectively integrating these considerations into the multi-node learning paradigm. We consider multiple nodes where each node has distinct training and test distributions. The goal here is train global model
4

---

Published as a conference paper at ICLR 2025
Algorithm 2 IW-ERM with VRLS in Distributed Learning
Require: Labeled training data i VE, Jin, each node k. for [K].
Require: Unlabeled test data [ at each node k. for k = [K].
Require: Initial global model hhe.
Ensure: Trained global model /ow optimized with IW-ERM.
1: Phase 1: Importance Ratio Estimation with VRLS
2: for each node l = 1 to K in parallel do
3: Train local predictor fa.o on local training data (< i vE,))
4: Use fu,0- to estimate the importance ratio pk,s" on unlabeled test data {ack) } at node k.
5: end for
6: Phase Importance Ratio Aggregation
7: for each node 1 K do
8:
Aggregate importance ratio using Equation ①I
9: end for
10: Phase 3: Global Model Training with IW-ERM
11: Train global model Aw using Equation (W-ERM with the aggregated importance ratios.
that utilizes local data and addresses overall test error. In this setup, each node uses its local data to estimate the required importance ratios, as outlined Section B and shares only low-dimensional ratio information, without the need to share any local data.
The process begins with each node training global model on its local data, independently estimating its importance ratios. These locally computed ratios are then shared amongst the nodes, allowing for the aggregated ratio required for IW-ERM to be computed centrally. This aggregated ratio is then used further refine the global model in second round of global training. This approach ensures minimal communication overhead and preserves node data privacy, as detailed Section 5 Our experimental results in Section 6 demonstrate that the IW-ERM framework significantly improves test error performance while minimizing communication and computation overhead compared to baseline ERM. The importance ratio estimation and IW-ERM are described in Algorithm 2
To provide a more comprehensive understanding of the multi-node environment, the following discussion delves into its details. Let X' RRdo be a compact metric space for input features, y be discrete label space with | = m, and be the number of nodes in an multi-node setting 2 Let S₁ =
{(0%, y b denote the training set node & with n. samples drawn i.i.d. from probability distribution Pl on x y. The test data node is drawn from another probability distribution pl on X' x V. We assume that the class-conditional distribution pk (y) p (fxl)) (y) := p>x:lyy remains the same for all nodes k. This is a common assumption and holds when label shifts primarily affect labels' prior distribution of the labels p(y) rather than the underlying feature distribution given the labels, e.g., when features that are generated given a label remains constant (Zadrozny [0044] Huang
al. 2006 Sugiyama et al. 2007). Note that pf (y) and (g)
(y) can be arbitrarily different, which
gives rise to intra- and inter-node label shifts (Zadrozny) 2004 Huang et al. 2006 Sugiyama et al. 2007 Garg et al. 2023).
In this multi-node environment, the aim is to find an unbiased estimate of the overall true risk minimizer across multiple nodes under both intra-node and inter-node label shifts. Specifically, we aim to find hypothesis hw E H X + V, represented by neural network parameterized by w, such that hw provides a good approximation of the label y E y corresponding to a new sample X drawn from the aggregated test data. Let é X + R4 denote loss function. Node k aims to learn a hypothesis hap that minimizes its true (expected) risk:
(k(hw) = (( (e(ho y)].
(Local Risk)
"Sets and scalars are represented by calligraphic and standard fonts, respectively. We use [m]] to denote (I...., m) for an integer m. We use S to ignore terms up to constants and logarithmic factors. We use E|-] to denote the expectation and || to represent the Fuclidean norm of vector. We use lower-case bold font to
denote vectors.
5

---

Published as a conference paper at ICLR 2025
We now modify the classical ERM and formulate IW-ERM to find predictor that minimizes the overall true risk over nodes under label shifts:
min Aow EH
E(1 (.
(IW-ERM)
P P
where nig is the number of training samples at node k.
To incorporate our VRLS importance ratio estimation method into the IW-ERM framework, we replace ratio term E (.)..) with our estimated importance ratios. This modification aims to PR(VE..)
align the empirical risk minimization with the risk minimization over all nodes. We formalize the convergence of this approach in Proposition 4.1
Proposition 4.1. Under the label shift setting described Section equation [W-ERM is consistent and the learned function hw converges in probability towards the optimal function that minimizes the overall true risk across nodes, Ek=1 Rg--
Proof. Due to space limitations, the proof is provided in Appendix C Convergence in probability established by applying the law of large numbers following (Shimodaira) (000))Section) and Sugiyama el al. [Sectin 2.2].
5 RATIO ESTIMATION BOUNDS AND CONVERGENCE RATES
In this section, we present bounds on ratio estimation and convergence rates for the finite sample errors incurred during the estimation, as further discussed in Appendices E In practice, only have access to finite number of labeled training samples, {(x;, y ) and finite number of unlabeled test samples, { These samples serve to compute the following estimates:
n
0, = arg min BEB 72
(CE(fo(z)), v.) +(((())
and AM arg max
log S0_0 (arz) Tr).
PERm
We will show that the emror of these estimates can be controlled. The following assumptions necessary to establish our results. Assumption 5.1 (Boundedness). The data and the parameter space O are bounded, i.e, there exists bx, be > 0 such that
Va E X, ||æ||2 S bxx
and
4O E O, |00| < be.
Assumption 5.2 (Calibration). Let 0* be as defined in Equation
(4. There exists μr 0 such that
E [fe- (x)
> plm
The calibration Assumption 5.2 first appears in (Garg al. 2020 It is necessary for the ratio
estimation procedure to be consistent and we refer the reader to Section 4.3 of Gf et al. (2020) for more details. We further need Assumption 5. because, unlike (Garg al. 2020), the empirical estimator Pr estimated using another estimator 0,... Uniform bounds therefore needed to control finite sample error as we cannot directly apply concentration inequalities, as is done in the proof of (Garg et al., 2020 Lemma 3), since we do have independence of the terms appearing in the empirical sums. We nonetheless prove a similar result in the following theorem.
Theorem 5.1 (Ratio Estimation Error Bound). Let 10, and == {ar + T fo(x), (r. E R. x θ}. Under Assumptions 5.1 1.2 there exist constants L > 0, B > 0 such that with probability
least
2
4
log( log(4/0g (4/5) n2"
4L pPain
A r5- pPmin
Rad(F) 4B
+
E E
0" [2]
(5)
Viye

---

Published as a conference paper at ICLR 2025
Here, Pmii miny ply) and
Rad(F
B 12
sup (r.)eRxe
Σ oir fo(az))
(6)
where T1y...,0 are Rademacher variables uniformly chosen from (-1,1}.
Proof. The proof of Theorem 5.1 is provided in Appendix E The Rademacher complexity appearing in the bound will depend on the function class chosen for f. Moreover as regularization often encourages lower complexity functions, this complexity can be reduced because of the presence of the regularization term in the estimation of & in our setting.
By estimating the ratios locally and incorporating them into local losses, the properties of the modified loss with respect to neural network parameters w remain unchanged, with data-dependent parameters like Lipschitz constants scaled linearly by Tmax- Our approach trains the predictor using only local data, ensuring IW-ERM with VRLS retains the same privacy guarantees as baseline ERM-solvers. Communication involves only the marginal label distribution, adding negligible overhead, as it is far smaller than model parameters and requires just round of communication. Overall, importance weighting does not impact communication guarantees during optimization.
Theorem 5.2 (Convergence- conprenccccoommmmmciiiiiono) Let maXgey sup, ry(y) = "max- Suppose Algo- rithum 2 e.g., IW-ERM with VRLS for multi-node environment, is run for T iterations. Then Algo- rithum accieves a convergence rate of (r (T)), where O(h(T)) denotes the rate of ERM--ollvr baseline without importance weighting. Throughout the course of optimization, Algorithm has the same overall communication guarantees as the baseline.
In the following, we establish tight convergence rates and communication guarantees for IW-ERM with VRLS in broad range of importance optimization settings including upper and lower bounds for convex optimization (Theorems [.3 5.4). second-orde differentiabilityy composite optimization with
proximal operator (Theorem
F33) optimization with adaptive step-sizes, and nonconvex optimization
(Theorems F.T F22). along the lines of e.g., (Woodworth et al. 2020 Haddadpour et al. 2021
Glasgow et al. 2022 Liu et al. 2023 Hu & Huang 2023] Wu et al. 2023] Liu et al.) (003).
Assumption 5.3 (Convex and Smooth). 1) minimizer w" exists with bounded ||w ||2: 2) The &o
3-smooohhess and convex w.r.t. w; 3) The stochastic gradient g( () V ((h is unbiased, i.e.,
E[g(w)] = V ((ww) for any w E W with bounded variance E [[[[g(w) (w) V (ww(h)))
For convex and smooth optimization, we establish convergence rates for IW-ERM with VRLS and local updating along the lines of e.g., (Woodworth et al., 2020 Theorem 2).
Theorem 5.3 (Upper Bound for Convex and Smooth). Let D ww w01| denote the number of local steps (number of stochastic gradients per round of communication per node). R denote the number of communication rounds, and maXyey sup, r
T (3) Tiaax- Under Assumption 5.3 with local updates is run for T = R total stochastic gradients per node with an optimally tuned and constant step-size. Then have the following upper bound:
suppose Algorithm 2
322
(D4N)3DAA113 (VFR)273
D
B(((herr) k(hw
+ VKTR
(7)
TR
Assumption 5.4 (Convex and Second-order Differentiable). 1) The (z)) y) is 3-smoothness
and convex w.r.t. w for any (x, y); 2) The stochastic gradient g(w) V Vw((hw) is unbiased. i.e. E[g(w)) V V (hr)) for any w W with bounded variance E ([]g(()) V., 112)-
Theorem 5.4 (Lower Bound for Convex and Second-order Differentiable). Let D Iw w'l T denote the number of local steps, R. denote the number of communication rounds. and maXyey sup r (y) Tmax. Under Assumption 5.4 suppose Algorithm with t local updates is run for T = R total stochastic gradients per node with tumed and constant step-size. Then we have the following lower bound:
T mmax 3 3D2
(Fmax.33 3D4)1/3
D
B(((hwr)) - ((D })
+
(8)
TR
(VTR)2/3
VKTR
7

---

Published as a conference paper at ICLR 2025
We finally establish high-probability convergence bounds for IW-ERM with VRLS along the lines of e.g., (Liu et al. 2023 Theorem 4.1). To show the impact of importance weighting on convergence rate decoupled from the impact of number of nodes and obtain the current SotA high--robabbility bounds for nonconvex optimization, we focus on IW-ERM with K =
Assumption 5.5 (Sub-Gaussian Noise). 1) A minimizer w exists; 2) The stochastic gradients
g(w)= = V ((ww) ) is unbiased, i.e. Elg( w)) = vw for any w E W; 3) The noise ([g(w) V (l ) 12 --sub-Gaussian (Vershynin 2018).
Theorem 5.5 (High-probabilltty Bound for Nonconvex Optimization). Let 8 E (0, and T E z4. Let K = and maxyey sup, Ty(y) = Tmax- Under Assumption 5.5 and 3-smoothness of nonconvex
f o hw suppose IW-ERM is run for T iterations with step-size min with probability 1 - 6, gradient norm squareds satisfy:
Then
F o
1 T
Σ
IV (h )12 =13=0(av
m
o² 2logg(168
T
T
Proof. We note that importance ratios do not depend on the model parameters w and the Lipschitz and smoothness constants for E fww w.r.t. w are scaled by mmax- The rest of the proof follows the arguments of (Liu et al. 2023 Theorem 4.1).
Theorem 5.5 shows that when the stochastic gradients are too noisy o Q log(1/0)) such that the second term in the rate dominates, then importance weighting does not have any negative impact on the convergence rate.
6 EXPERIMENTS
The experiments are divided into two main parts: evaluating VRLS's performance on a single node focusing on intra-node label shifts, and extending it to multi-node distributed learning scenarios with 5, 100, and 200 nodes. In the multi-node cases, we account for both inter-node and intra-node label shifts. Further experimental details, results, and discussions are provided in Appendix
Importance ratio estimation. We begin by evaluating VRLS on the MNIST (LeCun al. 1998) and CIFAR-10 (Krizhevsky) datasets in a single-node setting. Following the common experimental setup in the literature (Lipton et al., 2018), we simulate the test dataset using a Dirichlet distribution with varying a parameters. In this context, higher value indicates smoother transitions in the label distribution, while lower values reflect more abrupt shifts. The training dataset uniformly distributed across all classes. Initially, using sample size of 5,000, we investigate 20 a values within the range [10-1 1 J. Next, we fix a at either 1.0 or 0.1 and explore 50 different sample sizes ranging from 200 to 10,000. For each experiment,
run 100 trials and compute the mean squared error (MSE) between the true ratios and the estimated ratios. A two-layer MLP is used for MNIST, while ResNet-18 (He al. 2016) is applied for CIFAR-10.
Table 2: We utilize LeNet on Fashion MNIST to address label shifts across 5 nodes. For the baseline Feerodddd-eeAvg... FedBN, FedProx, and SCAF- FOLD-We run 15,000 iterations, while both the Upper Bound (IW-ERM with true ratios) and our IW-ERM with VRLS are limited 5,000 iterations. Notably, we employ a simple MLP with dropout for training the predictor. The model labeled Our IW-ERM (small) refers to our approach where the black-bor predictor is trained using only 10% of the available training data, balancing computa- tional efficiency with competitive performance.
Method
Avg.accuracy
Our IW-ERM
0.7520 + 0.0209
Our IW-EI EW-ERM (small)
0.7376 + 0.0099
FedAvg
0.5472 + 0.0297
FedBN
0.5359 + 0.0306
FedProx
0.5606 = 0.0070
SCAFFOLD
0.5774 0.0036
Upper Bound
0.8273 + 0.0041
8

---

Published as a conference paper at ICLR 2025
ca10
0015
0200
09925-
5015
0.00
PULL
0.75
2 3,25
8 10139
60335
0,20
1m
60310
D815
0015
30"
B paramaze
Shi
MNIST
CIFAR-10
CIFAR-10, 10, Relaxed CIFAR-10, Relax-m
Figure I: MSE analysis across different datasets and settings for VRLS (ours) compared to basclines, focusing on shift parameter (a) experiments. These subfigures include results from MNIST, CIFAR- 10, and relaxed label shift, illustrating the consistent superiority of VRLS. In the 'relaxed' setting, Gaussian blur (kernel size: 3; a: 0.1-0.5) and brightness adjustment (factor: =0.1) are applied with 30% probability to introduce real-world variability. In the 'relax-m' scenario, augmentations are applied with a 50% probability, with Gaussian blur (a: 0.1-0.7) and brightness (factor: +0.2).
VALS
VRLS
10
6-177
VRLS. EM
VILS 12
MLLS beselie
haseine
2000 1000 6in 1002 18000
7000 402D B0O0 R000 10503
2000 1005 6203
m950 19500
2000 4007 nim 19i0 10K00
Sumplesue
Sample 202
Sample
Sample sids
MNIST
CIFAR-10
CIFAR-10 10, Relaxed CIFAR-10 10, Relax-m
Figure 2: MSE analysis across different datasets and settings for VRLS (ours) compared to baselines, focusing on sample size experiments. These subfigures include results from MNIST, CIFAR- 10, and relaxed label shift conditions, highlighting VRLS's 's superior performance across varying test set sizes.
Table 3:
We deploy R8 on CIFAR-10 to address label shifts across nodes. The predictor is also a ResNet-18. 18, ensuring consistency with the single-node scenario. For a fair comparison, we limit IW-ERM with VRLS and the true ratios to 5,000 iterations, while FedAvg and FedBN are run for 10,000 iterations. Detailed results are provided in Table 7
CIFAR-10
Our IW-ERM
FedAvg
FedBN
Upper Bound
Avg. accuracy
0.5640 + 0.0241 0.4515 + 0.0148 0.4263 + 0.0975 0.5790 + 0.0103
Table 4: We present the average node accuracies from the CIFAR-10 target shift experiment conducted with 100 and 200 nodes, where nodes are randomly sampled to participate in each training round. Our IW-ERM with VRLS is run for 5,000 and 10,000 iterations, respectively, while both FedAvg and FedBN are run for 10,000 iterations each.
CIFAR-10
Our IW-ERM
FedAvg
FedBN
Avg. accuracy (100 nodes)
0.5354
0.3915
0.1537
Avg. accuracy (200 nodes)
0.6216
0.5942
0.1753
Figure I and Figure 2 compares our proposed VRLS with baselines (Garg al. 2020 Saerens et al.. 2002) under label shifts. MLLSSLD refers to the MLLS method using convex optimization via SGD (Garg et 2020), while MLLS EM employs the same objective function but is optimized using the EM algorithm (Sacrens et al. 2002). Our proposed VRLS is optimized in similar manner, resulting VRLSS 1.2 and VRLS EMM as shown in the figure. Our method consistently achieves lower MSE across different label shift intensities (a) and test sample sizes on both datasets. Notably, our importance ratio estimation experiments align with the error bound Theorem 5.1 demonstrating that increasing the number of test samples improves estimation error at rate proportional to the square root of the sample size. Additionally, the regularization term constrains the parameter space and reduces Rademacher complexity, leading to smoother predictions and improved model calibration, as supported by Section S2 (Guo et al. 2017a). Both of them contribute to reduced estimation error.

---

Published as a conference paper at ICLR 2025
We also tested importance ratio estimation under relaxed label shift conditions and found VRLS to exhibit greater robustness (see Appendix [.2] for detailed settings). Although this assumption holds broader potential for real-world upplications, its precise alignment with real-world datasets requires further investigation-22 an important direction for future research that extends beyond the scope of this work.
Distributed learning settings. We apply VRLS in distributed learning context, addressing both intra- and inter-node label shifts. The initial experiments involve 5 nodes, using predefined label distributions Fashion MNIST (Xiao et al. 2017) and CIFAR-IO, as shown in Tables g in Appendix !
We employ a simple MLP with dropout the predictor for Fashion MNIST. For global training with IW-ERM, LeNet (LeCun et al. 1998) is used on Fashion MNIST, and ResNet- 18 (Ramezani
Kebrya et al. 2023b) on CIFAR-TO. All experiments are run with three random seeds, reporting the average accuracy across nodes. We compare IW-ERM with VRLS against baseline methods,
including FedAvg (McMahan et al. 2017). FedBN (Li et al. 202Ib), FedProx (Li ct al. (020), and SCAFFOLD (Karimireddy et al., 2020a), as well as IW-ERM with true ratios serving as upper bound. Hyperparameters are kept consistent with those in (McMahan et al. 2017 et al. 20216
Ramezani-Kebrya et al. (2023b).
Each node's 's stochastic gradients are computed with a batch size of 64 and aggregated using the Adam optimizer. All experiments are on single GPU within internal cluster. Both MILS and VRLS use identical hyperparameters and training epochs for CIFAR-10 and Fashion MNIST, stopping once the classification loss reaches a predefined threshold 100 and 200 nodes on CIFAR-I0, 10, where five nodes are randomly sampled each iteration to simulate more realistic distributed learning. In this case, IW-ERM with true ratios does not act as the upper bound due to the stochastic node sampling. The experiment is run once, and average accuracy across nodes reported, with label distribution shown in Table 10 in Appendix T Despite FedBN's 's reported et al. 2023b), we maintain 15,000 and 10,000 iterations for
MNIST. We also conduct experiments with
slow convergence FedAvg and FedBN on Fashion MNIST and CIFAR-10, respectively, for fair comparison. However, IW-ERM is limited to 5,000 iterations using both true and estimated ratios due faster convergence.
As shown in Table 2 IW-ERM achieves over 20% higher average accuracy than all baselines on Fashion MNIST, with only third of the iterations. Notably, even with just 10% of the training data in the first round of global training, the performance remains comparable, demonstrating reduced training complexity. This improvement is attributed to the theoretical benefits of IW-ERM. the robustness of density estimation, and the fact that the aggregation of importance ratios reduces reliance on any single local estimate. Similarly, bound on CIFAR-10, outperforming the baselines. Individual node accuracies are detailed in Tables 7in Appendix In the 100-node scenario, IW-ERM continues to demonstrate superior performance. requiring only half the iterations, as shown in Table 4 It is important to note that using true ratios does not equate IW-ERM given the stochasticity of node selection during training.
that IW-ERM approaches the upper
7 CONCLUSIONS AND LIMITATIONS
We propose VRLS to address label shift in distributed learning. Paired with IW-ERM, VRLS improves intra- and inter-node label shifts multi-node settings. Empirically, VRL. consistently outperforms ML.LS-based baselines, and IW-ERM with VRLS exceeds all multi-node learning baselines. Theoretical bounds further strengthen our method's 's foundation. Future work will explore estimating ratios by relaxing the strict class-conditional assumption and optimizing IW-ERM to reduce time complexity while ensuring scalability and practicality in real-world distributed learning.
ETHICS STATEMENT
No ethical approval was needed as no human subjects were involved. All authors fully support the content and findings.
10

---

Published as a conference paper at ICLR 2025
REPRODUCIBILITY STATEMENT
We ensured reproducibility with publicly available datasets (MNIST. CIFAR-I0, and standard models (e.g., ResNet-18). Links to datasets, code, and configurations will be provided upon camera-ready submission. Experiments were run on NVIDIA 3090, A100 GPUs, and Google Colab, with average results and variances reported across multiple trials.
ACKNOWLEDGMENTS
The authors would like to thank Leello Tadesse Dadi and Thomas Pethick for helpful discussions. This work was supported by the Research Council of Norway (RCN) through its Centres Excellence scheme, Integreat: Norwegian Centre for knowledge-driven machine learning, project number 332645. The work of Changkyu Choi was funded by RCN under grant 309439. The work of Volkan Cevher was supported by Hasler Foundation Program: Hasler Responsible AI (project number 21043), the Army Research Office which was accomplished under Grant Number W91 W-F-24-1-0045 and the Swiss National Science Foundation (SNSF) under grant number 200021 ,220001 1. The computations were performed in part on resources provided by Sigma2 the National Infrastructure for High- Performance Computing and Data Storage in Norway.
REFERENCES
Rohit Agrawal and Thibaut Horel. Optimal bounds between f-divergences and integral probability metrics. In International Conference on Machine Learning (ICML), 2020.
Amr M. Alexandari, Anshul Kundaje, and Avanti Shrikumar. Maximum likelihood with bias-
corrected calibration is hard-to-beat at label shift adaptation. In International Conference Machine Learning (ICML), 2020.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International Conference on Machine Learning (ICML), 2017.
Kamyar Azizzadenesheli, Anqi Liu, Fanny Yang, and Animashree Anandkumar. Regularized learning for domain adaptation under label shifts. In International Conference on Learning Representations (ICLR), 2019.
Arindam Banerjee, Srujana Merugu, Inderjit S. Dhillon, and Joydeep Ghosh. Clustering with bregman divergences. Journal of Machine Learning Research (JMLR), 6(58): 2005.
Jeremiah Birrell, Paul Dupuis, Markos Katsoulakis, Yannis Pantazis, and Luc Bey-Bellet. (f. Divergences: Interpolating between /divergenncss and integral probability metrics. Journal of Machine Learning Research (JMLR), 23:1-70, 2022.
Jonathon Byrd and Zachary C. Lipton. What is the effect of importance weighting in deep learning? In International Conference on Machine Learning (ICML), 2019.
Artur Back Luca, Guojun Zhang, Xi Chen, and Yaoliang Yu. Mitigating data heterogeneity in federated learning with data augmentation. arXiv preprint arXiv:2206.099 09979, 2022.
Jérôme Dedecker, Clémentine Prieur, and Paul Raynaud de Fitte. Parametrized Kantorovich- Rubinštein theorem and application to the coupling of random variables. Springer, 2006.
Fartash Faghri, Iman Tabrizian, Ilia Markov, Dan Alistarh, Daniel M. Roy, and Ali Ramezani- Kebrya. Adaptive gradient quantization for data-parallel SGD. In Advances in neural information processing systems (NeurIPS), 2020.
Tongtong Fang, Nan Lu, Gang Niu, and Masashi Sugiyama. Rethinking importance weighting for deep learning under distribution shift. Advances in neural information processing systems (NeurIPS), 33:1 1991-19990007 2020.
Tongtong Fang. Nan Lu, Gang Niu, and Masashi Sugiyama. Generalizing importance weighting to a universal solver for distribution shift problems. Advances in neural information processing systems (NeurIPS), 36, 2023.
11

---

Published as a conference paper at ICLR 2025
Saurabh Garg, Yifan Wu, Sivaraman Balakrishnan, and Zachary C. Lipton. A unified view of label shift estimation. In Advances in neural information processing systems (NeurIPS), 2020.
Saurabh Garg, Sivaraman Balakrishnan, and Zachary C. Lipton. Domain adaptation under open set label shift. arXiv preprint arXiv:2207.13044 2022.
Saurabh Garg, Nick Erickson, James Sharpnack, Alexander J. Smola, Sivaraman Balakrishnan, and Zachary C. Lipton. Risbench: Domain adaptation under relaxed label shift. International Conference Machine Learning (ICML), 2023.
Margalit R. Glasgow, Honglin Yuan, and Tengyu Ma. Sharp bounds for federated averaging (local SGD) and continuous perspective. In International Conference Artificial Intelligence and Statistics (AISTATS), 2022.
Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel tao-s test. Journal of Machine Learning Research, 13(25):723-773, 2012.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q. Weinberger. On calibration of modern neural networks. Proceedings of the 34th International Conference on Machine Learning Volume 70, ICML'17, 17, pp. 1321-1330. JMLR.org. 2017a.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q. Weinberger. On calibration of modern neural networks. In Proceedings of the 34th International Conference on Machine Learning Volume 70, ICML'17, pp. 1321-1330. JMLR. 2017b.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q. Weinberger. On calibration of modern neural networks. Proceedings of the 34th International Conference on Machine Learning Volume 70, ICML'17, pp. 1321-1330. JMLR.org. 2017c
Jiaxian Guo, Mingming Gong, Tongliang Liu, Kun Zhang, and Dacheng Tao. LTF: A label transfor- mation framework for correcting label shift. In International Conference Machine Learning (ICML), 2020.
Sharut Gupta, Kartik Ahuja, Mohammad Havaei, Niladri Chatterjee, and Yoshua Bengio. Fl games: A federated learning framework for distribution shifts. arXiv preprint arXiv:2205.11100 2022.
Farzin Haddadpour, Mohammad Mahdi Kamani, Aryan Mokhtari, and Mehrdad Mahdavi. Federated learning with compression: Unified analysis and sharp guarantees. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2021.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Conference on Computer Vision and Pattern Recognition (CVPR), 2016.
Zhengmian Hu and Heng Huang. Tighter analysis for ProxSkip. International Conference on Machine Learning (ICML), 2023.
Jiayuan Huang, Arthur Gretton, Karsten Borgwardt, Bernhard Schölkopf, and Alex Smola. Correcting sample selection bias by unlabeled data. In Advances in neural information processing systems (NeurIPS), 2006.
Yutao Huang, Lingyang Chu, Zirui Zhou, Lanjun Wang, Jiangchuan Liu, Jian Pei, and Yong Zhang Personalized cross-silo federated learning on non-iid data. In AAAI Conference on Artificial Intelligence, 2021.
Kairouz, H. McMahan, B. Avent, Bellet, M. Bennis, N. Bhagoji, K. Bonawitz, Z. Charles, G. Cormode, R. Cummings. R. G. L. D'Oliveira, S. E. Rouayheb, D. Evans, J. Gardner, Z. Garrett, A. Gascón, B. Ghazi, B. Gibbons, M. Gruteser, Harchaoui, C. He, L. He, Z. Huo, B. Hutchin- son, J. Hsu, M. Jaggi, T. Javidi, G. Joshi, M. Khodak, J. Konečný, A. Korolova, F. Koushanfar, S. Koyejo, T. Lepoint, Liu, P., M. Mohri, R. Nock, A. Özgür, R. Pagh, M. Raykova, H. Qi, D. Ra- mage. R. Raskar, D. Song, W. Song, S. U. Stich, Z. Sun, A. T. Suresh, F. Tramèr, P. Vepakomma, Wang. L. Xiong, Z. Xu, Q. Yang, X. Yu, H. Yu, and S. Zhao. Advances and open problems in federated learning. Foundations and x in Machine Learning, 14(1-2):1-21 2021.
12

---

Published as a conference paper at ICLR 2025
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. SCAFFOLD: Stochastic controlled averaging for federated learning. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 1 119 of Proceedings of Machine Learning Research, pp. 2p.51132 5132-5143. PMLR, 13-18 Jul 2020a.
Sai Pranceth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. SCAFFOLD: Stochastic controlled averaging for federated learning. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 5132-5143. PMLR, 13-18 Jul 2020b.
Mikhail Khodak, Maria-Florine Balcan, and Ameet Talwalkar. Adaptive gradient-based meta-learning methods. In Advances in neural information processing systems (NeurIPS), 2019.
Jinkyu Kim, Gecho Kim, and Bohyung Han. Multi-level branched regularization for federated learning. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song. Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference Machine Learning. volume 162 of Proceedings of Machine Learning Research, pp. 11058- 11033. PMLR, 17-23 Jul 2022. URL https: httpsss:///oceeddiiee eggs ngs mlr
wtt.:ssmw.....rs...mmmm www.sp.essslll.. 62 2tt/2/kimm2.... htmmmmmmmmmm
A. Krizhevsky. Learning multiple layers features from tiny images. Technical report, University of Toronto, 2009.
Meelis Kull, Miquel Perello-Nieto, Markus Kängsepp, Telmo Silva Filho, Hao Song, and Peter Flach. Beyond temperature scaling: obtaining well-calibrated multiclass probabilities with Dirichlet calibration. Curran Associates Inc., Red Hook, NY, USA, 2019.
S. Kullback and R. A. Leibler. On Information and Sufficiency. The Annals of Mathematical Statistics, 22(1):79- - 86, 1951.
Gil Kur, Eli Putterman, and Alexnader Rakhlin. On the variance, admissibility, and stability of empirical risk minimization. In Proceedings of the 37th International Conference on Neural Information Processing Systems, NIPS '23, Red Hook, NY, USA, 2024. Curran Associates Inc.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Bradient-based learning applied to document recognition. Proceedings of the IEEE, 86:2278 -3244, 1998.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. In Inderjit S. Dhillon, Dimitris S. Papailiopou- los, and Vivienne Sze (eds.). Proceedings of Machine Learning and Systems 2020, MLSys 2020, Austin, TX. USA, March 2-4, 2020. mlsys.org, 2020.
Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and robust federated learning through personalization. In International Conference on Machine Learning (ICML), 2021a.
Xiaoxiao Li, Meirui Jiang, Xiaofei Zhang, Michael Kamp, and Qi Dou. FedBN; Federated learning on non-lID features via local batch normalization. In International Conference on Learning Representations (ICLR), 2021b.
Zachary C. Lipton, Yu-Xiang Wang, and Alexander Smola. Detecting and correcting for label shift with black box predictors. In International Conference on Machine Learning (ICML), 2018.
Zijian Liu, Ta Duy Nguyen, Thien Hang Nguyen, Alina Ene, and Huy Nguyen. High probability convergence of stochastic gradient methods. In International Conference on Machine Learning (ICML), 2023.
Kangyang Luo, Xiang Li, Yunshi Lan, and Ming Gao. Gradma: A gradient-memory-bue ac- celerated federated learning with alleviated catastrophic forgetting. In 2023 IEEE/CVF Con- ference on Computer Vision and Pattern Recognition (CVPR), pp. 3708-3717, 2023. doi: 10.1 10.1109CVVPS5779.220.00
13

---

Published as a conference paper at ICLR 2025
You-We Wei Luo and Chuan-Xian Ren. Generalized label shift correction via minimum uncer- tainty principle: Theory and algorithm. ArXiv, abs/2202.13043, 2022. URL https://api
semanticscholaxe wwtp:chhww.....ooooooooueeeeeeee http:/opuuuuo:::1577777
Pranav Mani, Manley Roberts, Saurabh Garg, and Zachary C. Lipton. Unsupervised learning under latent label shift. In ICML Workshop on Spurious Correlations, Invariance and Stability, 2022.
H. Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera Y. Arcas. -efficic learning of deep networks from decentralized data. In International
Communication- efficient Conference on Artificial Intelligence and Statistics (AISTATS), 2017.
Dexter Neo, Stefan Winkler, and Tsuhan Chen. Maxent loss: Constrained maximum entropy for calibration under out-of-diitributiin shift. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 21463-21472 2024.
Gabriel Pereyra, George Tucker, Jan Chorowski, Eukasz Kaiser, and Geoffrey Hinton. Regularizing neural networks by penalizing confident output distributions. arXiv preprint 2017.
1065488
Stephan Rabanser, Stephan Günnemann, and Zachary Lipton. Failing Loudly: An Empirical Study of Methods for Detecting Dataset Shift. Curran Associates Inc., Red Hook, NY, USA, 2019.
Anichur Rahman, Md Sazzad Hossain, Ghulam Muhammad, Dipanjali Kundu, Tanoy Debnath, Muaz Rahman, Md Saikat Islam Khan, Prayag Tiwari, and Shahab Band. Federated learning-based ai approaches in smart healthcare: concepts, taxonomies, challenges and open issues. Cluster computing, 26(4):2271-233 2311, 2023.
Suraj Rajendran, Zhenxing Xu, Weishen Pan, Arnab Ghosh, and Fei Wang. Data heterogeneity in federated learning with electronic health records: Case studies of risk prediction for acute kidney injury and sepsis diseases in critical care. PLOS Digital Health, 2(3):e00001 17,000111110010 2023.
Ali Ramezani- Bimnezani-Kebyya,, Fartash Faghri, Ilya Markov, Vitalii Aksenov, Dan Alistarh, and Daniel M. Roy. NUQSGD: Provably communication-effone effficient data tatta-arallll SGD via nonuniform quantization. Journal of Machine Learning Research (JMLR), 22(1 14):1-43, 2021.
Ali Ramezani-Kebrya, Kimon Antonakopoulos, Igor Krawezuk, Justin Deschenaux, and Volkan Cevher. Distributed extra-gradient with optimal complexity and communication guarantees. In International Conference Learning Representations (ICLR), 2023a.
Ali Ramezani- Kebrrya, Fanghui Liu, Thomas Pethick, Grigorios Chrysos, and Volkan Cevher. Fed- erated learning under covariate shifts with generalization guarantees. Transactions on Machine Learning Research (TMLR), 2023b.
Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do CIFAR-10 classifiers generalize to CIFAR-10? arXiv preprint / 201 8.
Marco Saerens, Patrice Latinne, and Christine Decaestecker. Adjusting the outputs of a classifier to new a priori probabilities: a simple procedure. In Neural Computation, 2002.
Hidetoshi Shimodaira. Improving predictive inference under covariate shift by weighting the log- likelihood function. Journal of Statistical Planning and Inference, 90(2):227-244, 2000.
Avanti Shrikumar, Amr M. Alexandari, and Anshul Kundaje. Adapting to label shift with bias- corrected calibration. arXiv preprint arXiv: 10..1190...6..55... 2019.
Virginia Smith, Chao-Kai Chiang, Maziar Sanjabi, and Ameet S. Talwalkar. Federated multi-task learning. Advances in neural information processing systems (NeurIPS), 2017.
Masashi Sugiyama, Benjamin Blankertz, Matthias Krauledat, Guido Dornhege, and Klbus-Robert Müller.
cros--alidatiion for covariate shift. In Joint Pattern Recognition Symposim, pp. 354-363. Springer, 2006.
Masashi Sugiyama, Matthias Krauledat, and Klaus-Robert Müller. Covariate shift adaptation by importance weighted cross validation. Journal of Machine Learning Research (JMLR), 8(5), 2007.
14

---

Published as a conference paper at ICLR 2025
Zeyu Sun, Dogyoon Song, and Alfred Hero. Minimum-risk recalibration of classifies. Advances in neural information processing systems (NeurIPS), 36, 2024.
Antonio Torralba, Rob Fergus, and William T. Freeman. 80 million tiny images: A large data set for nonparametric object and scene recognition. IEEE Transactions Pattern Analysis and Machine Intelligence, 30(11): 1958--9999 2008.
Roman Vershynin.
High-dimensional probability: An introduction with applications in data science. Cambridge university press, 2018.
Cédric Villani. The Wasserstein distances. Springer Berlin Heidelberg, 2009.
Deng-Bao Wang, Lei Feng, and Min-Ling Zhang. Rethinking calibration of deep neural networks: Do not be afraid of overconfidence. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in neural information processing systems (NeurIPS), volume 34, pp. 11809-11820 Curran Associates, Inc., 2021.
Kaibin Wang, Qiang He, Feifei Chen, Chunyang Chen, Faliang Huang, Hai Jin, and Yun Yang. Flexifed: Personalized federated learning for edge clients with heterogeneous model architectures. In Proceedings of the ACM Web Conference 2023, pp. 2979-2990, 2023.
Jie Wen, Zhixia Zhang. Yang Lan, Zhihua Cui, Jianghui Cai, and Wensheng Zhang. A survey on federated learning: challenges and applications. International Journal of Machine Learning and Cybernetics, 14(2):513-55) 2023.
Blake Woodworth, Kumar Kshitij Patel, Sebastian Stich, Zhen Dai, Brian Bullins, Brendan Mcmahan, Ohad Shamir, and Nathan Srebro. local SGD better than minibatch SGD? In International Conference on Machine Learning (ICML), 2020.
Xidong Wu, Feihu Huang, Zhengmian Hu, and Heng Huang, Faster adaptive federated learning, In AAAI Conference on Artificial Intelligence, 2023.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-MNIST: novel image dataset for bench- marking machine learning algorithms. arXiv preprint arXiv: 170.:0::000007777, 2017.
Mang Ye, Xiuwen Fang, Bo Du, Pong C Yuen, and Dacheng Tao. Heterogeneous federated learning: Statc-of-the-an and research challenges. ACM Computing Surveys, 56(3):1-44, 2023.
Mingzhang Yin, Yixin Wang, and David M Blei. Optimization-basel causal estimation from
heterogeneous environments. Journal of Machine Learning Research, 25:1-44. 2024.
Bianca Zadrozny. Learning and evaluating classifiers under sample selection bias. In International Conference Machine Learning (ICML), 2004.
Helen Zhou, Sivaraman Balakrishnan, and Zachary Lipton. Domain adaptation under missingness shift. Artificial Intelligence and Statistics (AISTATS), 2023.
15

---

Published as a conference paper at ICLR 2025
The Appendix part is organized as follows:
All related work are provided in Appendix A
Additional details of prior work of BBSE and MLLS are in Appendix B
Mathematical proof for label shifts with multiple nodes and IW-ERM given in Appendix C
General algorithmic description is in Appendix D
Proof of Theorem 5.1 is in Appendix
Proof of Theorem 5.2 and
Privacy guarantees for IW-ERM
in Equation (WWEERM) are provided in Appendix F
Complexity analysis is in Appendix G
Mathematical notations are summarized in Appendix H
Limitations are discussed in Appendix |
Additional experiments and experimental details are provided in Appendix
16

---

Published as a conference paper at ICLR 2025
A RELATED WORK
In the context of distributed learning with label shifts, importance ratio estimation is tackled either by solving linear system as in (Lipton et al. 2018 Azizzadenesheli et 2019) by minimizing distribution divergence as in (Garg et al. 2020). In this section, we overview complete related work.
Federated learning (FL.). Much of the current research in FL. predominantly centers around the minimization of empirical risk, operating under the assumption that each node maintains the same training/test data distribution (Kairouz et al. (02T). Prominent methods in FL include FedAvg (McMa-
han et al., 2017), FedBN (Li et | 2021b), FedProx (Li et al., (020) and SCAFFOLD (Karimireddy et al. 2020a). FedAvg and its variants such as (Huang et al. [021]] Karimireddy et al. 2020b)
have been the subject of thorough investigation in optimization literature, exploring facets such as communication efficiency, node participation, and privacy assurance (Faghri et al. 2020 Ramezani Kebrya et al.) 202T 2023a b). Subsequent work, such as the study by by Luca al. (2022). D explores
Federated Domain Generalization and introduces data augmentation to the training. This model aims to generalize to both in-domain datasets from participating nodes and an out-of-domain dataset from a non-participating node. Additionally, Gupta et al. (002)) introduces FL Games, game-theoretic framework designed learn causal features that remain invariant across nodes. This is achieved by employing ensembles over nodes' historical actions and enhancing local computation, under the assumption of consistent training/test data distribution across nodes. The existing strategies to address statistical heterogeneity across nodes during training primarily rely on heuristic- based personalization methods, which currently lack theoretical backing in statistical learning (Smith et al. 2017 Khodak al. 2019 Li et al. 202 In contrast, aim to minimize overall test error amid both intra-node and inter-node distribution shifts, a situation frequently observed in real-world scenarios. Techniques ensuring communication efficiency, robustness, and secure aggregations serve as complementary.
Importance ratio estimation Classical Empirical Risk Minimization (ERM) seeks to minimize the expected loss over the training distribution using finite samples. When faced with distribution shifts, the goal shifts to minimizing the expected loss over the target distribution, leading to the development of Importance- Weighted Empirical Risk Minimization ([V-ERMy(Shimoddnn 2000 Sugiyama al. 2006; Byrd & C. Lipton 2019 Fang et al. 2020 Shimodaira (2000) established that the IW-ERM estimator is asymptotically unbiased. Moreover, Ramezani-Kelll et al. (002bb)
introduced FTW-ERM, which integrates importance ratio estimation.
Label shift and MLLS f'amily For theoretical analysis, the conditional distribution (2) y) is held strictly constant across all distributions (Lipton al., 2018 Garg et al. 20200] Saerens et al.] 2002). Both BBSE (Lipton et al. 2018) and RLLS (Azizzadeneshel et al. (2019) designate a discrete latent space z and introduce a confusion matrix based estimation method to compute the ratio ur by solving a linear system (Saerens et al., 2002; Lipton et al.) 2018)). This approach is straightforward and has been proven consistent, even when the predictor is not calibrated. However, its subpar performance is attributed to the information loss inherent in the confusion matrix (Garg et al. 2220)).
Consequently, MLLS (Garg et al. 2020)) introduces a continuous latent space, resulting in a significant enhancement in estimation performance, especially when combined with a post-hoc calibration method (Shrikumar et al., 2019). It also provides a consistency guarantee with a canonically calibrated predictor. This EM-based MLLS method is both concave and can be solved efficiently.
Discrepancy Measure In information theory and statistics, discrepancy measures play a critical role in quantifying the differences between probability distributions. One such measure is the Bregman Divergence (Banerjee et al. 2005), defined as
( (g)) = <(x) - p(y) (V)(y), x- y).
which encapsulates the difference between the value of a convex function ø at two points and the value of the linear approximation of at one point, leveraging the gradient another point.
Discrepancy measures are generally categorized into two main families: Integral Probability Metrics (IPMs) and f-divergences. IPMs, including Maximum Mean Discrepancy (Gretton et al,, 2012) and Wasserstein distance (Villani 2209),, focus on distribution differences p- Q. In contrast, f divergences, such as KI KL--iivereence (Kullback & Leibler 1951) and Total Variation distance, operate
17

---

Published as a conference paper at ICLR 2025
ratios P/Q and do not satisfy the triangular inequality. Interconnections and variations between these families are explored in studies like (f. Diviveergencss (Birrell et al. 2022). which interpolate between f-divergences and IPMs, and research outlining optimal bounds between them (Agrawal & Horel 2020).
MLLS (Garg et al. 2020) employs f-divergence, notably the KL divergence, which is not metric as doesn't i satisty the triangular inequality, and requires distribution to be absolutely continuous with respect Q. Concerning IPMs, while MMD is reliant on a kernel function, it can suffer from the curse of dimensionality when faced with high-dimensional data. On the other hand, the Wasserstein distance can be reformulated using Kantorovich-Rubinsteinn duality (Dedecker et al. 2006 Arjovsky et al. 2017) as a maximization problem subject to Lipschitz constrained function R 7 R.
18

---

Published as a conference paper at ICLR 2025
B BBSE AND MLLS FAMILY
In this section, we summarize the contributions of BBSE Lipton et al. 2018) and MLLS (Garg et al. 2020). Our objective to estimate the ratio pd (v)/p"(u) (y). We consider scenario with possible label classes, where y = c for c E [m] Let r [r]. mm] represent the true ratios, with each GGarg al. 2020). We then define family of distributions over Z, parameterized by [ 7'
(y==) r2 defined r- p (ymo)
R where Te is the e-th element of the ratio vector.
m
Pr(z)
p (z p c) -r
(10)
Here, Te for [m] and E re p= c) = Σm (f 1 constraints. When
r = r", e.g., Te = rg for e E [m], we have (2) = Pr (2) = (2) (Garg et al. 2020). So our task
find such that
p (2ly = c) pp" = c) rex
c=1
(11)
m
Σ pd (z, c)-re = (2)
Lipton et al. (2018) introduced Black Box Shift Estimation (BBSE) to address this issue. With pre- pre-trained classifier f for the classification task, BBSE assumes that the latent space Z is discrete and defines p(=[x) = Sarg max f(m). where the output of f(x) is probability vector (or simplex) over m classes. BBSE estimates (zly) as a confusion matrix, using both the training and validation data. It calculates p" (y c) from the training set and pl (z) from the test data, The problem then reduces to solving the following equation:
Aw = B
(12)
where |2] = m, A E Rmxm with Ajo = (= = jjy= c) pp c). and B E R" with B, (2 = 5) for e, [m].
The estimation of the confusion matrix in terms of (² leads to the loss of calibration information (Garg et al. 2020). Furthermore, when defining Z continuous latent space, the confusion matrix becomes intractable since = has infinitely many values. Therefore, MLLS directly minimizes the divergence between (2) and (=(2), instead of solving the linear system in Equation (12).
Within the f-divergence family, MLLS seeks to find a weight vector r by minimizing the KL- divergence DKL (2)(2), Pr)))) Ege [log Dd (z)/pr(z)], for (z( defined in Equation (10). Lever- aging on the properties of the logarithm, this is equivalent to maximizing the log-likelihood: = arg max,ER Eme [log (=)). Expanding (=),), have
Eye [z)])(2)) Ere
log
(= =)ro)
(13)
Er
log(
("( = z)rc) + (z)
Therefore the unified form of MLLS can be formulated as:
arg max Ene log Σ p- 2)ro) PER
(14)
This is a convex optimization problem and can be solved efficiently using methods such as EM, an analytic approach, and also iterative optimization methods like gradient descent with labeled training data and unlabeled test data. MLLS defines the (z) as ò plugs in the pre-defined f to approximate yd (y((a)) and optimizes the following objective:
19

---

Published as a conference paper at ICLR 2025
ry == arg max f(r, := arg max Er [arg([(x)")) [r)]
(15)
PER
PER
With the Cias-Corrected Calibration (BCT) (hrrikuaar al. 2019) strategy, they adjust the logits f() of f()) element-wise for each class, and the objective becomes:
Ty arg max E(r, = arg max En [ooglg f(z))r rER
(16)
PER
where g is calibration function.
20

---

Published as a conference paper at ICLR 2025
C PROOF OF PROPOSITION 4.1
In the following, consider four typical scenarios under various distribution shifts described in and formulate their IW-ERM with focus on minimizing Ra.
C.I No INTRA-NODE LABEL SHIFT
For simplicity, we assume that there are only nodes, our results can be extended to multiple
nodes. This scenario assumes p)(y) (z) for l = 1, 2, but p * (3)) Node aims to learn
fww assuming (v) is given. We consider the following IW-ERM that is consistent in minimizing R1;
pv(v)
min Rww EH ni
(ho (x v..))
(17)
(v
(h (x. v1,)).
12
(1
Here H. is the hypothesis class of hww This scenario is referred to as No-LS.
C.2 LABEL SHIFT ONLY FOR NODE
Here we consider label shift only for node 1, i.e., (T())) + p3)) and po(y) =p% We consider the following IW-ERM:
(VT))) (B)l.))
min h
E(h (a vi,s)
(18)
n
pl (v
E( (2/,,) v..))
n
p5
This scenario is referred to as LS on single.
C.3 LABEL SHIFT FOR BOTH NODES
Here we assume (y) + (v) and (1) + ( i.e., label shift for both nodes. The cor- responding IW-ERM the both.
Eq. Equation (88) This scenario is referred LS on
Without loss of generality and for simplicity, set / 1. various distribution shifts and formulate their IW-ERM with focus on minimizing Rj. The details of these scenarios are summarized Table
We consider four typical scenarios under
C.4 MULTIPLE NODES
Here we consider general scenario with K nodes. We assume both intra-node and inter-node label shifts by the following IW-ERM:
Pl (1
min Few
(x vK,))
(19)
p (
This scenario is referred to as LS on multi.
For the scenario without intra-node label shift, the IW-ERM in Equation 17
be expressed
21

---

Published as a conference paper at ICLR 2025
1
Pi
elo (22.)).
2
(((
(1( (z),
p5(()) Bp(nl))) [e
[C(h (2:), (y)dy)
(20)
pt (() [0( y)ldy
pf (U) (z),
(x (a),, vil
Rí (hwo)
where the second equality holds due to the assumption of the label shift setting and Bayes' theorem: p(x, y) p(xlly) p(y). and the fourth equality holds by the assumption that ( p (g) in the No-LS setting.
For the scenario with label shift only for Node for both nodes, the IW-ERM in Equation [8] admits
(Nhw (z02,)) y².i)
(21)
2
(3
mg-000 (x
p9(y)) E(h (x), 3)
(22)
(15(3)
1 (V)
Bp(crle) (ettyileeenn (e( (x)), (
(23)
pb)(y)
[. pf Bypeeel (2)) d)
(24)
(pp((n)) (((hew 4)l
(25)
Rí (hw.)
(26)
For multiple nodes, let ls E [K]. Similarly, we have
(a vi..) o -to Ri (hw).
(27)
k PK(())))
Then we have
K
ni.....k-- +00
Elhw (a%k (1 vh,i)
Ry (hw))
(28)
Note that to solve Equation
19 node I needs to estimate (((v) for all nodes & in Equation 19
Pf
The consistency of Equation arguments in e.g., (Shimodaira Seetone 3] and (Sugiyama et al. (007) 2.2] using the law of large numbers.
[W-ERM) i.e., convergence in probability, is followed the standard
22

---

Published as a conference paper at ICLR 2025
D ALGORITHMIC DESCRIPTION
2. Split the training dataset on each node
trainsets = target target_shift. spl.tttsplit..mm www.shttttttttatttteaeeeeeeeeeea saraattttt data, trainset. targets, node mode_Label Label dist_train, transforrtranstomm train)
5 Split the test dataset on each node
testsets - target_shift. split_datast (tesssst.ttttst testset. targets, Lode_label_dist_ee test, transfomsttrasfor_ttes
Initialize K local models (nets) for each node
9 nets = (Initialize_modell () for in range (node_num)
Initialize the estimator for each local model
estimators [1s_RatioModel (nets [k]) for in range (node_num))
13
14 Initialize tensors to store the estimated ratios, values, and marginal values for each pair of nodes.
estimated_rate ratios torch. zeros (node (ncde_num, num, node num, aclass) estimated vstina values torch. zeros (node, num, node num, nclass) marginal_values torch. zeros (node_num, nelass)
17
Phase 1: Compute the estimated ratios for each node pair (k, d) for in range (node_num)) for j range (node_num)
21
22
# Perform test on node k using node jts testset estimated_ration
23
jl estimators [k] (testsets data. ()
numpy ())
24
25
Compute the marginal values on each node's a training set for i, trainset in enumerate (trainsets): marginal_values[]]
Phase 2:
36
marginal (trainset. strainset.targets:
29 Phase 3:
Compute the final estimated values for each node 30 for k in range (node (node_num):
31 for j in range (node_num): : astimated_vallest [k, jl marginal_values [j] ostimated
jl
34 Aggregate the estimated values across nodes gggregated_valuse torch. sum (esttimmtted_vv peatted_values, dim-l)
39 Compute the #inal ratios for each node
38 ratios (aggregated_ values / pargins__valles to device)
Listing 1: Our VRLS in distributed learning. It is the implementation of Algorithm
23

---

Published as a conference paper at ICLR 2025
E PROOF OF THEOREM 5.1
Proof. Let (r, 0) x) = - Log fxx T))r) From the strong convexity in Lemma E.7
we have that
Pn 1 K
(Ca. (f La (rg))
(29)
pPmmn
Now focusing on the term on the right-hand side, we find by invoking Lemma E.4 that
Le (fnz)
Lo (ry.)
H H x2) B
H Hr H(rr,, x2)
+ 2LE
110 0-12
n"
x
H(fn=, ((+n=,,n=t,2)) 2))
H One, x))) Σ n
H (f pn xzg) Σ
r2
j=1
j=1
E
H Onr, x))
+ +LE
10 0-1 1|2
n""
< H(Ane, k))
Σ n
1
H(P æ);) +
Σ
H (r 0n,, 2.3))
n
j=1
==1
18
E
([rs One,
+ +LLE 10 0-1
(30)
where in the last inequality we used the fact that A, minimizer of r + ドー (r, 0z, 2z,). Finally by using Lemma E.5 and Lemma E.6 with 5/2 each, we have that with probability 1 8,
log(4/0)
La (Pa)) La- (ry.)=
<
Rad (F) + 2LE 11ām.. - 0-112 + VBV
(31)
n
nto
Plugging this back into Equation
29 have that
Am 1
2
4
Rad(F) + 4B
log (0g(44)) 12
+
4L
[10 - 1'1 E
(32)
pPmin
O, x E X, we have that Lemma E.1. For any m
1 Pmiin
r [(x 0) S
Proof. Applying Hölder's inequality have that
f(x, 0) < |lll 0)1
= =||||
Moreover, since Thm have that Σ, (y) This implies that ||||| yields the result.
which
Panin
Lemma E.2 (Implication of Assumption Assumption 5.1 such that for any r E RT, O x X.
Under Assumption 5.1 there exists B
|log(r' f(x, 0)) S B.
Proof. Since R.m. it has at least non-zero coordinate and f(x 0) is the output of sofumax layer so all of its coordinates are non-zero. Consequently,
r f(x, 0
So by Assumption 5.1 the function (r, 0, x) + log(r f(x, 0)) is defined and continuous over a compact set, there exists constant giving the result.
24

---

Published as a conference paper at ICLR 2025
Lemma E.3 (Population Strong Convexity). Let H (rr,0, 2r) -log(r' fog( 0)). Under Assump- tion Assumption 5.2 the function
C.. r B H( 0 xx)
is /APmiu 1 stronngy convex.
Proof. We first compute the Hessian of L to find that
(2((r) = (r f(x, 0.))2 f(x 0* (xx,
1
Since by Lemma E.1 we have that r f(x, Pmiin we conclude that
v2r(() ž Pmin f(æ, (() = PPmiu
Lemma E.4 (Lipschitz Parametrization) Let (r, 0,,,0 (xr There exists L L>> such that for any 0,, O2 E G and r E IRT we have that
[H(r 01, x) (r, 02 x) S L| 02112.
Proof. The gradient of H with respect to O is given by
1
Vo
(xx Vef(2,0)
Reasoning like in Lemma F.1 we know that T defined and continuous over the compact set of its parameters, we also know that is a neural network parametrized by 0, hence [a 0) is bounded when and æ are bounded. Consequently, under Assumption 5.1 there exists a constant L >> 0 such that
(r S L.
Lemma E.5 (Uniform Bound 1). Let 8 E (0, 1), with probability 1 8, we have that
E H 0,,
H( 0,,
J=1
(33)
2 Vn
log(4/8)
S
Rad(J)) + +BB
71
Proof. Let 6 (0, 1). Since f. is learned from the samples zy, do not have independence, which would have allowed us to apply a concentration inequality. Hence, we derive uniform bound as follows. We begin by observing that:
n Σ j=1
E
H 21.
H(fn. Og x.)
< sup r
E H
Σ H(r 0 x,,
Now since Lemma E.2 holds, we can apply McDiarmid's Inequality to get that with probability 1 5, we have:
n Σ
sup r
E (1(r,0,0)
(r, 0,x;)
n
j=1
H(r, 022)) )
log (2/8)
E sup
E B[H(r, m))
Σ
H
2B 2BV
r.0
25

---

Published as a conference paper at ICLR 2025
The expectation of the supremum the right-hand side can be bounded by the Rademacher com- plexity F = { r f(x, 0), (r,0) E RT x ,) and we obtain:
E (r. 0, 0,0)]
Σ H(r, 0, 0,
sup r,
(34)
2 Rad(FF + 2B
log (2/6)
Lemma E.6 (Uniform Bound 2). Let ō E (0, 1), with probability 1 6, have that
E H(rs.,
((rr-,,,.20)
==1
(35)
2 2B
log (2/(2
Rad(F)
nn
n
Proof.
The proof is identical to that of Lemma E.5
Lemma E.7 (Strong Convexity of Population Loss). Let L(r, 0) be the population loss as defined in Lemma E.7 We establish that C(r, 0) ("Pmin-Stoongh convex under the assumptions of calibration
(Assumption 5.2
Proof. We compute the Hessian of the population loss as in Lemma E.7 obtaining that:
[r((r) = f(x, 0)}= f( () OT
From Lemma E.1 we have that T f(x, 0) S Pmin Therefore, we conclude:
v2L(r) > Pmmi E 1(, 0)7
2 PPmin I-
Lemma E.8 (Bound on Empirical Loss). Under Assumption 5.1 the empirical loss C, (r, 0n.) satisfies the following concentration bound:
P sup FERM
(r, Owx) C(r, C(r, Ow.) > s 2exp (-cn"t
Proof. This result follows from standard concentration inequalities, such as McDiarmid's inequality. together with the Lipschitz continuity of the loss function with respect to the samples.
26

---

Published as a conference paper at ICLR 2025
F PROOF OF THEOREM 5.2 AND GUARANTEES FOR IW-ERM WITH VRLS
We now establish convergence rates for IW-ERM with VRLS and show our proposed importance weighting achieves the same rales with the data-dependent constant terms increase linearly with m sup, r Tmax under negligible communication overhead over the baseline ERM- solvers without importance weighting. In Appendix FF] we establish tight convergence rates and communication guarantees for IW-ERM with VRLS in broad range of importance optimization settings including convex optimization. second-order differentiabilityy composite optimization with proximal operator, optimization with adaptive step-sizes, and nonconvex optimization, along the lines of e.g., (Woodworth et al. 2020 Haddadpour et al. 2021 | Glasgow et al. 2022 Liu et al. 2023 Hu
Huang 2023 Wu al.
2023 Liu et al. 2023).
By estimating the ratios locally and absorbing into local losses, we note that the properties of the modified local loss w.r.t. the neural network parameters w, e.g., convexity and smoothness, do not change. The data-dependent parameters such as Lipschitz and smoothness constants for & o hew w.r.t. w are scaled linearly by T'aax. Our method of importance ratio estimation trains the pre-defined predictor exclusively using local training data, which implies IW-ERM with VRLS achieves the same privacy guarantees as the baseline ERM-solvers without importance weighting. For ratio estimation, the communication between clients involves only the estimated marginal label distribution, instead of data, ensuring negligible communication overhead. Given the size of variables to represent marginal distributions, which is by orders of magnitude smaller than the number of parameters of the underlying neural networks for training and the fact that ratio estimation involves only one round of communication, the overall communication overhead for ratio estimation is masked by the communication costs of model training. The communication costs for IW-ERM with VRLS over the course of optimization are exactly the same those of the baseline ERM-solvers without importance weighting. All in all, importance weighting does not negatively impact communication guarantees throughout the course of optimization, which proves Theorem 5.2
In the following, we establish tight convergence rates and communication guarantees for IW-ERM with VRLS in a broad range of importance optimization settings including convex optimization. second-order differentiability. composite optimization with proximal operator, optimization with adaptive step-sizes, and nonconvex optimization.
For convex and second-order Differentiable optimization, we establish a lower bound on the conver- gence rates for IW-ERM in with VRLS and local updating along the lines of e.g., (Glasgow et al.. 2022 Theorem 3)).
Assumption F.1 (PL with Compression). 1) The (zb v) B-smoothness and convex w.r.t. w for any (x, y) and satisfies Lojyyk--ojasewiizz (PL) condition (there exists at >> 0 such that, for all w W. we have ((hw) < Vr (22/(2xx): 2) The compression scheme Q is unbiased with bounded variance, i.e., E[[(x)) = and [|]| [(x) 2 < |||: [[2]; 3) The stochastic gradient g(w)= V ((ww)) is unbiased, Le.. E[g(w)) = V (h for any E W with bounded variance Blllg(w) - V w (hvw) ||]].
For nonconvex optimization with PI. condition and communication compression, we establish convergence and communication guarantees for IW-ERM with VRLS, compression, and local updating along the lines of e.g., (Haddadpour et al. 2021 Theorem 5.1).
Theorem F.1 (Convergence and Communication Bounds for Nonconvex Optimization with PL). Let denote the condition number, t denote the number of local steps, R denote the number of communi- cation rounds, and max sup, ry(y) = Tmaxx Under Assumption F.I. suppose Algorithmm 2 with local updates and communication compression (Haddadpour et af 2027 Algorithm 1) is for TR total stochastic grudients per node with fixed step-sizes 1) = 1/2rmaxxyT(()
1)) and
7 K. Then we have E(((hwrr) [(h >) S by setting
9 K
+1)xcogg (:) and T K(q/K + 1)e)
9+1
R
(36)
Assumption F.2 (Nonconvex Optimization with Adaptive Step-sizes). The l o haw is 3-smoothness
with bounded gradients; 2) The stochastic gradients g(w) & (h(( is unbiased with bounded
variance E[llg(w) - V & ||2;; 3) Adaptive matrices A, constructed as in (Wu et al. 2023
Algorithm 2) are diagonal and the minimum eigenvalues satisfy Amin (At) 2 p > 0 for some p R4.
27

---

Published as a conference paper at ICLR 2025
For nonconvex optimization with adaptive step-sizes, establish convergence and communication guarantees for IW-ERM with VRLS and local updating along the lines of e.g., (Wu et al. 2023 Theorem 2).
Theorem F.2 (Convergence and Communication Guarantees for Nonconvex Optimization with Adaptive Step-sizes). Let T denote the number of local steps, R denote the number of communication rounds, and maXyey s f (y) Tmax. Under Assumption F2 suppose Algorithm 2 with local
updates is run for T TR total stochastic gradients per node with an adaptive step-size similar (Wu et al.
2023
Algorithm 2). Then EE| [ Vell( )l2 by setting:
Tmax T Ke
and R
Tmax
(37)
Assumption F.3 (Composite Optimization with Proximal Operator). 1) The E hw is smooth and strongly convex with condition number K; 2) The stochastic gradients g(w) ww((hw) is unbiased.
For composite optimization with strongly convex and smooth functions and proximal operator, establish an upper bound on oracle complexity to achieve error on the Lyapunov function defined as in (Hu & Huang 2023 Section 4) for Gradient Flow-type transformation of IW-ERM with VRLS in the limit of infinitssimal step-size.
Theorem F.3 (Oracle Complexity of Proximal Operator for Composite Optimization). Let i. denote the condition number. Under Assumption E3 suppose Gradient Flow-type transformation of Al- gorithm with VRLS and Proximal Operator evolves the limit of infinitesimal step-size (Hu & Huang 2023 Algorithm 3). Then it achieves O VR log(1/c)) Proximal Operator Complexity.
28

---

Published as a conference paper at ICLR 2025
G COMPLEXITY ANALYSIS
In our algorithm, the ratio estimation is performed once in parallel before the IW-ERM step.
In the experiments, used simple network to estimate the ratios in advance, which required significantly less computational effort compared to training the global model. Although IW-ERM with VRLS introduces additional computational complexity compared the baseline FedAvg, results in substantial improvements in overall generalization, particularly under challenging label shift conditions.
29

---

Published as a conference paper at ICLR 2025
H
MATHEMATICAL NOTATIONS
In this appendix, we provide a summary of mathematical notations used in this paper in Table
Table 5: Math Symbols
Math Symbol
Definition
X
Compact metric space for features
Discrete label space with Y = m
K
Number of clients in an FL setting All samples in the training set client
Se
hw
Hypothesis function h.. X + y Hypothesis class for liw Mapping space from A, which be discrete continuous
Z
30

---

Published as a conference paper at ICLR 2025
I LIMITATIONS
The distribution shifts observed in real-world data are often not fully captured by the label shift or relaxed distribution shift assumptions. In our experiments, we applied mild test data augmentation approximate the relaxed label shift and manage ratio estimation errors for both the baselines and our method. However, the label shift assumption remains overly restrictive, and the relaxed label shift lacks robust empirical validation in practical scenarios.
Additionally, IW-ERM's parameter estimation relies local predictors at each client, which limits its scalability. In practice, simpler global predictor could be sufficient for parameter estimation and IW-ERM training. Future research could explore VRLS variants capable of effectively handling more complex distribution shifts in challenging datasets, such as CIFAR-10.1 (Recht et al., [2018] Torralba (Garg et al. 2023).
al. 2008), as suggested
31

---

Published as a conference paper at ICLR 2025
J EXPERIMENTAL DETAILS AND ADDITIONAL EXPERIMENTS
In this section, we provide experimental details and additional experiments. In particular, we validate our theory on multiple clients in federated setting and show that our IW-ERM outperforms FedAvg and FedBN baselines under drastic and challenging label shifts.
J.1 EXPERIMENTAL DETAILS
In single-client experiments, a simple MLP without dropout is used as the predictor for MNIST, and ResNet- 18 for CIFAR-10.
federated learning setting, both MNIST (LeCun et al., 1998 and Fashion MNIST (Xiao al., 2017) datasets are employed, each containing 60,000 training samples and 10,000 test samples. with each sample being a 28 by 28 pixel grayscale image. The CIFAR-10 dataset (Krizhevsky) comprises 60,000 colored images, sized 32 by 32 pixels, spread across 10 classes with 6,000 images per class; it is divided into 50,000 training images and 10,000 test images. In this setting, the objective is to minimize the cross-entropy loss. Stochastic gradients for each client are calculated with a batch size of 64 and aggregated on the server using the Adam optimizer. LeNet used for experiments on MNIST and Fashion MNIST with learning rate of 0.001 and a weight decay of x 10 6, For CIFAR-10, 10, ResNet-18 18 is employed with a learning rate of 0.0001 and weight decay of 0.0001. Three independent runs are implemented for 5-client experiments on Fashion MNIST and CIFAR-10, 10, while for 10 clients, one run is conducted on CIFAR-10. The regularization coefficient& Ç in Equation 4 is set to for all experiments. All experiments are performed using a single GPU on an internal cluster and Colab.
For experiments
Importantly, the training of the predictor for ratio estimation on both the baseline MLLS and our VRLS is executed with identical hyperparameters and epochs for CIFAR-10 and Fashion MNIST. The training is halted once the classification loss reaches a predefined threshold on MNIST.
J.2 RELAXED LABEL SHIFT EXPERIMENTS
In conventional label shift, is assumed that p( y) remains unchanged across training and test data. However, this assumption is often too strong for real-world applications, such as in healthcare, where different hospitals may use varying equipment, leading to shifts in p(z y) even with the same labels (Rajendran et al.) 2023). Relaxed label shift loosens this assumption by allowing small changes
the conditional distribution (Garg et al. 2023 Luo & Ren 2022).
To formalize this, we use the distributional distance 'D and a relaxation parameter 0, as defined by Garg et al. (2023): max, D (pr((: y), (t: v)) S This allows for slight differences in feature distributions between training and testing, capturing more realistic scenario where the conditional distribution is not strictly invariant.
In our case, visual inspection suggests that the differences between temporally distinct datasets, such as CIFAR-10 and CIFAR-10. 1_v6.. (Torralba et al. 2008) Recht et al., 2018), may not meet the assumption of small To address this, we instead simulate controlled shifts using test data augmentation, allowing us to regulate the degree of relaxation, following the approach outlined in Garg et al. 2023)
J.3 ADDITIONAL EXPERIMENTS
In this section, we provide supplementary results, visualizations of accuracy across clients and tables showing dataset distribution FI. setting and relaxed label shift.
32

---

Published as a conference paper at ICLR 2025
0.10
MLLS, EM, baseline
10
MLLS, EM baseline baseline MLLS_ CG
10
WAW
0.08
MILS baseline
MLLS, EM baseline
MLLS CG
MSE
MLLS, 12 baseline
10-22
10-22
MLLS, baseline
MLLS, CG
0.00
10
10a Shift paramater
4000 6000
10000
2000 8000 10000
Sample
Sample
Figure 3: MSE analysis on MNIST for MLLS baselines. Left: Performance evaluation across various alpha values, comparing different methods: MLLS EM, MLLS.I, MLLS_L2, and MLLS_CG, MLLS LI and MLLS L2 utilize convex optimization with L1 and L2 regularization for estimating our limited test sample problem, respectively, and are solved directly with a convex solver. In contrast, MLLS CG uses conjugate gradient descent and MLLS EM solves this convex optimization problem with EM algorithm. Both the EM and convex optimization methods (MLLS LI, MLLS 1.2) demonstrate superior and more consistent performance, especially under severe label shift conditions, when compared to MLLS.CG. Middle: At an alpha value of 1.0. the MSE analysis shows comparable performance across most methods, with the exception of MLLS. ___G,, which lags behind. Right: For alpha=0.1, MILS CG performs significantly worse than the EM and convex optimization methods, consistent with the trends observed in the left plot.
VRLS_ EM
VRLS EM
0.008
VRLS 12
VRLS_ 12 MLLS EM EM_baseline
MLLS EM baseline
MSE
0.006
MLLS 12 baseline
MSE
-0
MLLS 12 bascline
RLLS
var RLLS
0.004
0.002
10-3
10-1
10c
10a
0
2000 4000 6000 8000 10000
Shift paramater
Sample size
Figure 4: In our detailed analysis with the MNIST dataset, conduct thorough comparison of VRLS alongside MLLS (Garg et al. 2020), EM (Saerens et al. 2002). and also RLLS (Azizzade) nesheli al. 2019).
Table 6: LeNet on Fashion MNIST with label shift across 5 clients. 15,000 iterations for FedAvg and FedBN; 5,000 for Upper Bound (FTW-ERM) using true ratios and our IW-ERM. To mention, to train our predictor, we use a simpliest MLP and employ linear kernel.
FMNIST
Our IW-ERM
FedAvg
FedBN
Upper Bound
Avg. accuracy
0.7520 + 0.0209
0.5472
0.0297
0.5359 + 0.0306
0.8273 0.0041
Client I accuracy
0.7162 + 0.0059
0.3616 0.0527
0.3261 + 0.0296
0.8590 + 0.0062
Client accuracy
0.9266 + 0.0125
0.9060
0.0157
0.9035 + 0.0162
0.9357
0.0037
Client 3 accuracy
0.6724 = 0.0467
0.3279 0.0353
0.3612 = 0.0814
0.7896
0.0109
Client 4 accuracy
0.7979 = 0.0448
0.6858 0.0105
0.6654 = 0.0121
0.8098 = 0.0112
Client 5 accuracy
0.6468 + 0.0248
0.4548 + 0.0655
0.4234 = 0.0387
0.7426 0.0257
33

---

Published as a conference paper at ICLR 2025
VRLS_EM
VRLS EM
MLLS_EW baseline
MLLL SESS_EE___ baseline
0.05
10-1
MSE
0.04
MSE
0.03
0.02
1
10"
10a
2000 4000
8000 10000
Shift paramater
Shift paramater
Figure In this experiment with Fashion MNIST, simple MLP with dropout were employed.
Table 7: ResNet-18 18 CIFAR-10 with label shift across 5 clients. For fair comparison, we run 5,000 iterations for our method and Upper Bound, while 10000 for FedAvg and FedBN.
CIFAR-10
Our IW-ERM
FedAvg
FedBN
Upper Bound
Avg. accuracy
0.5640 + 0.0241
0.4515 + 0.0148
0.4263 + 0.0975
0.5790 0.0103
Client I accuracy
0.6410 + 0.0924
0.5405 + 0.1845
0.5321 + 0.0620
0.7462 + 0.0339
Client 2 accuracy
0.8434 = 0.0359
0.3753 = 0.0828
0.4656 = 0.2158
0.7509
0.0534
Client 3 accuracy
0.4591 = 0.1131
0.3973 0.1333
0.2838 = 0.1055
0.5845
0.0854
Client 4 accuracy
0.4751 = 0.1241
0.5007 = 0.1303
0.5256 = 0.1932
0.3507
0.0578
Client 5 accuracy
0.4013 + 0.0430
0.4429 + 0.1195
0.5603 + 0.1581
0.4627 0.0456
Figure 6: The average, best-client, and worst-client accuracy, along with their standard deviations, derived from Table Our method exhibits the lowest standard deviation, showcasing the most robust accuracy amongst the compared methods.
Comper son Mestees on CItA8110
Petthoda
Figure 7: The average, best-client, and worst-client accuracy, along with their standard deviations, are derived from Table
34

---

Published as a conference paper at ICLR 2025
Table 8: Label distribution on Fasion MNIST with 5 nodes, with the majority of classes possessing limited number of training and test images across each node.
Class
2
4
6
7
Train
34
34
34
34
34
5862
34
34
34
34
Node I
Test
977
5
5
5
5
Train
34
34
34
34
34
34
5862
34
34
34
Node 2
Test
977
5
5
5
Train
34
34
34
34
34
34
34
5862
34
34
Node 3
Test
5
977
5
5
5
Train
34
34
34
34
34
34
34
34
5862
34
Node 4
Test
5
5
5
977
5
5
5
5
Train
34
34
34
34
34
34
34
5862
Node 5
Test
5
5
5
977
5
5
5
5
Table 9: Label distribution on CIFAR-10 with 5 clients, with the majority of classes possessing limited number of training and test images across each client.
Class
2
4
6
7
Train
34
34
34
34
34
5862
34
34
34
34
Node 1
Test
977
5
5
Train
34
34
34
34
34
34
5862
34
34
34
Node 2
Test
977
5
5
5
Train
34
34
34
34
34
34
34
5862
34
34
Node 3
Test
5
5
977
5
5
Train
34
34
34
34
34
34
34
34
5862
34
Node 4
Test
5
5
977
5
5
5
Train
34
34
34
34
34
34
34
34
5862
Node 5
Test
5
5
5
977
5
5
5
35

---

Published as a conference paper at ICLR 2025
Table Label distribution on CIFAR-I0 with 100 clients, wherein groups of 10 clients share the same distribution and ratios. The majority of classes possess limited quantity of training and test images on each client.
Class
0
2
3
4
Train
95/100
5/9
0/9
6/9
5/9
Client 1-10
Test
5/9
6/9
f/9
6/9
6/9
Train
5/9
95/100
5/9
5/9
5/9
Client 11-20
Test
5/9
5/9
5/9
5/9
5/9
Train
5/9
5/9
96/100
5/9
5/9
Client 21-30
Test
5/9
5/9
5/9
5/9
5/9
Train
5/9
5/9
5/9
95//100
5/9
Client 31-40
Test
5/9
6/9
6/9
6/9
6/9
Train
6/9
6/9
6/9
6/9
96 100
Client 41-50
Test
5/9
5/9
5/9
5/9
5/9
Train
5/9
5/9
5/9
5/9
5/9
Client 51-60
Test
5/9
5/9
5/9
5/9
95/100
Train
5/9
5/9
6/9
5/D
5/g
Client 61-70
Test
5/9
5/9
5/9
95/100
5/9
Train
5/9
5/9
5/9
5/9
5/9
Client 71-80
Test
6/9
5/9
95/100
6/9
5/9
Train
5/9
5/9
5/9
5/9
5/9
Client 81-90
Test
6/9
95/100
5/9
5/9
5/9
Train
5/9
5/9
6/9
5/9
6/9
Client 91-100
Test
95/100
5/9
f/9
6/9
5/9
Class
6
7
9
Train 5/9
6/9
5/9
5/9
5/9
Client 1-10
Test
5/9
5/9
5/9
5/9
95/100
Train
3/9
5/9
5/g
3/9
Client 11-20
Test
5/9
5/9
5/9
95/100
5/9
Train
5/9
6/9
5/9
b/g
5/9
Client 21-30
Test
5/9
5/9
95/100
5/9
5/9
Train
5/9
5/9
5/g
5/9
Client 31-40
Test
5/9
95/100
5/9
5/9
5/9
Train
5/9
5/9
5/9
5/9
Client 41-50
Test
95/100
5/9
5/9
5/9
5/9
Train
95/100
6/9
5/9
6/9
5/9
Client 51-60
Test
5/9
5/9
5/9
5/9
5/9
Train
5/9
M6/100
6/9
5/g
6/9
Client 61-70
Test
5/9
5/9
5/9
5/g
3/9
Train
5/9
5/9
95/100
5/9
5/9
Client 71-80
Test
5/9
5/9
5/9
5/g
5/9
Train
5/9
5/9
5/9
95/100
5/9
Client 81-90
Test
6/9
6/9
5/g
6/9
Train
6/9
6/9
9/9
6/9
93 100
Client 91-100
Test
5/9
5/9
5/9
5/9
5/9
36


Reviews:
[{"review_id": "bjzuiNBCEH", "reviewer": "Reviewer_og3p", "summary": "This paper solves the problem where distributed data nodes have different label distributions in their training and test data. The authors propose VRLS - a method that estimates how these distributions differ between nodes and use this information to improve model training. Their approach maintains data privacy while achieving 20% better accuracy than existing methods.", "questions": "same as the third point of weakness: why minimizing equation 2 leads to better conditional distribution estimation?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. This paper clearly articulates its research focus by addressing both inter-node and intra-node label shifts in distributed learning systems, which represents a more comprehensive approach than prior works.\n2. The paper provides extensive theoretical proofs to support the proposed algorithm, though I have not thoroughly examined all of them.", "weaknesses": "1. The paper lacks clear logical flow and detailed algorithm explanations. While the title emphasizes distributed learning, Section 4 (covering VLRS in distributed environments) is surprisingly brief - only one-third of a page. It simply states that minimizing equation (IW_ERM) enables VLRS in distributed settings, without proper elaboration. In contrast, Section 5 on VLRS bounds spans two full pages. It's puzzling why the authors prioritized theoretical bounds over making their work more accessible to readers.\n2. Lines 193-194 contain an inconsistency where the authors claim \"Our proposed VRLS objective aims to better approximate the conditional distribution ptr(y|x)\", but equation 1 (VRLS objective) actually estimates r. The conditional distribution approximation corresponds to equation 2.\n3. The authors fail to adequately explain why minimizing equation 2 leads to better conditional distribution estimation. While Ω(fθ) represents entropy and promotes more dispersed neural network outputs as claimed, there's no clear justification for how this ensures f(x) better estimates ptr(y|x).\n4. The paper suffers from poor notation management. Symbols like Ω(fθ) should be defined at first appearance rather than later, as this significantly increases the time needed to understand the paper.\n5. The experimental validation is unconvincing, relying only on synthetic datasets (MNIST and CIFAR-10) with small images and just 10 classes. The evaluation should include larger datasets like ImageNet and real-world scenarios. The lack of real-world validation raises concerns about whether the paper's assumptions are practical.", "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": "1zheso3yT8", "reviewer": "Reviewer_4bHP", "summary": "The paper proposes the Versatile Robust Label Shift (VRLS) method to address label shift in distributed learning systems. VRLS employs entropy regularization to improve density ratio estimation between training and test label distributions, making it resilient to both intra-node and inter-node label shifts. The paper further extends VRLS to a multi-node setup through an Importance Weighted ERM (IW-ERM) framework, which balances label distribution estimation across nodes while preserving data privacy. Empirical results show VRLS outperforming baselines by up to 20% in imbalanced settings, and theoretical bounds confirm VRLS's estimation accuracy and convergence rates.", "questions": "1-  Could the authors discuss how VRLS would perform with other regularization methods, such as focal loss, or compare its performance with maxent loss for calibration? Would such alternatives affect calibration or robustness?\n\n2- Can the authors provide empirical or theoretical evidence on VRLS’s performance under non-uniform or more complex label shift distributions, such as class-conditional label shift?\n\n3- How does VRLS handle scalability in networks with significantly more nodes or nodes with varying computational capacities? Are there adjustments to VRLS or IW-ERM to address such scenarios?\n\n4-  Given that assumptions like bounded data and sub-Gaussian noise are not always valid in practice, could the authors discuss how VRLS might be adapted or how results may vary when these assumptions do not hold?\n\n5-  VRLS is designed for classification, but can it be adapted for regression tasks where label distributions may also shift? If so, what adjustments would be needed?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1-VRLS introduces a unique entropy regularization term that enhances density ratio estimation by penalizing degenerate solutions, thus yielding well-calibrated predictions. Unlike existing methods such as MLLS and BBSE, which do not adequately capture uncertainty, VRLS adapts the estimation process to real-world variability. This improvement aligns with recent work in calibration under distribution shifts (e.g., Guo et al., 2017), offering both theoretical and practical robustness.\n\n2-The paper’s thorough empirical evaluation spans single-node and multi-node scenarios, with tests on MNIST, Fashion MNIST, and CIFAR-10. In distributed settings, VRLS with IW-ERM adapts effectively, providing significant accuracy improvements across nodes while handling inter-node and intra-node label shifts. Compared to methods such as FedAvg and FedProx, which struggle with label shift, VRLS maintains robust performance in both settings, a key distinction from prior approaches.\n\n3- Theoretical error bounds and convergence rates for VRLS are rigorously established. The authors show that the VRLS method’s estimation error decreases proportionally to the square root of test sample size, a valuable addition to the literature where empirical results often lack formal guarantees. This theoretically grounded approach helps ensure VRLS's performance consistency in both high and low label-shift scenarios, a substantial contribution over prior studies that assume stability without formal validation.\n\n4- By using locally computed density ratios without requiring data exchange, IW-ERM maintains data privacy across nodes, a marked improvement over traditional federated approaches. This reduction in communication rounds, coupled with enhanced convergence rates, renders IW-ERM with VRLS practical for real-world applications where privacy and resource efficiency are paramount, setting it apart from resource-intensive federated models like FedBN.", "weaknesses": "1- The paper effectively shows VRLS's reliance on Shannon entropy regularization but does not explore how it compares to other regularization techniques, such as focal loss or maxent loss. A comparative analysis across regularization methods on the same datasets, particularly with focal loss regularization, would help clarify VRLS's robustness and calibration capabilities. Additionally, a theoretical discussion on the effects of different regularizations on model calibration and uncertainty estimation could deepen the understanding of VRLS's effectiveness.\n\n2-VRLS primarily addresses scenarios under symmetric or Dirichlet-based label shifts, with limited discussion on its performance with highly non-uniform or asymmetric label shifts. Extending empirical analysis to include more diverse label distributions would provide a clearer picture of VRLS’s generalizability to complex, real-world shifts.\n\n3- The theoretical error bounds rely on assumptions that may not fully align with real-world distributed data characteristics, such as bounded data spaces and sub-Gaussian noise. Providing examples of scenarios where these assumptions might fail, along with discussing VRLS’s potential performance in such cases, would add rigor to the theoretical claims. Additionally, testing VRLS with artificially unbounded or heavy-tailed data distributions would offer practical insights into the method’s robustness beyond the idealized assumptions, adding depth to its theoretical grounding.\n\n4- The paper’s focus on label shift in classification limits its applicability to other supervised learning tasks, such as regression or multi-label classification, where label distributions differ. Addressing how VRLS might generalize to other tasks, or clarifying any modifications needed, could extend the method’s relevance across broader distributed learning applications.\n\n5- The VRLS method’s reliance on entropy-based regularization (inspired by Shannon entropy) is theoretically motivated by preventing degenerate solutions. However, there is limited theoretical analysis on whether this regularization could inadvertently introduce biases, especially in settings where label shifts deviate from the assumptions about class-conditional stability. This could be explored further in the theoretical results or discussed as a potential limitation.", "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": "gohiJQz9kM", "reviewer": "Reviewer_Ym52", "summary": "This paper addresses the challenge of inter-node and intra-node label shifts in distributed learning. Test-to-train density ratio estimation plays an important role in addressing the label-shift problem. The authors proposed the Versatile Robust Label Shift (VRLS) method, which incorporates a regularization term based on Shannon entropy. This regularization term significantly improves the estimation of the conditional distribution $p^{tr}(\\mathbf{y}|\\mathbf{x})$, thereby enhancing the estimation of the label density ratio. The authors conduct experiments on MNIST, Fashion-MNIST, and CIFAR-10 datasets to verify the effectiveness of the proposed method. The authors also provide theoretical analysis for the estimation error and convergence.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The authors proposed a simple but effective method to address the problem of label shift in distributed learning.\n2. The authors conduct experiments to verify their proposed method. The MSE for the density ratio of the proposed method is lower than that of baselines. The test accuracy of the proposed method is also better than that of baselines.\n3. Bounds on ratio estimation and convergence rates are analyzed theoretically.", "weaknesses": "1. The definition of inter-node and intra-node label shifts is ambiguous. I suggest that the authors define them more clearly.\n2. In Lines 196-198, why penalizing the Shannon entropy of the softmax output can improve the approximation of $p^{tr}(\\mathbf{y}|\\mathbf{x})$? In other words, why the smoother output is a better approximation of $p^{tr}(\\mathbf{y}|\\mathbf{x})$ (Line 207-208)?\n3. What is the meaning of $\\lambda_k$ in the risk of IW-ERM?\n4. Lack of sensitivity test for the hyperparameter $\\zeta$ in Equation (2).\n5. The font sizes in Fig 1 and Fig 2 are inconsistent.", "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": "ER3s3N5FUt", "reviewer": "Reviewer_L7cy", "summary": "This paper proposes the Versatile Robust Label Shift (VRLS) method, which enhances the maximum likelihood estimation of the test-to-train label density ratio, for multi-node distributed learning. It incorporates Shannon entropy-based regularization and adjusts the density ratio during training to better handle label shifts at the test time. Experiments conducted on MNIST, Fashion MNIST, and CIFAR-10 demonstrate the effectiveness of the proposed VRLS. The main contributions of this paper are:\n\n- It proposes VRLS to enhance the training of the predictor of the conditional distribution by incorporating a novel regularization term based on Shannon entropy.\n\n- It establishes high-probability estimation error bounds for VRLS, as well as high-probability convergence bounds for IW-ERM with VRLS in nonconvex optimization settings.", "questions": "1. The concept of \"underlying uncertainty\" needs to be further elaborated. Specifically, why does the failure of a learned predictor to inherently capture the underlying uncertainty have an impact on model performance? What is the Role of \"underlying uncertainty\" in multi-node distributed learning?\n2. The motivation for the introduction of Shannon entropy-based regularization also should be further explained by the authors. What is the advantage of the leveraged technology compared with others?\n3. The comparison methods in this paper are all proposed before 2021, and lack the latest algorithms. The comparisons between the proposed method and the latest algorithms are necessary.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The problem studied in this paper is interesting.\n\n- The paper is well-organized, which is easy to follow.\n\n- This paper provides some novel perspectives for multi-node distributed learning.", "weaknesses": "- In my opinion, the concept of \"underlying uncertainty\" needs to be further elaborated. Specifically, why does the failure of a learned predictor to inherently capture the underlying uncertainty have an impact on model performance? What is the Role of \"underlying uncertainty\" in multi-node distributed learning?\n- The motivation for the introduction of Shannon entropy-based regularization also should be further explained by the authors. What is the advantage of the leveraged technology compared with others?\n- For the experimental evaluation, this paper compares the proposed method with some baseline methods. The experimental results seem promising. However, the comparison methods in this paper are all proposed before 2021, and lack the latest algorithms (e.g. [1]-[2]). I think the comparisons between the proposed method and the latest algorithms are necessary.\n\n\n\n[1] [GradMA: A Gradient-Memory-based Accelerated Federated Learning with Alleviated Catastrophic Forgetting. CVPR 2023: 3708-3717.](https://ieeexplore.ieee.org/document/10203330/)\n\n[2] [Multi-Level Branched Regularization for Federated Learning. ICML 2022: 11058-11073.](https://proceedings.mlr.press/v162/kim22a.html)", "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)