REVIEW 3 major objections 6 minor 28 references
Addressing Label Shift in Distributed Learning via Entropy Regularization
T0 review · 3 major / 6 minor · reviewed 2026-08-09 · deepseek-v4-flash
Pith's one-line read Adding a Shannon-entropy penalty to the predictor used for label-shift estimation yields more accurate test-to-train density ratios, and weighting distributed ERM by these ratios improves accuracy by up to 20% over standard baselines…
desk verdict Clean empirical win for entropy-regularized label-shift ratio estimation, but the theory does not actually show the regularizer helps and the distributed experiments lack a no-regularizer baseline. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The load-bearing object is the entropy-regularized predictor. The regularizer is the Shannon entropy of the softmax output, $\Omega(f_\theta) = \sum_c \phi(f_\theta(x))_c \log \phi(f_\theta(x))_c$, added to cross-entropy with strength $\zeta$. Its job is to prevent the predictor from producing over-confident probabilities, so that $f_\theta(x)$ stays close to the true conditional $p_{\mathrm{tr}}(y|x)$; the MLLS estimator then uses these calibrated probabilities in the concave objective $\arg\max_r \mathbb{E}_{\mathrm{te}}[\log(f_{\theta^*}(x)^\top r)]$. The second piece is the aggregated ratio $r_k(y)$ that converts locally estimated ratios into the global importance weight used in the IW-ERM objective. The theoretical engine is the strong-convexity-and-calibration analysis: under the assumption that $\mathbb{E}[f_{\theta^*}(x)f_{\theta^*}(x)^\top] \succeq \mu I_m$, the population estimation loss is $\mu p_{\min}$-strongly convex, and uniform Rademacher bounds control the finite-sample error caused by estimating the ratio with an estimated predictor.
What would settle it
Take a dataset whose class-conditional feature distribution is shifted between train and test, for instance Gaussian class conditionals with different means or stronger blur and brightness changes than the mild augmentations used in the paper, compute the true test-to-train label ratios by sampling, and compare VRLS ratio estimates with plain MLLS estimates.
Extended reading notes
Core claim
The paper's central claim is that entropy regularization of the predictor makes the maximum-likelihood label-shift machinery work well enough to be useful in distributed learning. VRLS trains a predictor by minimizing $\mathbb{E}_{\mathrm{tr}}[\ell_{\mathrm{CE}}(f_\theta(x), y) + \zeta \Omega(f_\theta)]$, where $\Omega(f_\theta) = \sum_{c=1}^m \phi(f_\theta(x))_c \log \phi(f_\theta(x))_c$ is the Shannon entropy of the softmax output, and then estimates the test-to-train ratio as $\hat{r} = \arg\max_{r\in\mathbb{R}^m_+} \mathbb{E}_{\mathrm{te}}[\log(f_{\theta^*}(x)^\top r)]$. In the multi-node setting, each node estimates ratios locally and shares only the small marginal-label vector; the aggregated ratio $r_k(y) = \frac{\sum_{j=1}^K p_{\mathrm{te}}^j(y)}{p_{\mathrm{tr}}^k(y)}$ weights each node's contribution to the global IW-ERM objective. The theory, stated as Theorem 5.1, bounds $\|\hat{r}_{n_{\mathrm{te}}} - r_{f^*}\|_2$ by terms involving the Rademacher complexity of the ratio-predictor class, the calibration constant $\mu$, and the distance between the empirical and ideal predictor parameters. The experimental claim is that these estimates are accurate enough that the resulting global model approaches the upper bound obtained by feeding true ratios into the same objective.
Load-bearing premise
The load-bearing assumption is that the distribution of input features within each class, $p(x|y)$, is exactly the same at every node and in the test data; only the label proportions may change.
Editorial extensions
If this is right
- VRLS-regularized ratio estimates should remain more accurate than plain MLLS estimates as the test sample grows, since the Theorem 5.1 error bound decreases at rate $1/\sqrt{n_{\mathrm{te}}}$.
- Distributed nodes can correct each other's label shifts while sharing only a marginal label vector, so the method preserves the privacy profile of local-only training.
- Importance weighting scales the smoothness and Lipschitz constants of the global objective by $r_{\max}$ but does not change the convergence rate or communication rounds compared with an unweighted ERM solver, as stated in Theorems 5.2 through 5.5.
- The global IW-ERM model trained with VRLS estimates should approach the accuracy of the same objective with true ratios, which the experiments confirm on Fashion MNIST and CIFAR-10.
- In the large-scale setting with 100 and 200 nodes and random participation, IW-ERM with VRLS still improves over standard distributed baselines while using fewer training iterations.
Reading between the lines
- Because the benefit comes from counteracting over-confidence, VRLS may help most when the base predictor is poorly calibrated; on already well-calibrated predictors the entropy penalty could in principle wash out or slightly hurt, a regime the paper does not isolate.
- The paper's theory and central experiments keep the class-conditional feature distribution fixed; a natural extension is to combine VRLS with feature-level correction so that the ratio estimates remain meaningful under relaxed label shift, where the paper's own experiments only apply mild augmentation.
- The one-round, low-dimensional ratio sharing suggests VRLS is compatible with secure aggregation or differential-privacy mechanisms applied to the shared marginal vector, though the paper does not analyze that composition.
- The Theorem 5.1 error bound suggests a practical diagnostic: monitor the entropy of the local predictor and the estimated ratio vector, because a saturated entropy signal would indicate that the regularization coefficient needs tuning.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper proposes VRLS, a label-shift ratio estimator that augments MLLS-style maximum-likelihood density-ratio estimation with Shannon entropy regularization of the predictor, and embeds VRLS in an importance-weighted ERM framework for multi-node distributed learning. The authors report that VRLS achieves lower mean-squared error than MLLS baselines on MNIST and CIFAR-10 under simulated and relaxed label shift, and that IW-ERM with VRLS improves average accuracy over FedAvg/FedBN/FedProx/SCAFFOLD in 5-, 100-, and 200-node experiments, approaching an upper bound using true ratios. The paper also states high-probability bounds on ratio estimation and convergence/communication guarantees for IW-ERM.
Significance. If the theoretical claims were fully established, the paper would make a useful contribution to label-shift estimation and distributed learning under distribution shift: the single-node MSE experiments are a clean, externally grounded comparison, and the 100- and 200-node experiments address a realistic-scale regime that is uncommon in this literature. The proposed method is simple and easy to reuse. However, the theoretical support is currently weaker than the abstract suggests: the main ratio-estimation bound does not quantify the effect of entropy regularization, and the distributed consistency proof targets a different objective from the one in the main text. The empirical distributed evaluation also omits the key ablation needed to attribute gains to VRLS rather than to importance weighting alone.
major comments (3)
- [Section 4, Proposition 4.1 and Appendix C.4] Proposition 4.1 claims that the displayed IW-ERM objective is consistent and converges to the minimizer of the overall true risk sum_k R_k, but the proof in Appendix C.4 establishes convergence only for the alternative objective in Eq. (19), whose weights are pte_1(y)/ptr_k(y) and whose limit is R_1(h_w), not sum_k R_k. These objectives agree only under a special choice of test marginals or if the target is node 1's risk. As written, the consistency of the actual IW-ERM objective used in the experiments is not proven; the proposition should be proved for the displayed objective or restated to match the appendix.
- [Theorem 5.1, Eq. (5)] The claimed high-probability bound on ratio estimation error contains an additive term (4L/(mu p_min)) E[||theta - theta_star||_2] that is never quantified, and the theorem gives no reason to expect this term to shrink when zeta > 0. The remark following the proof that regularization 'often encourages lower complexity functions' is informal, and the Rademacher term in Eq. (6) is defined over the joint class after regularized training, so it is not comparable with the zeta = 0 MLLS case. Thus Theorem 5.1 does not, as stated, support the paper's central claim that entropy regularization improves density-ratio estimation; a quantitative control of E[||theta - theta_star||] and a comparison with the MLLS bound are needed.
- [Section 6, Tables 2-4] The distributed experiments compare IW-ERM with VRLS ratios only against FedAvg, FedBN, FedProx, SCAFFOLD, and an upper bound using true ratios. There is no IW-ERM baseline using plain MLLS ratios under the same predictor training protocol. Because the single-node results show that MLLS is already a strong ratio estimator, the reported 20% gains over federated baselines may be attributable to importance weighting itself rather than to the proposed entropy regularization. Adding IW-ERM with MLLS ratios is necessary to isolate the contribution of VRLS in the distributed setting.
minor comments (6)
- [Section 5 / Appendix E] In the proof of Theorem 5.1, the notation for the empirical estimator is inconsistent: the main text defines hat r_{n_te}, while Appendix E uses hat r_n and also an undefined hat r_n in Eq. (30); please unify the notation and define all quantities before use.
- [Appendix E, Lemma E.7] Lemma E.7's statement says 'Let L(r, theta) be the population loss as defined in Lemma E.7', which is a self-reference; it should refer to Lemma E.3 or define L before the lemma.
- [Section 4] The proposition refers to 'equation IW-ERM' as if it were numbered, but the displayed objective later in the same section has no number; please add equation numbers and use them consistently in the text.
- [Appendix C] Appendix C says 'Without loss of generality and for simplicity, we set l = 1' and repeatedly states that the IW-ERM formulations focus on minimizing R_1, which conflicts with the main text's claim about the overall risk; please clarify the target risk in each scenario.
- [Section 6, Table 4] The 100-node and 200-node experiments report a single run with no variance or seed information, so it is difficult to assess the statistical significance of the differences in Table 4; please report multiple seeds or state the computational constraint explicitly.
- [Table 8 caption] The caption contains the typo 'Fasion MNIST'; it should be 'Fashion MNIST'.
Circularity Check
No significant circularity: the central VRLS claim is an empirical extension of MLLS with entropy regularization and is benchmarked against external baselines; self-citations are contextual, and theory gaps are evidential rather than circular.
full rationale
The paper does not reduce any claimed prediction to its inputs by construction. VRLS is defined as MLLS ratio estimation (Equation 3) using a predictor trained with an explicit Shannon-entropy regularizer (Equation 4); the estimator objective is the standard MLLS log-likelihood, and the regularization coefficient zeta is varied only in experiments, not fitted to the reported ratio estimates. Both single-node density-ratio experiments and multi-node IW-ERM experiments compare against external or oracle baselines (MLLS EM/L2, FedAvg, FedBN, FedProx, SCAFFOLD, and an upper bound using true ratios), so the reported gains are not forced by the method's definition. Theorem 5.1 is a high-probability bound on the regularized estimator's error following Garg et al.'s MLLS analysis; the added term (4L/(mu p_min))E[||theta-hat-theta*||] is left unbounded, and the remark that regularization 'often encourages lower complexity functions' is informal, so the theoretical support for the entropy regularizer's causal role is incomplete rather than circular. The consistency proof of Proposition 4.1 invokes standard external results (Shimodaira 2000; Sugiyama et al. 2007); the self-citation to Ramezani-Kebrya et al. (2023b) supplies prior Federated IW-ERM context and hyperparameters, not the load-bearing derivation. The appendix's proof for Equation (19) targets R1 rather than the Section 4 IW-ERM objective's overall risk, and Appendix I states the label-shift assumption is restrictive; these are correctness/scope issues, not instances of a result equaling its input. No circular step meeting the quoting standard was found.
Assumptions & free parameters
free parameters (1)
- Regularization coefficient zeta =
1
assumptions (5)
- domain assumption Class-conditional feature distribution is shared across nodes and train/test: p_te_k(x|y) = p_tr_k(x|y) = p(x|y) for all k
- domain assumption Assumption 5.2: E[f_{theta*}(x)f_{theta*}(x)^T] >= mu I_m for some mu > 0
- standard math Assumption 5.1: data and parameter space are bounded
- standard math Assumptions 5.3-5.5: smoothness, convexity/PL, sub-Gaussian gradient noise, unbiased gradients
- domain assumption Each node has access to unlabeled test data from its test distribution during ratio estimation
Cite this review
Pith. "Pith review of Addressing Label Shift in Distributed Learning via Entropy Regularization." pith.science (2026). https://pith.science/paper/MQBJMK2Y
@misc{pith2026250202544,
author = {Pith},
title = {Pith review of: Addressing Label Shift in Distributed Learning via Entropy Regularization},
year = {2026},
howpublished = {\url{https://pith.science/paper/MQBJMK2Y}},
note = {Machine review of arXiv:2502.02544}
}
read the original abstract
We address the challenge of minimizing true risk in multi-node distributed learning. 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 density ratio. VRLS incorporates Shannon entropy-based regularization and adjusts the density 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 density 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.
Figures
Figures from the paper (4 more)
Reference graph
Works this paper leans on
-
[1]
C.4 M ULTIPLE NODES Here we consider a general scenario with K nodes. We assume both intra-node and inter-node label shifts by the following IW-ERM: min hw∈H KX k=1 1 ntr k ntr kX i=1 pte 1 (ytr k,i) ptr k(ytr k,i) ℓ(hw(xtr k,i), ytr k,i), (19) This scenario is referred to as LS on multi . For the scenario without intra-node label shift, the IW-ERM in Equ...
work page 2000
-
[2]
are diagonal and the minimum eigenvalues satisfy λmin(At) ≥ ρ >0 for some ρ ∈ R+. 27 For nonconvex optimization with adaptive step-sizes, we 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 O...
work page 2023
-
[3]
Fl games: A federated learning framework for distribution shifts
Sharut Gupta, Kartik Ahuja, Mohammad Havaei, Niladri Chatterjee, and Yoshua Bengio. Fl games: A federated learning framework for distribution shifts. arXiv preprint arXiv:2205.11101,
-
[4]
Theorem F .3(Oracle Complexity of Proximal Operator for Composite Optimization)
for Gradient Flow-type transformation of IW-ERM with VRLS in the limit of infinitesimal step-size. Theorem F .3(Oracle Complexity of Proximal Operator for Composite Optimization). Let κ denote the condition number. Under Assumption F .3, suppose Gradient Flow-type transformation of Al- gorithm 2 with VRLS and Proximal Operator evolves in the limit of infi...
work page 2023
-
[6]
You-Wei Luo and Chuan-Xian Ren
doi: 10.1109/CVPR52729.2023.00361. You-Wei Luo and Chuan-Xian Ren. Generalized label shift correction via minimum uncer- tainty principle: Theory and algorithm. ArXiv, abs/2202.13043,
arXiv 2023
-
[10]
Antonio Torralba, Rob Fergus, and William T. Freeman. 80 million tiny images: A large data set for nonparametric object and scene recognition. IEEE Transactions on Pattern Analysis and Machine Intelligence, 30(11):1958–1970,
work page 1958
-
[11]
Flexifed: Personalized federated learning for edge clients with heterogeneous model architectures
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,
work page 2023
-
[12]
Fashion-MNIST: A novel image dataset for bench- marking machine learning algorithms
Han Xiao, Kashif Rasul, and Roland V ollgraf. Fashion-MNIST: A novel image dataset for bench- marking machine learning algorithms. arXiv preprint arXiv:1708.07747,
Show all 28 references
-
[13]
• Additional details of prior work of BBSE and MLLS are in Appendix B
15 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 is given in Appendix C. • General algorithmi...
2018
-
[14]
In this section, we overview complete related work
or 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 t...
2020
-
[15]
and SCAFFOLD (Karimireddy et al., 2020a). FedAvg and its variants such as (Huang et al., 2021; Karimireddy et al., 2020b) have been the subject of thorough investigation in optimization literature, exploring facets such as communication efficiency, node participation, and priv...
2022
-
[16]
This approach is straightforward and has been proven consistent, even when the predictor is not calibrated
designate a discrete latent space z and introduce a confusion matrix-based estimation method to compute the ratio w 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 ...
2002
-
[17]
It also provides a consistency guarantee with a canonically calibrated predictor
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 ...
2019
-
[20]
Concerning IPMs, while MMD is reliant on a kernel function, it can suffer from the curse of dimensionality when faced with high-dimensional data
employs f-divergence, notably the KL divergence, which is not a metric as it doesn’t satisfy the triangular inequality, and requires distribution P to be absolutely continuous with respect to Q. Concerning IPMs, while MMD is reliant on a kernel function, it can suffer from the...
2006
-
[23]
The stochastic gradient g(w) = e∇wℓ(hw) is unbiased, i.e., E[g(w)] = ∇wℓ(hw) for any w ∈ Wwith bounded variance E[∥g(w) − ∇wℓ(hw)∥2 2]. For nonconvex optimization with PL condition and communication compression, we establish convergence and communication guarantees for IW-ERM ...
2021
-
[24]
Adaptive matrices At constructed as in (Wu et al., 2023, Algorithm
2023
-
[26]
The stochastic gradients g(w) = e∇wℓ(hw) is unbiased. For composite optimization with strongly convex and smooth functions and proximal operator, we establish an upper bound on oracle complexity to achieve ϵ error on the Lyapunov function defined as in (Hu & Huang, 2023, Section
2023
-
[1951]
and Total Variation distance, operate 17 on ratios P/Q and do not satisfy the triangular inequality. Interconnections and variations between these families are explored in studies like (f, Γ)-Divergences (Birrell et al., 2022), which interpolate between f-divergences and IPMs,...
2022
-
[1998]
Federated optimization in heterogeneous networks
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, MLSy...
2020
-
[2000]
Alexandari, and Anshul Kundaje
Avanti Shrikumar, Amr M. Alexandari, and Anshul Kundaje. Adapting to label shift with bias- corrected calibration. arXiv preprint arXiv:1901.06852v5,
1901 arXiv
-
[2012]
In contrast, f- divergences, such as KL-divergence (Kullback & Leibler,
and Wasserstein distance (Villani, 2009), focus on distribution differences P − Q. In contrast, f- divergences, such as KL-divergence (Kullback & Leibler,
2009
-
[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...
2023
-
[2018]
log( mX c=1 ptr(z, y= c)rc) # = Ete
and MLLS (Garg et al., 2020). Our objective is to estimate the ratio pte(y)/ptr(y). We consider a scenario with m possible label classes, where y = c for c ∈ [m]. Let r⋆ = [r⋆ 1, . . . , r⋆ m]⊤ represent the true ratios, with each r⋆ c defined as r⋆ c = pte(y=c) ptr(y=c) (Garg...
2018
-
[2019]
Mitigating data heterogeneity in federated learning with data augmentation
Artur Back de Luca, Guojun Zhang, Xi Chen, and Yaoliang Yu. Mitigating data heterogeneity in federated learning with data augmentation. arXiv preprint arXiv:2206.09979,
-
[2020]
Saurabh Garg, Sivaraman Balakrishnan, and Zachary C. Lipton. Domain adaptation under open set label shift. arXiv preprint arXiv:2207.13048,
-
[2021]
Distributed extra-gradient with optimal complexity and communication guarantees
Ali Ramezani-Kebrya, Kimon Antonakopoulos, Igor Krawczuk, Justin Deschenaux, and V olkan Cevher. Distributed extra-gradient with optimal complexity and communication guarantees. In International Conference on Learning Representations (ICLR), 2023a. Ali Ramezani-Kebrya, Fanghui...
-
[2023]
Gradma: A gradient-memory-based ac- celerated federated learning with alleviated catastrophic forgetting
Kangyang Luo, Xiang Li, Yunshi Lan, and Ming Gao. Gradma: A gradient-memory-based 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
-
[2024]
Regularizing neural networks by penalizing confident output distributions
13 Gabriel Pereyra, George Tucker, Jan Chorowski,Łukasz Kaiser, and Geoffrey Hinton. Regularizing neural networks by penalizing confident output distributions. arXiv preprint arXiv:1701.06548,
Reviewed August 9, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.