# Rethinking Architecture Design for Tackling Data Heterogeneity in Federated Learning

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

## Key facts

- Cited by 111 papers in the map (data through 14 Oct 2026).
- Builds on 53 of its 77 references that are in the map.
- First posted on arXiv 10 Jun 2021.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Emerging Properties in Self-Supervised Vision Transformers](https://mlsanity.com/paper/2104.14294/) (2021)
- [EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks](https://mlsanity.com/paper/1905.11946/) (2019)
- [Overcoming catastrophic forgetting in neural networks](https://mlsanity.com/paper/1612.00796/) (2016)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [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)
- [UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction](https://mlsanity.com/paper/1802.03426/) (2018)
- [Extracting Training Data from Large Language Models](https://mlsanity.com/paper/2012.07805/) (2020)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [Federated Learning: Challenges, Methods, and Future Directions](https://mlsanity.com/paper/1908.07873/) (2019)
- [ImageNet-trained CNNs are biased towards texture; increasing shape bias improves accuracy and robustness](https://mlsanity.com/paper/1811.12231/) (2018)
- [Continual Learning with Deep Generative Replay](https://mlsanity.com/paper/1705.08690/) (2017)
- [Federated Learning with Non-IID Data](https://mlsanity.com/paper/1806.00582/) (2018)
- [Learning Robust Global Representations by Penalizing Local Predictive Power](https://mlsanity.com/paper/1905.13549/) (2019)
- [Measuring the Effects of Non-Identical Data Distribution for Federated Visual Classification](https://mlsanity.com/paper/1909.06335/) (2019)
- [LEAF: A Benchmark for Federated Settings](https://mlsanity.com/paper/1812.01097/) (2018)
- [Image Transformer](https://mlsanity.com/paper/1802.05751/) (2018)
- [Model-Contrastive Federated Learning](https://mlsanity.com/paper/2103.16257/) (2021)
- [Inverting Gradients -- How easy is it to break privacy in federated learning?](https://mlsanity.com/paper/2003.14053/) (2020)
- [Ensemble Distillation for Robust Model Fusion in Federated Learning](https://mlsanity.com/paper/2006.07242/) (2020)
- [Overcoming catastrophic forgetting with hard attention to the task](https://mlsanity.com/paper/1801.01423/) (2018)
- [Multimodal Transformer for Unaligned Multimodal Language Sequences](https://mlsanity.com/paper/1906.00295/) (2019)
- [The Future of Digital Health with Federated Learning](https://mlsanity.com/paper/2003.08119/) (2020)

## Built on by (most cited first, 30 of 111)

- [Heterogeneous Federated Learning: State-of-the-art and Research Challenges](https://mlsanity.com/paper/2307.10616/) (2023)
- [Federated Learning from Pre-Trained Models: A Contrastive Learning Approach](https://mlsanity.com/paper/2209.10083/) (2022)
- [PromptFL: Let Federated Participants Cooperatively Learn Prompts Instead of Models -- Federated Learning in Age of Foundation Model](https://mlsanity.com/paper/2208.11625/) (2022)
- [Federated Learning for Generalization, Robustness, Fairness: A Survey and Benchmark](https://mlsanity.com/paper/2311.06750/) (2023)
- [Selective Aggregation for Low-Rank Adaptation in Federated Learning](https://mlsanity.com/paper/2410.01463/) (2024)
- [When Federated Learning Meets Pre-trained Language Models' Parameter-Efficient Tuning Methods](https://mlsanity.com/paper/2212.10025/) (2022)
- [On the Importance and Applicability of Pre-Training for Federated Learning](https://mlsanity.com/paper/2206.11488/) (2022)
- [A Comprehensive Survey of Forgetting in Deep Learning Beyond Continual Learning](https://mlsanity.com/paper/2307.09218/) (2023)
- [When Adversarial Training Meets Vision Transformers: Recipes from Training to Architecture](https://mlsanity.com/paper/2210.07540/) (2022)
- [Efficient Model Personalization in Federated Learning via Client-Specific Prompt Generation](https://mlsanity.com/paper/2308.15367/) (2023)
- [No One Left Behind: Real-World Federated Class-Incremental Learning](https://mlsanity.com/paper/2302.00903/) (2023)
- [Fair Federated Medical Image Segmentation via Client Contribution Estimation](https://mlsanity.com/paper/2303.16520/) (2023)
- [Label-Efficient Self-Supervised Federated Learning for Tackling Data Heterogeneity in Medical Imaging](https://mlsanity.com/paper/2205.08576/) (2022)
- [FedTP: Federated Learning by Transformer Personalization](https://mlsanity.com/paper/2211.01572/) (2022)
- [Cross-Silo Prototypical Calibration for Federated Learning with Non-IID Data](https://mlsanity.com/paper/2308.03457/) (2023)
- [FedTune: A Deep Dive into Efficient Federated Fine-Tuning with Pre-trained Transformers](https://mlsanity.com/paper/2211.08025/) (2022)
- [Why Batch Normalization Damage Federated Learning on Non-IID Data?](https://mlsanity.com/paper/2301.02982/) (2023)
- [Federated Learning for Computer Vision](https://mlsanity.com/paper/2308.13558/) (2023)
- [FedPerfix: Towards Partial Model Personalization of Vision Transformers in Federated Learning](https://mlsanity.com/paper/2308.09160/) (2023)
- [Unlocking the Potential of Prompt-Tuning in Bridging Generalized and Personalized Federated Learning](https://mlsanity.com/paper/2310.18285/) (2023)
- [SplitAVG: A heterogeneity-aware federated deep learning method for medical imaging](https://mlsanity.com/paper/2107.02375/) (2021)
- [PILoRA: Prototype Guided Incremental LoRA for Federated Class-Incremental Learning](https://mlsanity.com/paper/2401.02094/) (2024)
- [Multimodal Federated Learning with Missing Modality via Prototype Mask and Contrast](https://mlsanity.com/paper/2312.13508/) (2023)
- [FedMM: Federated Multi-Modal Learning with Modality Heterogeneity in Computational Pathology](https://mlsanity.com/paper/2402.15858/) (2024)
- [An Upload-Efficient Scheme for Transferring Knowledge From a Server-Side Pre-trained Generator to Clients in Heterogeneous Federated Learning](https://mlsanity.com/paper/2403.15760/) (2024)
- [FedBEVT: Federated Learning Bird's Eye View Perception Transformer in Road Traffic Systems](https://mlsanity.com/paper/2304.01534/) (2023)
- [Non-IID data in Federated Learning: A Survey with Taxonomy, Metrics, Methods, Frameworks and Future Directions](https://mlsanity.com/paper/2411.12377/) (2024)
- [Federated Learning Attacks and Defenses: A Survey](https://mlsanity.com/paper/2211.14952/) (2022)
- [On the Efficacy of Differentially Private Few-shot Image Classification](https://mlsanity.com/paper/2302.01190/) (2023)
- [Towards Multi-modal Transformers in Federated Learning](https://mlsanity.com/paper/2404.12467/) (2024)

## Abstract (arXiv)

Federated learning is an emerging research paradigm enabling collaborative training of machine learning models among different organizations while keeping data private at each institution. Despite recent progress, there remain fundamental challenges such as the lack of convergence and the potential for catastrophic forgetting across real-world heterogeneous devices. In this paper, we demonstrate that self-attention-based architectures (e.g., Transformers) are more robust to distribution shifts and hence improve federated learning over heterogeneous data. Concretely, we conduct the first rigorous empirical investigation of different neural architectures across a range of federated algorithms, real-world benchmarks, and heterogeneous data splits. Our experiments show that simply replacing convolutional networks with Transformers can greatly reduce catastrophic forgetting of previous devices, accelerate convergence, and reach a better global model, especially when dealing with heterogeneous data. We release our code and pretrained models at https://github.com/Liangqiong/ViT-FL-main to encourage future exploration in robust architectures as an alternative to current research efforts on the optimization front.