# Overcoming label shift with target-aware federated learning

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 39 of its 52 references that are in the map.
- First posted on arXiv 6 Nov 2024.

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

- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [Advances and Open Problems in Federated Learning](https://mlsanity.com/paper/1912.04977/) (2019)
- [Domain-Adversarial Training of Neural Networks](https://mlsanity.com/paper/1505.07818/) (2015)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Federated Learning: Challenges, Methods, and Future Directions](https://mlsanity.com/paper/1908.07873/) (2019)
- [WILDS: A Benchmark of in-the-Wild Distribution Shifts](https://mlsanity.com/paper/2012.07421/) (2020)
- [Federated Learning with Non-IID Data](https://mlsanity.com/paper/1806.00582/) (2018)
- [In Search of Lost Domain Generalization](https://mlsanity.com/paper/2007.01434/) (2020)
- [Deeper, Broader and Artier Domain Generalization](https://mlsanity.com/paper/1710.03077/) (2017)
- [LEAF: A Benchmark for Federated Settings](https://mlsanity.com/paper/1812.01097/) (2018)
- [Ditto: Fair and Robust Federated Learning Through Personalization](https://mlsanity.com/paper/2012.04221/) (2020)
- [Federated Learning on Non-IID Data Silos: An Experimental Study](https://mlsanity.com/paper/2102.02079/) (2021)
- [An Efficient Framework for Clustered Federated Learning](https://mlsanity.com/paper/2006.04088/) (2020)
- [Clustered Federated Learning: Model-Agnostic Distributed Multi-Task Optimization under Privacy Constraints](https://mlsanity.com/paper/1910.01991/) (2019)
- [Differentially Private Federated Learning: A Client Level Perspective](https://mlsanity.com/paper/1712.07557/) (2017)
- [Agnostic Federated Learning](https://mlsanity.com/paper/1902.00146/) (2019)
- [Bayesian Nonparametric Federated Learning of Neural Networks](https://mlsanity.com/paper/1905.12022/) (2019)
- [Detecting and Correcting for Label Shift with Black Box Predictors](https://mlsanity.com/paper/1802.03916/) (2018)
- [FedDG: Federated Domain Generalization on Medical Image Segmentation via Episodic Learning in Continuous Frequency Space](https://mlsanity.com/paper/2103.06030/) (2021)
- [Improving Federated Learning Personalization via Model Agnostic Meta Learning](https://mlsanity.com/paper/1909.12488/) (2019)
- [Federated Meta-Learning with Fast Convergence and Efficient Communication](https://mlsanity.com/paper/1802.07876/) (2018)
- [Minibatch vs Local SGD for Heterogeneous Distributed Learning](https://mlsanity.com/paper/2006.04735/) (2020)
- [Generalized Federated Learning via Sharpness Aware Minimization](https://mlsanity.com/paper/2206.02618/) (2022)
- [Federated Learning with Label Distribution Skew via Logits Calibration](https://mlsanity.com/paper/2209.00189/) (2022)
- [Improving Generalization in Federated Learning by Seeking Flat Minima](https://mlsanity.com/paper/2203.11834/) (2022)
- [Differentially Private Secure Multi-Party Computation for Federated Learning in Financial Applications](https://mlsanity.com/paper/2010.05867/) (2020)
- [Revisiting Weighted Aggregation in Federated Learning with Neural Networks](https://mlsanity.com/paper/2302.10911/) (2023)
- [FedAvg with Fine Tuning: Local Updates Lead to Representation Learning](https://mlsanity.com/paper/2205.13692/) (2022)

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

- [FedCF: Fair Federated Conformal Prediction](https://mlsanity.com/paper/2509.22907/) (2025)
- [Reviving Stale Updates: Data-Free Knowledge Distillation for Asynchronous Federated Learning](https://mlsanity.com/paper/2511.00655/) (2025)

## Abstract (arXiv)

Federated learning enables multiple actors to collaboratively train models without sharing private data. Existing algorithms are successful and well-justified in this task when the intended target domain, where the trained model will be used, shares data distribution with the aggregate of clients, but this is often violated in practice. A common reason is label shift -- that the label distributions differ between clients and the target domain. We demonstrate empirically that this can significantly degrade performance. To address this problem, we propose FedPALS, a principled and practical model aggregation scheme that adapts to label shifts to improve performance in the target domain by leveraging knowledge of label distributions at the central server. Our approach ensures unbiased updates under federated stochastic gradient descent which yields robust generalization across clients with diverse, label-shifted data. Extensive experiments on image classification tasks demonstrate that FedPALS consistently outperforms baselines by aligning model aggregation with the target domain. Our findings reveal that conventional federated learning methods suffer severely in cases of extreme label sparsity on clients, highlighting the critical need for target-aware aggregation as offered by FedPALS.