# Gradient Masked Averaging for Federated Learning

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

## Key facts

- Cited by 22 papers in the map (data through 14 Oct 2026).
- Builds on 44 of its 55 references that are in the map.
- First posted on arXiv 28 Jan 2022.

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

- [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)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Federated Learning: Challenges, Methods, and Future Directions](https://mlsanity.com/paper/1908.07873/) (2019)
- [Group Normalization](https://mlsanity.com/paper/1803.08494/) (2018)
- [Federated Learning with Non-IID Data](https://mlsanity.com/paper/1806.00582/) (2018)
- [Adaptive Federated Optimization](https://mlsanity.com/paper/2003.00295/) (2020)
- [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)
- [LEAF: A Benchmark for Federated Settings](https://mlsanity.com/paper/1812.01097/) (2018)
- [Out-of-Distribution Generalization via Risk Extrapolation (REx)](https://mlsanity.com/paper/2003.00688/) (2020)
- [Ditto: Fair and Robust Federated Learning Through Personalization](https://mlsanity.com/paper/2012.04221/) (2020)
- [Local SGD Converges Fast and Communicates Little](https://mlsanity.com/paper/1805.09767/) (2018)
- [Federated Learning on Non-IID Data Silos: An Experimental Study](https://mlsanity.com/paper/2102.02079/) (2021)
- [Federated Learning with Matched Averaging](https://mlsanity.com/paper/2002.06440/) (2020)
- [Federated Learning with Differential Privacy: Algorithms and Performance Analysis](https://mlsanity.com/paper/1911.00222/) (2019)
- [Adaptive Federated Learning in Resource Constrained Edge Computing Systems](https://mlsanity.com/paper/1804.05271/) (2018)
- [Data-Free Knowledge Distillation for Heterogeneous Federated Learning](https://mlsanity.com/paper/2105.10056/) (2021)
- [A Unified Theory of Decentralized SGD with Changing Topology and Local Updates](https://mlsanity.com/paper/2003.10422/) (2020)
- [FedPAQ: A Communication-Efficient Federated Learning Method with Periodic Averaging and Quantization](https://mlsanity.com/paper/1909.13014/) (2019)
- [Personalized Federated Learning: A Meta-Learning Approach](https://mlsanity.com/paper/2002.07948/) (2020)
- [Parallel Restarted SGD with Faster Convergence and Less Communication: Demystifying Why Model Averaging Works for Deep Learning](https://mlsanity.com/paper/1807.06629/) (2018)
- [Gradient Starvation: A Learning Proclivity in Neural Networks](https://mlsanity.com/paper/2011.09468/) (2020)
- [Domain Generalization for Object Recognition with Multi-task Autoencoders](https://mlsanity.com/paper/1508.07680/) (2015)
- [EMNIST: an extension of MNIST to handwritten letters](https://mlsanity.com/paper/1702.05373/) (2017)
- [Invariant Risk Minimization Games](https://mlsanity.com/paper/2002.04692/) (2020)
- [FedBABU: Towards Enhanced Representation for Federated Image Classification](https://mlsanity.com/paper/2106.06042/) (2021)
- [Learning explanations that are hard to vary](https://mlsanity.com/paper/2009.00329/) (2020)

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

- [FedCLIP: Fast Generalization and Personalization for CLIP in Federated Learning](https://mlsanity.com/paper/2302.13485/) (2023)
- [Test-Time Robust Personalization for Federated Learning](https://mlsanity.com/paper/2205.10920/) (2022)
- [DiPrompT: Disentangled Prompt Tuning for Multiple Latent Domain Generalization in Federated Learning](https://mlsanity.com/paper/2403.08506/) (2024)
- [Benchmarking Algorithms for Federated Domain Generalization](https://mlsanity.com/paper/2307.04942/) (2023)
- [Task Arithmetic Through The Lens Of One-Shot Federated Learning](https://mlsanity.com/paper/2411.18607/) (2024)
- [Re-Weighted Softmax Cross-Entropy to Control Forgetting in Federated Learning](https://mlsanity.com/paper/2304.05260/) (2023)
- [Making Batch Normalization Great in Federated Deep Learning](https://mlsanity.com/paper/2303.06530/) (2023)
- [Understanding and Improving Model Averaging in Federated Learning on Heterogeneous Data](https://mlsanity.com/paper/2305.07845/) (2023)
- [Beyond the Federation: Topology-aware Federated Learning for Generalization to Unseen Clients](https://mlsanity.com/paper/2407.04949/) (2024)
- [UPFL: Unsupervised Personalized Federated Learning towards New Clients](https://mlsanity.com/paper/2307.15994/) (2023)
- [Federated Learning and Differential Privacy Techniques on Multi-hospital Population-scale Electrocardiogram Data](https://mlsanity.com/paper/2405.00725/) (2024)
- [Many-Task Federated Fine-Tuning via Unified Task Vectors](https://mlsanity.com/paper/2502.06376/) (2025)
- [Federated Generalised Variational Inference: A Robust Probabilistic Federated Learning Framework](https://mlsanity.com/paper/2502.00846/) (2025)
- [Hypernetwork-Driven Model Fusion for Federated Domain Generalization](https://mlsanity.com/paper/2402.06974/) (2024)
- [Employing Federated Learning for Training Autonomous HVAC Systems](https://mlsanity.com/paper/2405.00389/) (2024)
- [DataFreeShield: Defending Adversarial Attacks without Training Data](https://mlsanity.com/paper/2406.15635/) (2024)
- [Fast convergence of a Federated Expectation-Maximization Algorithm](https://mlsanity.com/paper/2408.05819/) (2024)
- [An Architecture Built for Federated Learning: Addressing Data Heterogeneity through Adaptive Normalization-Free Feature Recalibration](https://mlsanity.com/paper/2410.02006/) (2024)
- [PARDON: Privacy-Aware and Robust Federated Domain Generalization](https://mlsanity.com/paper/2410.22622/) (2024)
- [Unified Alignment Protocol: Making Sense of the Unlabeled Data in New Domains](https://mlsanity.com/paper/2505.21010/) (2025)
- [FTTE: Enabling Federated and Resource-Constrained Deep Edge Intelligence](https://mlsanity.com/paper/2510.03165/) (2025)
- [Forget to Generalize: Iterative Adaptation for Generalization in Federated Learning](https://mlsanity.com/paper/2602.04536/) (2026)

## Abstract (arXiv)

Federated learning (FL) is an emerging paradigm that permits a large number of clients with heterogeneous data to coordinate learning of a unified global model without the need to share data amongst each other. A major challenge in federated learning is the heterogeneity of data across client, which can degrade the performance of standard FL algorithms. Standard FL algorithms involve averaging of model parameters or gradient updates to approximate the global model at the server. However, we argue that in heterogeneous settings, averaging can result in information loss and lead to poor generalization due to the bias induced by dominant client gradients. We hypothesize that to generalize better across non-i.i.d datasets, the algorithms should focus on learning the invariant mechanism that is constant while ignoring spurious mechanisms that differ across clients. Inspired from recent works in Out-of-Distribution generalization, we propose a gradient masked averaging approach for FL as an alternative to the standard averaging of client updates. This aggregation technique for client updates can be adapted as a drop-in replacement in most existing federated algorithms. We perform extensive experiments on multiple FL algorithms with in-distribution, real-world, feature-skewed out-of-distribution, and quantity imbalanced datasets and show that it provides consistent improvements, particularly in the case of heterogeneous clients.