# Implicit Gradient Alignment in Distributed and Federated Learning

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

## Key facts

- Cited by 29 papers in the map (data through 14 Oct 2026).
- Builds on 31 of its 46 references that are in the map.
- First posted on arXiv 25 Jun 2021.

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

- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [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)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Understanding deep learning requires rethinking generalization](https://mlsanity.com/paper/1611.03530/) (2016)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [Adaptive Federated Optimization](https://mlsanity.com/paper/2003.00295/) (2020)
- [On First-Order Meta-Learning Algorithms](https://mlsanity.com/paper/1803.02999/) (2018)
- [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)
- [Federated Optimization: Distributed Machine Learning for On-Device Intelligence](https://mlsanity.com/paper/1610.02527/) (2016)
- [Sharp Minima Can Generalize For Deep Nets](https://mlsanity.com/paper/1703.04933/) (2017)
- [Stochastic Gradient Descent as Approximate Bayesian Inference](https://mlsanity.com/paper/1704.04289/) (2017)
- [Three Factors Influencing Minima in SGD](https://mlsanity.com/paper/1711.04623/) (2017)
- [Tighter Theory for Local SGD on Identical and Heterogeneous Data](https://mlsanity.com/paper/1909.04746/) (2019)
- [Don't Use Large Mini-Batches, Use Local SGD](https://mlsanity.com/paper/1808.07217/) (2018)
- [Measuring the Effects of Data Parallelism on Neural Network Training](https://mlsanity.com/paper/1811.03600/) (2018)
- [On the Origin of Implicit Regularization in Stochastic Gradient Descent](https://mlsanity.com/paper/2101.12176/) (2021)
- [Is Local SGD Better than Minibatch SGD?](https://mlsanity.com/paper/2002.07839/) (2020)
- [The Power of Interpolation: Understanding the Effectiveness of SGD in Modern Over-parametrized Learning](https://mlsanity.com/paper/1712.06559/) (2017)
- [Implicit Gradient Regularization](https://mlsanity.com/paper/2009.11162/) (2020)
- [A Bayesian Perspective on Generalization and Stochastic Gradient Descent](https://mlsanity.com/paper/1710.06451/) (2017)
- [Stochastic gradient descent performs variational inference, converges to limit cycles for deep networks](https://mlsanity.com/paper/1710.11029/) (2017)
- [Hessian-based Analysis of Large Batch Training and Robustness to Adversaries](https://mlsanity.com/paper/1802.08241/) (2018)
- [Accelerating Federated Learning via Momentum Gradient Descent](https://mlsanity.com/paper/1910.03197/) (2019)
- [Fast Convergence of Stochastic Gradient Descent under a Strong Growth Condition](https://mlsanity.com/paper/1308.6370/) (2013)
- [Stiffness: A New Perspective on Generalization in Neural Networks](https://mlsanity.com/paper/1901.09491/) (2019)
- [Coherent Gradients: An Approach to Understanding Generalization in Gradient Descent-based Optimization](https://mlsanity.com/paper/2002.10657/) (2020)
- [Extrapolation for Large-batch Training in Deep Learning](https://mlsanity.com/paper/2006.05720/) (2020)

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

- [Federated Learning for Generalization, Robustness, Fairness: A Survey and Benchmark](https://mlsanity.com/paper/2311.06750/) (2023)
- [Personalized Federated Learning with Feature Alignment and Classifier Collaboration](https://mlsanity.com/paper/2306.11867/) (2023)
- [Stochastic Training is Not Necessary for Generalization](https://mlsanity.com/paper/2109.14119/) (2021)
- [Learning to Retain while Acquiring: Combating Distribution-Shift in Adversarial Data-Free Knowledge Distillation](https://mlsanity.com/paper/2302.14290/) (2023)
- [FRAug: Tackling Federated Learning with Non-IID Features via Representation Augmentation](https://mlsanity.com/paper/2205.14900/) (2022)
- [Learning to Unlearn while Retaining: Combating Gradient Conflicts in Machine Unlearning](https://mlsanity.com/paper/2503.06339/) (2025)
- [Data-heterogeneity-aware Mixing for Decentralized Learning](https://mlsanity.com/paper/2204.06477/) (2022)
- [Implicit meta-learning may lead language models to trust more reliable sources](https://mlsanity.com/paper/2310.15047/) (2023)
- [Neural Thermodynamics: Entropic Forces in Deep and Universal Representation Learning](https://mlsanity.com/paper/2505.12387/) (2025)
- [FREE: Faster and Better Data-Free Meta-Learning](https://mlsanity.com/paper/2405.00984/) (2024)
- [A Layer Selection Approach to Test Time Adaptation](https://mlsanity.com/paper/2404.03784/) (2024)
- [Distributed Personalized Empirical Risk Minimization](https://mlsanity.com/paper/2310.17761/) (2023)
- [Deep Fusion: Efficient Network Training via Pre-trained Initializations](https://mlsanity.com/paper/2306.11903/) (2023)
- [Implicit biases in multitask and continual learning from a backward error analysis perspective](https://mlsanity.com/paper/2311.00235/) (2023)
- [Leveraging Function Space Aggregation for Federated Learning at Scale](https://mlsanity.com/paper/2311.10291/) (2023)
- [FAGH: Accelerating Federated Learning with Approximated Global Hessian](https://mlsanity.com/paper/2403.11041/) (2024)
- [Contrastive Federated Learning with Tabular Data Silos](https://mlsanity.com/paper/2409.06123/) (2024)
- [Bi-discriminator Domain Adversarial Neural Networks with Class-Level Gradient Alignment](https://mlsanity.com/paper/2310.13959/) (2023)
- [pFedSOP : Accelerating Training Of Personalized Federated Learning Using Second-Order Optimization](https://mlsanity.com/paper/2506.07159/) (2025)
- [Uncertainty Minimization for Personalized Federated Semi-Supervised Learning](https://mlsanity.com/paper/2205.02438/) (2022)
- [Personalized Federated Learning with Contextual Modulation and Meta-Learning](https://mlsanity.com/paper/2312.15191/) (2023)
- [Majority Kernels: An Approach to Leverage Big Model Dynamics for Efficient Small Model Training](https://mlsanity.com/paper/2402.05033/) (2024)
- [DataFreeShield: Defending Adversarial Attacks without Training Data](https://mlsanity.com/paper/2406.15635/) (2024)
- [ATR-Bench: A Federated Learning Benchmark for Adaptation, Trust, and Reasoning](https://mlsanity.com/paper/2505.16850/) (2025)
- [Accelerated Training of Federated Learning via Second-Order Methods](https://mlsanity.com/paper/2505.23588/) (2025)
- [Client Clustering Meets Knowledge Sharing: Enhancing Privacy and Robustness in Personalized Peer-to-Peer Learning](https://mlsanity.com/paper/2506.20413/) (2025)
- [FedDPC : Handling Data Heterogeneity and Partial Client Participation in Federated Learning](https://mlsanity.com/paper/2512.20329/) (2025)
- [Personalized Federated Learning via Gaussian Generative Modeling](https://mlsanity.com/paper/2603.11620/) (2026)
- [Class-Grouped Normalized Momentum and Faster Hyperparameter Exploration to Tackle Class Imbalance in Federated Learning](https://mlsanity.com/paper/2607.01474/) (2026)

## Abstract (arXiv)

A major obstacle to achieving global convergence in distributed and federated learning is the misalignment of gradients across clients, or mini-batches due to heterogeneity and stochasticity of the distributed data. In this work, we show that data heterogeneity can in fact be exploited to improve generalization performance through implicit regularization. One way to alleviate the effects of heterogeneity is to encourage the alignment of gradients across different clients throughout training. Our analysis reveals that this goal can be accomplished by utilizing the right optimization method that replicates the implicit regularization effect of SGD, leading to gradient alignment as well as improvements in test accuracies. Since the existence of this regularization in SGD completely relies on the sequential use of different mini-batches during training, it is inherently absent when training with large mini-batches. To obtain the generalization benefits of this regularization while increasing parallelism, we propose a novel GradAlign algorithm that induces the same implicit regularization while allowing the use of arbitrarily large batches in each update. We experimentally validate the benefits of our algorithm in different distributed and federated learning settings.