# Stabilizing and Improving Federated Learning with Non-IID Data and Client Dropout

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 41 of its 62 references that are in the map.
- First posted on arXiv 11 Mar 2023.

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

- [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)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [Federated Learning: Challenges, Methods, and Future Directions](https://mlsanity.com/paper/1908.07873/) (2019)
- [Federated Learning with Non-IID Data](https://mlsanity.com/paper/1806.00582/) (2018)
- [Model-Contrastive Federated Learning](https://mlsanity.com/paper/2103.16257/) (2021)
- [Ensemble Distillation for Robust Model Fusion in Federated Learning](https://mlsanity.com/paper/2006.07242/) (2020)
- [Personalized Federated Learning with Moreau Envelopes](https://mlsanity.com/paper/2006.08848/) (2020)
- [Federated Learning on Non-IID Data Silos: An Experimental Study](https://mlsanity.com/paper/2102.02079/) (2021)
- [Federated Learning Based on Dynamic Regularization](https://mlsanity.com/paper/2111.04263/) (2021)
- [Exploiting Shared Representations for Personalized Federated Learning](https://mlsanity.com/paper/2102.07078/) (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)
- [Adaptive Federated Learning in Resource Constrained Edge Computing Systems](https://mlsanity.com/paper/1804.05271/) (2018)
- [Balanced Meta-Softmax for Long-Tailed Visual Recognition](https://mlsanity.com/paper/2007.10740/) (2020)
- [Data-Free Knowledge Distillation for Heterogeneous Federated Learning](https://mlsanity.com/paper/2105.10056/) (2021)
- [Three Approaches for Personalization with Applications to Federated Learning](https://mlsanity.com/paper/2002.10619/) (2020)
- [No Fear of Heterogeneity: Classifier Calibration for Federated Learning with Non-IID Data](https://mlsanity.com/paper/2106.05001/) (2021)
- [FedBABU: Towards Enhanced Representation for Federated Image Classification](https://mlsanity.com/paper/2106.06042/) (2021)
- [FedBE: Making Bayesian Model Ensemble Applicable to Federated Learning](https://mlsanity.com/paper/2009.01974/) (2020)
- [Fine-tuning Global Model via Data-Free Knowledge Distillation for Non-IID Federated Learning](https://mlsanity.com/paper/2203.09249/) (2022)
- [Inverting Visual Representations with Convolutional Networks](https://mlsanity.com/paper/1506.02753/) (2015)
- [Survey of Personalization Techniques for Federated Learning](https://mlsanity.com/paper/2003.08673/) (2020)
- [FedPD: A Federated Learning Framework with Optimal Rates and Adaptivity to Non-IID Data](https://mlsanity.com/paper/2005.11418/) (2020)
- [Federated Visual Classification with Real-World Data Distribution](https://mlsanity.com/paper/2003.08082/) (2020)
- [Federated Learning with Label Distribution Skew via Logits Calibration](https://mlsanity.com/paper/2209.00189/) (2022)
- [FedMix: Approximation of Mixup under Mean Augmented Federated Learning](https://mlsanity.com/paper/2107.00233/) (2021)
- [Clustered Sampling: Low-Variance and Improved Representativity for Clients Selection in Federated Learning](https://mlsanity.com/paper/2105.05883/) (2021)
- [SAFA: a Semi-Asynchronous Protocol for Fast Federated Learning with Low Overhead](https://mlsanity.com/paper/1910.01355/) (2019)

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

- [Federated Neural Nonparametric Point Processes](https://mlsanity.com/paper/2410.05637/) (2024)
- [Overcoming label shift with target-aware federated learning](https://mlsanity.com/paper/2411.03799/) (2024)
- [The More is not the Merrier: Investigating the Effect of Client Size on Federated Learning](https://mlsanity.com/paper/2504.08198/) (2025)

## Abstract (arXiv)

The label distribution skew induced data heterogeniety has been shown to be a significant obstacle that limits the model performance in federated learning, which is particularly developed for collaborative model training over decentralized data sources while preserving user privacy. This challenge could be more serious when the participating clients are in unstable circumstances and dropout frequently. Previous work and our empirical observations demonstrate that the classifier head for classification task is more sensitive to label skew and the unstable performance of FedAvg mainly lies in the imbalanced training samples across different classes. The biased classifier head will also impact the learning of feature representations. Therefore, maintaining a balanced classifier head is of significant importance for building a better global model. To this end, we propose a simple yet effective framework by introducing a prior-calibrated softmax function for computing the cross-entropy loss and a prototype-based feature augmentation scheme to re-balance the local training, which are lightweight for edge devices and can facilitate the global model aggregation. The improved model performance over existing baselines in the presence of non-IID data and client dropout is demonstrated by conducting extensive experiments on benchmark classification tasks.