# Addressing Label Shift in Distributed Learning via Entropy Regularization

arXiv: https://arxiv.org/abs/2502.02544 · HTML: https://mlsanity.com/paper/2502.02544/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 43 of its 74 references that are in the map.
- First posted on arXiv 4 Feb 2025.

## Builds on (most cited first, 30 of 43)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [Federated Multi-Task Learning](https://mlsanity.com/paper/1705.10467/) (2017)
- [Ditto: Fair and Robust Federated Learning Through Personalization](https://mlsanity.com/paper/2012.04221/) (2020)
- [FedBN: Federated Learning on Non-IID Features via Local Batch Normalization](https://mlsanity.com/paper/2102.07623/) (2021)
- [Regularizing Neural Networks by Penalizing Confident Output Distributions](https://mlsanity.com/paper/1701.06548/) (2017)
- [Detecting and Correcting for Label Shift with Black Box Predictors](https://mlsanity.com/paper/1802.03916/) (2018)
- [Beyond temperature scaling: Obtaining well-calibrated multiclass probabilities with Dirichlet calibration](https://mlsanity.com/paper/1910.12656/) (2019)
- [Personalized Cross-Silo Federated Learning on Non-IID Data](https://mlsanity.com/paper/2007.03797/) (2020)
- [What is the Effect of Importance Weighting in Deep Learning?](https://mlsanity.com/paper/1812.03372/) (2018)
- [Failing Loudly: An Empirical Study of Methods for Detecting Dataset Shift](https://mlsanity.com/paper/1810.11953/) (2018)
- [Do CIFAR-10 Classifiers Generalize to CIFAR-10?](https://mlsanity.com/paper/1806.00451/) (2018)
- [Adaptive Gradient-Based Meta-Learning Methods](https://mlsanity.com/paper/1906.02717/) (2019)
- [Heterogeneous Federated Learning: State-of-the-art and Research Challenges](https://mlsanity.com/paper/2307.10616/) (2023)
- [Federated Learning with Compression: Unified Analysis and Sharp Guarantees](https://mlsanity.com/paper/2007.01154/) (2020)
- [Is Local SGD Better than Minibatch SGD?](https://mlsanity.com/paper/2002.07839/) (2020)
- [Regularized Learning for Domain Adaptation under Label Shifts](https://mlsanity.com/paper/1903.09734/) (2019)
- [A Unified View of Label Shift Estimation](https://mlsanity.com/paper/2003.07554/) (2020)
- [Rethinking Importance Weighting for Deep Learning under Distribution Shift](https://mlsanity.com/paper/2006.04662/) (2020)
- [High Probability Convergence of Stochastic Gradient Methods](https://mlsanity.com/paper/2302.14843/) (2023)
- [Adaptive Gradient Quantization for Data-Parallel SGD](https://mlsanity.com/paper/2010.12460/) (2020)
- [Sharp Bounds for Federated Averaging (Local SGD) and Continuous Perspective](https://mlsanity.com/paper/2111.03741/) (2021)
- [Faster Adaptive Federated Learning](https://mlsanity.com/paper/2212.00974/) (2022)
- [Multi-Level Branched Regularization for Federated Learning](https://mlsanity.com/paper/2207.06936/) (2022)
- [Domain Adaptation under Open Set Label Shift](https://mlsanity.com/paper/2207.13048/) (2022)
- [RLSbench: Domain Adaptation Under Relaxed Label Shift](https://mlsanity.com/paper/2302.03020/) (2023)
- [Optimal Bounds between f-Divergences and Integral Probability Metrics](https://mlsanity.com/paper/2006.05973/) (2020)

## Built on by (most cited first, 1 of 1)

- [From matrix inversion to constraints: provably tighter confidence regions for importance weights in label shift](https://mlsanity.com/paper/2609.14802/) (2026)

## Abstract (arXiv)

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.