# FRAug: Tackling Federated Learning with Non-IID Features via Representation Augmentation

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

## Key facts

- Cited by 24 papers in the map (data through 14 Oct 2026).
- Builds on 57 of its 86 references that are in the map.
- First posted on arXiv 30 May 2022.

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

- [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)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [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)
- [Unsupervised Domain Adaptation by Backpropagation](https://mlsanity.com/paper/1409.7495/) (2014)
- [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)
- [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)
- [Temporal Ensembling for Semi-Supervised Learning](https://mlsanity.com/paper/1610.02242/) (2016)
- [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)
- [Domain Generalization: A Survey](https://mlsanity.com/paper/2103.02503/) (2021)
- [Ensemble Distillation for Robust Model Fusion in Federated Learning](https://mlsanity.com/paper/2006.07242/) (2020)
- [Ditto: Fair and Robust Federated Learning Through Personalization](https://mlsanity.com/paper/2012.04221/) (2020)
- [Personalized Federated Learning with Moreau Envelopes](https://mlsanity.com/paper/2006.08848/) (2020)
- [Federated Learning for Mobile Keyboard Prediction](https://mlsanity.com/paper/1811.03604/) (2018)
- [FedBN: Federated Learning on Non-IID Features via Local Batch Normalization](https://mlsanity.com/paper/2102.07623/) (2021)
- [Domain Generalization with MixStyle](https://mlsanity.com/paper/2104.02008/) (2021)
- [Federated Learning with Personalization Layers](https://mlsanity.com/paper/1912.00818/) (2019)
- [Exploiting Shared Representations for Personalized Federated Learning](https://mlsanity.com/paper/2102.07078/) (2021)
- [FedMD: Heterogenous Federated Learning via Model Distillation](https://mlsanity.com/paper/1910.03581/) (2019)
- [Domain Generalization via Model-Agnostic Learning of Semantic Features](https://mlsanity.com/paper/1910.13580/) (2019)
- [Data-Free Knowledge Distillation for Heterogeneous Federated Learning](https://mlsanity.com/paper/2105.10056/) (2021)
- [A Survey of Unsupervised Deep Domain Adaptation](https://mlsanity.com/paper/1812.02849/) (2018)
- [FedProto: Federated Prototype Learning across Heterogeneous Clients](https://mlsanity.com/paper/2105.00243/) (2021)
- [Federated Learning on Non-IID Data: A Survey](https://mlsanity.com/paper/2106.06843/) (2021)
- [See through Gradients: Image Batch Recovery via GradInversion](https://mlsanity.com/paper/2104.07586/) (2021)
- [Learning to Generate Novel Domains for Domain Generalization](https://mlsanity.com/paper/2007.03304/) (2020)

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

- [Visual Prompt Based Personalized Federated Learning](https://mlsanity.com/paper/2303.08678/) (2023)
- [FedBiP: Heterogeneous One-Shot Federated Learning with Personalized Latent Diffusion Models](https://mlsanity.com/paper/2410.04810/) (2024)
- [Non-IID data in Federated Learning: A Survey with Taxonomy, Metrics, Methods, Frameworks and Future Directions](https://mlsanity.com/paper/2411.12377/) (2024)
- [Advances in Robust Federated Learning: A Survey with Heterogeneity Considerations](https://mlsanity.com/paper/2405.09839/) (2024)
- [A Comprehensive Survey of Federated Transfer Learning: Challenges, Methods and Applications](https://mlsanity.com/paper/2403.01387/) (2024)
- [Fake It Till Make It: Federated Learning with Consensus-Oriented Generation](https://mlsanity.com/paper/2312.05966/) (2023)
- [Bridging Data Islands: Geographic Heterogeneity-Aware Federated Learning for Collaborative Remote Sensing Semantic Segmentation](https://mlsanity.com/paper/2404.09292/) (2024)
- [FedSC: Federated Learning with Semantic-Aware Collaboration](https://mlsanity.com/paper/2506.21012/) (2025)
- [Federated Hybrid Training and Self-Adversarial Distillation: Towards Robust Edge Networks](https://mlsanity.com/paper/2412.19354/) (2024)
- [FedDAG: Federated Domain Adversarial Generation Towards Generalizable Medical Image Analysis](https://mlsanity.com/paper/2501.13967/) (2025)
- [FedPall: Prototype-based Adversarial and Collaborative Learning for Federated Learning with Feature Drift](https://mlsanity.com/paper/2507.04781/) (2025)
- [FedPop: Federated Population-based Hyperparameter Tuning](https://mlsanity.com/paper/2308.08634/) (2023)
- [CAPT: Class-Aware Prompt Tuning for Federated Long-Tailed Learning with Vision-Language Model](https://mlsanity.com/paper/2503.06993/) (2025)
- [Towards Unbiased Federated Graph Learning: Label and Topology Perspectives](https://mlsanity.com/paper/2504.09963/) (2025)
- [Towards Unified Modeling in Federated Multi-Task Learning via Subspace Decoupling](https://mlsanity.com/paper/2505.24185/) (2025)
- [Distribution Alignment for One-Shot Federated Learning via Optimal Transport](https://mlsanity.com/paper/2606.16655/) (2026)
- [Heterogeneous Federated Learning via Personalized Generative Networks](https://mlsanity.com/paper/2308.13265/) (2023)
- [Stratify: Rethinking Federated Learning for Non-IID Data through Balanced Sampling](https://mlsanity.com/paper/2504.13462/) (2025)
- [Adaptive Latent-Space Constraints in Personalized Federated Learning](https://mlsanity.com/paper/2505.07525/) (2025)
- [FedMP: Tackling Medical Feature Heterogeneity in Federated Learning from a Manifold Perspective](https://mlsanity.com/paper/2508.09174/) (2025)
- [Personalized federated prototype learning in mixed heterogeneous data scenarios](https://mlsanity.com/paper/2510.03726/) (2025)
- [CEPerFed: Communication-Efficient Personalized Federated Learning for Multi-Pulse MRI Classification](https://mlsanity.com/paper/2510.17584/) (2025)
- [FedCausal-Dyn: A Causal-Dynamic Paradigm for Federated Learning under Dynamic Feature Drift](https://mlsanity.com/paper/2607.09695/) (2026)
- [CRIP: Channel Level Representation Injection for Personalized One-Shot Federated Learning](https://mlsanity.com/paper/2608.02222/) (2026)

## Abstract (arXiv)

Federated Learning (FL) is a decentralized learning paradigm, in which multiple clients collaboratively train deep learning models without centralizing their local data, and hence preserve data privacy. Real-world applications usually involve a distribution shift across the datasets of the different clients, which hurts the generalization ability of the clients to unseen samples from their respective data distributions. In this work, we address the recently proposed feature shift problem where the clients have different feature distributions, while the label distribution is the same. We propose Federated Representation Augmentation (FRAug) to tackle this practical and challenging problem. Our approach generates synthetic client-specific samples in the embedding space to augment the usually small client datasets. For that, we train a shared generative model to fuse the clients knowledge learned from their different feature distributions. This generator synthesizes client-agnostic embeddings, which are then locally transformed into client-specific embeddings by Representation Transformation Networks (RTNets). By transferring knowledge across the clients, the generated embeddings act as a regularizer for the client models and reduce overfitting to the local original datasets, hence improving generalization. Our empirical evaluation on public benchmarks and a real-world medical dataset demonstrates the effectiveness of the proposed method, which substantially outperforms the current state-of-the-art FL methods for non-IID features, including PartialFed and FedBN.