# Mitigating Data Heterogeneity in Federated Learning with Data Augmentation

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

## Key facts

- Cited by 21 papers in the map (data through 14 Oct 2026).
- Builds on 46 of its 69 references that are in the map.
- First posted on arXiv 20 Jun 2022.

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

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [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)
- [Domain Randomization for Transferring Deep Neural Networks from Simulation to the Real World](https://mlsanity.com/paper/1703.06907/) (2017)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [On the Convergence of FedAvg on Non-IID Data](https://mlsanity.com/paper/1907.02189/) (2019)
- [Deep Hashing Network for Unsupervised Domain Adaptation](https://mlsanity.com/paper/1706.07522/) (2017)
- [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)
- [Tackling the Objective Inconsistency Problem in Heterogeneous Federated Optimization](https://mlsanity.com/paper/2007.07481/) (2020)
- [Measuring the Effects of Non-Identical Data Distribution for Federated Visual Classification](https://mlsanity.com/paper/1909.06335/) (2019)
- [Generalizing to Unseen Domains: A Survey on Domain Generalization](https://mlsanity.com/paper/2103.03097/) (2021)
- [Out-of-Distribution Generalization via Risk Extrapolation (REx)](https://mlsanity.com/paper/2003.00688/) (2020)
- [Domain Generalization via Invariant Feature Representation](https://mlsanity.com/paper/1301.2115/) (2013)
- [Federated Learning with Matched Averaging](https://mlsanity.com/paper/2002.06440/) (2020)
- [Causal inference using invariant prediction: identification and confidence intervals](https://mlsanity.com/paper/1501.01332/) (2015)
- [Agnostic Federated Learning](https://mlsanity.com/paper/1902.00146/) (2019)
- [Generalizing to Unseen Domains via Adversarial Data Augmentation](https://mlsanity.com/paper/1805.12018/) (2018)
- [Self-Supervised Learning with Data Augmentations Provably Isolates Content from Style](https://mlsanity.com/paper/2106.04619/) (2021)
- [The Non-IID Data Quagmire of Decentralized Machine Learning](https://mlsanity.com/paper/1910.00189/) (2019)
- [Learning to Learn Single Domain Generalization](https://mlsanity.com/paper/2003.13216/) (2020)
- [Data Augmentation Generative Adversarial Networks](https://mlsanity.com/paper/1711.04340/) (2017)

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

- [When Foundation Model Meets Federated Learning: Motivations, Challenges, and Future Directions](https://mlsanity.com/paper/2306.15546/) (2023)
- [Federated Domain Generalization: A Survey](https://mlsanity.com/paper/2306.01334/) (2023)
- [Federated Generative Learning with Foundation Models](https://mlsanity.com/paper/2306.16064/) (2023)
- [A Review of the Role of Causality in Developing Trustworthy AI Systems](https://mlsanity.com/paper/2302.06975/) (2023)
- [Multimodal Federated Learning in Healthcare: a Review](https://mlsanity.com/paper/2310.09650/) (2023)
- [Non-IID data in Federated Learning: A Survey with Taxonomy, Metrics, Methods, Frameworks and Future Directions](https://mlsanity.com/paper/2411.12377/) (2024)
- [Towards Achieving Near-optimal Utility for Privacy-Preserving Federated Learning via Data Generation and Parameter Distortion](https://mlsanity.com/paper/2305.04288/) (2023)
- [Federated Learning under Covariate Shifts with Generalization Guarantees](https://mlsanity.com/paper/2306.05325/) (2023)
- [Federated Deconfounding and Debiasing Learning for Out-of-Distribution Generalization](https://mlsanity.com/paper/2505.04979/) (2025)
- [Federated attention consistent learning models for prostate cancer diagnosis and Gleason grading](https://mlsanity.com/paper/2302.06089/) (2023)
- [Topology Learning for Heterogeneous Decentralized Federated Learning over Unreliable D2D Networks](https://mlsanity.com/paper/2312.13611/) (2023)
- [FedBM: Stealing Knowledge from Pre-trained Language Models for Heterogeneous Federated Learning](https://mlsanity.com/paper/2502.16832/) (2025)
- [FLAG: Fast Label-Adaptive Aggregation for Multi-label Classification in Federated Learning](https://mlsanity.com/paper/2302.13571/) (2023)
- [Addressing Label Shift in Distributed Learning via Entropy Regularization](https://mlsanity.com/paper/2502.02544/) (2025)
- [FedPhD: Federated Pruning with Hierarchical Learning of Diffusion Models](https://mlsanity.com/paper/2507.06449/) (2025)
- [FedVeca: Federated Vectorized Averaging on Non-IID Data with Adaptive Bi-directional Global Objective](https://mlsanity.com/paper/2209.13803/) (2022)
- [Global Layers: Non-IID Tabular Federated Learning](https://mlsanity.com/paper/2305.19290/) (2023)
- [Heterogeneous Federated Learning via Personalized Generative Networks](https://mlsanity.com/paper/2308.13265/) (2023)
- [Accelerating Federated Learning by Selecting Beneficial Herd of Local Gradients](https://mlsanity.com/paper/2403.16557/) (2024)
- [Topology-Aware Knowledge Propagation in Decentralized Learning](https://mlsanity.com/paper/2505.11760/) (2025)
- [FedASAP: Activation Statistics-driven Structured Adaptive Pruning for Efficient Personalized Federated Learning for Lesion Segmentation on brain MRI](https://mlsanity.com/paper/2609.19042/) (2026)

## Abstract (arXiv)

Federated Learning (FL) is a prominent framework that enables training a centralized model while securing user privacy by fusing local, decentralized models. In this setting, one major obstacle is data heterogeneity, i.e., each client having non-identically and independently distributed (non-IID) data. This is analogous to the context of Domain Generalization (DG), where each client can be treated as a different domain. However, while many approaches in DG tackle data heterogeneity from the algorithmic perspective, recent evidence suggests that data augmentation can induce equal or greater performance. Motivated by this connection, we present federated versions of popular DG algorithms, and show that by applying appropriate data augmentation, we can mitigate data heterogeneity in the federated setting, and obtain higher accuracy on unseen clients. Equipped with data augmentation, we can achieve state-of-the-art performance using even the most basic Federated Averaging algorithm, with much sparser communication.