# On the effectiveness of partial variance reduction in federated learning with heterogeneous data

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

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Builds on 34 of its 46 references that are in the map.
- First posted on arXiv 5 Dec 2022.

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

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Densely Connected Convolutional Networks](https://mlsanity.com/paper/1608.06993/) (2016)
- [Intriguing properties of neural networks](https://mlsanity.com/paper/1312.6199/) (2013)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Advances and Open Problems in Federated Learning](https://mlsanity.com/paper/1912.04977/) (2019)
- [Similarity of Neural Network Representations Revisited](https://mlsanity.com/paper/1905.00414/) (2019)
- [Model-Contrastive Federated Learning](https://mlsanity.com/paper/2103.16257/) (2021)
- [Ensemble Distillation for Robust Model Fusion in Federated Learning](https://mlsanity.com/paper/2006.07242/) (2020)
- [Federated Optimization: Distributed Machine Learning for On-Device Intelligence](https://mlsanity.com/paper/1610.02527/) (2016)
- [Federated Learning Based on Dynamic Regularization](https://mlsanity.com/paper/2111.04263/) (2021)
- [SAGA: A Fast Incremental Gradient Method With Support for Non-Strongly Convex Composite Objectives](https://mlsanity.com/paper/1407.0202/) (2014)
- [Classification with Valid and Adaptive Coverage](https://mlsanity.com/paper/2006.02544/) (2020)
- [A Unified Theory of Decentralized SGD with Changing Topology and Local Updates](https://mlsanity.com/paper/2003.10422/) (2020)
- [Uncertainty Sets for Image Classifiers using Conformal Prediction](https://mlsanity.com/paper/2009.14193/) (2020)
- [A Field Guide to Federated Optimization](https://mlsanity.com/paper/2107.06917/) (2021)
- [Do Wide and Deep Networks Learn the Same Things? Uncovering How Neural Network Representations Vary with Width and Depth](https://mlsanity.com/paper/2010.15327/) (2020)
- [No Fear of Heterogeneity: Classifier Calibration for Federated Learning with Non-IID Data](https://mlsanity.com/paper/2106.05001/) (2021)
- [FedBABU: Towards Enhanced Representation for Federated Image Classification](https://mlsanity.com/paper/2106.06042/) (2021)
- [FedDC: Federated Learning with Non-IID Data via Local Drift Decoupling and Correction](https://mlsanity.com/paper/2203.11751/) (2022)
- [ProxSkip: Yes! Local Gradient Steps Provably Lead to Communication Acceleration! Finally!](https://mlsanity.com/paper/2202.09357/) (2022)
- [Communication Efficient Distributed Optimization using an Approximate Newton-type Method](https://mlsanity.com/paper/1312.7853/) (2013)
- [Acceleration for Compressed Gradient Descent in Distributed and Federated Optimization](https://mlsanity.com/paper/2002.11364/) (2020)
- [A Unified Theory of SGD: Variance Reduction, Sampling, Quantization and Coordinate Descent](https://mlsanity.com/paper/1905.11261/) (2019)
- [Unified Optimal Analysis of the (Stochastic) Gradient Method](https://mlsanity.com/paper/1907.04232/) (2019)
- [FedAvg with Fine Tuning: Local Updates Lead to Representation Learning](https://mlsanity.com/paper/2205.13692/) (2022)
- [On the Ineffectiveness of Variance Reduced Optimization for Deep Learning](https://mlsanity.com/paper/1812.04529/) (2018)
- [RelaySum for Decentralized Deep Learning on Heterogeneous Data](https://mlsanity.com/paper/2110.04175/) (2021)
- [On the Discrepancy between the Theoretical Analysis and Practical Implementations of Compressed Communication for Distributed Deep Learning](https://mlsanity.com/paper/1911.08250/) (2019)

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

- [Federated Learning for Generalization, Robustness, Fairness: A Survey and Benchmark](https://mlsanity.com/paper/2311.06750/) (2023)
- [No Fear of Classifier Biases: Neural Collapse Inspired Federated Learning with Synthetic and Fixed Classifier](https://mlsanity.com/paper/2303.10058/) (2023)
- [Geometric Knowledge-Guided Localized Global Distribution Alignment for Federated Learning](https://mlsanity.com/paper/2503.06457/) (2025)
- [Scaff-PD: Communication Efficient Fair and Robust Federated Learning](https://mlsanity.com/paper/2307.13381/) (2023)
- [FedBug: A Bottom-Up Gradual Unfreezing Framework for Federated Learning](https://mlsanity.com/paper/2307.10317/) (2023)
- [Synthetic data shuffling accelerates the convergence of federated learning under data heterogeneity](https://mlsanity.com/paper/2306.13263/) (2023)
- [Federated Learning for Blind Image Super-Resolution](https://mlsanity.com/paper/2404.17670/) (2024)
- [FedSPC: Shared Parameter Correction for Personalized Federated Learning](https://mlsanity.com/paper/2606.13748/) (2026)
- [Global Federated Learning Strategies for Building Efficient Personalized Models](https://mlsanity.com/paper/2608.15107/) (2026)

## Abstract (arXiv)

Data heterogeneity across clients is a key challenge in federated learning. Prior works address this by either aligning client and server models or using control variates to correct client model drift. Although these methods achieve fast convergence in convex or simple non-convex problems, the performance in over-parameterized models such as deep neural networks is lacking. In this paper, we first revisit the widely used FedAvg algorithm in a deep neural network to understand how data heterogeneity influences the gradient updates across the neural network layers. We observe that while the feature extraction layers are learned efficiently by FedAvg, the substantial diversity of the final classification layers across clients impedes the performance. Motivated by this, we propose to correct model drift by variance reduction only on the final layers. We demonstrate that this significantly outperforms existing benchmarks at a similar or lower communication cost. We furthermore provide proof for the convergence rate of our algorithm.