Correcting Before Adapting: Gradient Correction for Federated Adaptive Optimization

Evan Chen; Shiqiang Wang; Dong-Jun Han; Shahryar Zehtabi; Christopher Brinton

cs.LGFeb 4, 2025

Correcting Before Adapting: Gradient Correction for Federated Adaptive Optimization

Authors: Evan Chen, Shiqiang Wang, Dong-Jun Han, Shahryar Zehtabi, Christopher Brinton

Organizations: Elmore Family School of Electrical and Computer Engineering, Purdue University, West Lafayette, IN · IBM T. J. Watson Research Center, Yorktown Heights, NY, USA · Department of Computer Science and Engineering, Yonsei University, Seoul, Korea

Abstract

Client-drift compensation methods mitigate the impact of data heterogeneity in federated learning (FL). While drift correction is well studied for stochastic gradient descent (SGD)-based FL, its interaction with adaptive client optimizers remains less understood. Adaptive optimizers such as Adam build each update from running averages of past gradients, so correcting only the final update direction leaves those averages, and the drift they contain, unchanged. We propose Federated Adaptive Gradient Correction ({\tt FA-GC}), which corrects each gradient before it enters these running averages, and compare it with Federated Adaptive Naive Tracking ({\tt FA-NT}), which adds the correction to the final direction. We show that, in an idealized setting without gradient noise, the direction FA-GC follows vanishes only at stationary points of the global objective, whereas the direction FA-NT follows can vanish elsewhere when the optimizer is nonlinear. We also prove convergence of FA-GC for Adam-family methods, gradient normalization, and Muon, which orthogonalizes matrix updates. Unlike our bound for FA-NT, these bounds contain no term for client heterogeneity. On seven image and language tasks, FA-GC with Adam needs the fewest rounds on average to reach the target accuracy, and on the three image tasks it also uses less total communication than uncorrected local Adam and FA-NT. On the QNLI question-entailment task with SOAP, which applies Adam in a rotated basis, and Muon under strong heterogeneity, FA-GC gives the highest accuracy, while FA-NT falls below the uncorrected optimizers.

Figures & tables

Panel of Figure 1 — Fig. 1: Where the correction enters the update. Top: related methods grouped by correction placement. Bottom: FA-GC corrects the gradient before moment estimation, while FA-NT corrects the resulting direction. Here $g_{i}$ is the raw client gradient, $c_{i}=y-y_{i}$ the difference between server and client trackers, $(m_{i},v_{i})$ the first and second moments, and $d_{i}=\mathcal{P}_{i}(m_{i};v_{i})$ the optimizer direction. Both placements are analyzed for the three optimizer families shown; finite-step Muon approximates the polar map with a fixed number of Newton–Schulz iterations (Section IV-B ).

Fig. 2: Client drift under local SGD, adaptive optimization, and adaptive optimization with gradient correction. Gray and red dashed contours represent client and global objectives. Black dots $x_{i}^{*}$ mark client minimizers, and the red diamond $x^{*}$ marks the global minimizer. Trajectories start from $x^{0}$ , and the magenta point $\bar{x}$ is the average local model. Gradient correction changes the directions used by the adaptive optimizer.

Method	$\phi(g)$	$C_{2}$	$\\|\nabla f\\|_{*}^{q}$	$C_{3}$	$\chi_{E}$
Adam, AMSGrad, SOAP †	$g$	$\lambda_{p}/2$	$\\|\nabla f\\|_{2}^{2}$	$\frac{1}{2\epsilon}$	$\frac{\epsilon}{2}$
Normalization	$N(g)$	$1$	$\\|\nabla f\\|_{2}$	$2$	$0$
Exact Muon	$P_{\rm M}(g)$	$1$	$\\|\nabla f\\|_{\rm tr}$	$2\sqrt{r}$	$0$
Finite-step Muon	$P_{\rm M}(g)$	$r^{-1/2}$	$\\|\nabla f\\|_{\rm tr}$	$1+\sqrt{r}$	$0$
Adam family, reference	$g$	$\frac{1}{2\epsilon}$	$\\|\nabla f\\|_{2}^{2}$	$\frac{1}{\epsilon}$	$0$
† SOAP uses the Frobenius norm.

Table I: Optimizer-specific quantities in Lemma V.5 . The comparison direction is

d_{\star}^{t}=C_{2}\,\phi(\nabla f(x^{t}))

, and the rate is measured by

\|\nabla f\|_{*}^{q}

. Adam-family rows have the small term

\nu=\lambda_{p}G^{2}p/2

from the event that

\mathcal{H}_{p}

fails; the others have

\nu=0

. The last row compares Adam-family directions with the fixed scaling

\epsilon^{-1}I

; it is used for FA-NT and in Appendix C-B , and its bound ( 14 ) also contains the residual

\epsilon\|(D-\epsilon^{-1}I)m\|_{2}^{2}

, where

D

is the optimizer’s scaling matrix. For a

d_{1}\times d_{2}

Muon parameter,

r=\min(d_{1},d_{2})

is its largest possible rank.

Settings	Task Type	Dataset	FedAvg-M [ 14 ]	SCAFFOLD-M [ 14 ]	FedAdam [ 15 ]	FedAMS [ 16 ]	LocalAdam	PAdaMFed [ 17 ]	FA-NT (Adam)	FA-GC (Adam)
Communication Rounds	Image Tasks	CIFAR-10	1014.5 $\pm$ 250.3	544.3 $\pm$ 59.9	2532.5 $\pm$ 343.3	2388.8 $\pm$ 286.6	589.5 $\pm$ 74.0	360.3 $\pm$ 20.8	394.8 $\pm$ 31.3	310.0 $\pm$ 16.8
		CIFAR-100	998.5 $\pm$ 88.5	621.8 $\pm$ 55.4	1854.0 $\pm$ 185.3	1654 $\pm$ 156.4	678.3 $\pm$ 40.6	511.0 $\pm$ 15.2	530.3 $\pm$ 17.6	323.8 $\pm$ 16.3
		TinyImageNet	215.5 $\pm$ 10.4	242.2 $\pm$ 15.4	543.2 $\pm$ 45.2	463.5 $\pm$ 35.7	177.3 $\pm$ 8.3	87.5 $\pm$ 6.4	157.0 $\pm$ 6.4	66.3 $\pm$ 4.4
	Language Tasks	20NewsGroups	245.5 $\pm$ 27.6	214.0 $\pm$ 17.7	247.0 $\pm$ 8.1	224.5 $\pm$ 5.4	156.8 $\pm$ 7.8	149.3 $\pm$ 3.4	155.0 $\pm$ 7.0	143.3 $\pm$ 4.1
		QNLI	171.3 $\pm$ 41.2	162.5 $\pm$ 37.4	137.5 $\pm$ 10.2	145.0 $\pm$ 9.8	117.0 $\pm$ 10.4	87.3 $\pm$ 13.2	99.8 $\pm$ 11.2	55.5 $\pm$ 16.3
		QQP	316.5 $\pm$ 64.3	299.0 $\pm$ 70.3	245.3 $\pm$ 20.5	267.8 $\pm$ 22.1	213.0 $\pm$ 53.8	84.7 $\pm$ 4.3	196.3 $\pm$ 5.1	63.0 $\pm$ 4.6

Table II: Communication rounds to target, reported as mean

\pm

standard deviation over four trials. Unless otherwise stated, each round samples

S=10

of the

n=100

clients. Targets are 75% for CIFAR-10, 50% for CIFAR-100, 30% for TinyImageNet, 85% for SST-2, and 75% for the other language tasks.

Panel of Figure 5 — Fig. 3: Adam sensitivity on CIFAR-100. In each pair, the left plot reports communication rounds and the right plot reports communication per client. (a) Cost to reach 50% accuracy as the Dirichlet parameter varies. Smaller $\alpha$ means stronger heterogeneity. Each point in both panels is a single run, so the values at $\alpha=0.1$ differ from the four-trial means in Table II . (b) With $n=100$ and $S=50$ , the horizontal axis $Y$ denotes the refresh count $\widetilde{S}$ . Communication counts total transmitted model-sized vectors divided by $S$ . Global iterations are communication rounds. The horizontal lines show LocalAdam and FedAvg (local SGD with server averaging), which use no trackers and so do not depend on $\widetilde{S}$ .

Fig. 4: QNLI with GPT-2 and fixed $\beta_{1},\beta_{2}$ : practical SOAP at $\alpha=0.1$ and Muon with five Newton–Schulz iterations at $\alpha\in\{0.1,0.3,1.0\}$ (Section VI ). Local is the same optimizer without tracking. All runs use $n=100$ , $S=10$ , one local epoch, 500 rounds, and four shared seeds. Curves and bands show the mean and one standard deviation; printed values are rounds-401–500 means in Local/FA-NT/FA-GC order.

Dataset	LocalAdam	FA-NT	FA-GC
CIFAR-10	$551.77\pm 69.26$	$646.68\pm 51.27$	$\mathbf{507.78}\pm 27.52$
CIFAR-100	$634.89\pm 38.00$	$868.63\pm 28.83$	$\mathbf{530.38}\pm 26.70$
TinyImageNet	$165.95\pm 7.77$	$257.17\pm 10.48$	$\mathbf{108.60}\pm 7.21$

Table III: Aggregate image-training payload to target (decimal GB), derived from the round means and standard deviations.

Dataset	$\beta_{2}$
	$0$	$0.3$	$0.7$	$0.99$
CIFAR-10	$358.8\!\pm\!21.4$	$334.3\!\pm\!20.2$	$318.8\!\pm\!15.3$	$\mathbf{310.0\!\pm\!16.8}$
CIFAR-100	$527.0\!\pm\!16.5$	$413.3\!\pm\!15.5$	$370.5\!\pm\!14.7$	$\mathbf{323.8\!\pm\!16.3}$
TinyImageNet	$85.5\!\pm\!5.7$	$75.3\!\pm\!5.7$	$74.0\!\pm\!6.6$	$\mathbf{66.3\!\pm\!4.4}$

Table IV: Rounds to target for FA-GC (Adam) under fixed

\beta_{2}

: mean

\pm

standard deviation over four trials;

\beta_{2}=0.99

is the default.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Group	Notation	Description
System	$n,S,\widetilde{S};s,\widetilde{s}$	Client count, training and refresh counts, and participation fractions $S/n$ and $\widetilde{S}/n$ .
	$T,K;t,k$	Communication rounds, local steps, and their indices.
	$\mathcal{S}^{t},\widetilde{\mathcal{S}}^{t}$	Training and tracker-refresh client sets in round $t$ .
	$f,f_{i},f_{*}$	Global objective, client objective, and finite lower bound.
	$\mathcal{D}_{i},\xi_{i}$	Client data distribution and a stochastic sample.
	$d;L,L_{g}$	Number of scalar parameters; smoothness constants of $f$ and of each $\nabla f_{i}$ (Assumption V.1 ).

Appendix

Table V: Consolidated notation used in the main text and proofs.

Optimizer / $\alpha$	Local	FA-NT	FA-GC
SOAP / 0.1	$(.003,.3)$	$(.003,.3)$	$(.001,.3)$
Muon / 0.1	$(.003,.3)$	$(.001,.9)$	$(.003,.1)$
Muon / 0.3	$(.003,.3)$	$(.001,.9)$	$(.003,.3)$
Muon / 1.0	$(.003,.9)$	$(.003,.9)$	$(.003,.3)$

Appendix

Table VII: Selected

(\eta_{\ell},\eta_{g})

for the QNLI optimizer study.

Fig. 5: Five-step Newton–Schulz Muon with fixed $\beta_{1}$ on federated SST-2 with GPT-2 and LoRA at Dirichlet $\alpha\in\{0.1,0.3\}$ . All runs use $n=100$ , $S=10$ , one complete local epoch with batch size 32 and sequence length 128, 500 communication rounds, and four shared seeds. Curves show mean held-out accuracy, and bands show one standard deviation. Values below the panels are rounds-401–500 means in Local/FA-NT/FA-GC order.

Dirichlet $\alpha$	Method	Local LR	Server LR
0.1	Local	$3\times 10^{-3}$	0.9
0.1	FA-NT	$3\times 10^{-3}$	0.9
0.1	FA-GC	$3\times 10^{-3}$	0.3
0.3	Local	$3\times 10^{-3}$	0.9
0.3	FA-NT	$3\times 10^{-3}$	0.9
0.3	FA-GC	$3\times 10^{-3}$	0.3

Appendix

Table VIII: Selected Learning Rates for the SST-2 Muon Runs

Fig. 6: Additional CIFAR-100 results under varying client sampling rates with $K=60$ , Dirichlet $\alpha=0.1$ , and $\widetilde{S}=S$ (every participating client refreshes). Each curve is a single run.

Fig. 7: Additional CIFAR-100 results under varying client sampling rates with $K=10$ , Dirichlet $\alpha=0.1$ , and $\widetilde{S}=S$ (every participating client refreshes). Each curve is a single run.

Fig. 8: Additional CIFAR-10 results under varying client sampling rates with $K=60$ , Dirichlet $\alpha=0.1$ , and $\widetilde{S}=S$ (every participating client refreshes). Each curve is a single run.

Fig. 9: Additional CIFAR-10 results under varying client sampling rates with $K=10$ , Dirichlet $\alpha=0.1$ , and $\widetilde{S}=S$ (every participating client refreshes). Each curve is a single run.

Fig. 10: Mean and variance across four random trials under different levels of data heterogeneity.

Settings	Dataset	LocalAdam	FA-NT ( $\beta_{2}=0$ )	FA-GC ( $\beta_{2}=0$ )	FA-GC ( $\beta_{2}=0.3$ )	FA-GC ( $\beta_{2}=0.7$ )	FA-NT	FA-GC
Communication Rounds	CIFAR-10	589.5 $\pm$ 74.0	401.5 $\pm$ 33.9	358.8 $\pm$ 21.4	334.3 $\pm$ 20.2	318.8 $\pm$ 15.3	394.8 $\pm$ 31.3	310.0 $\pm$ 16.8
	CIFAR-100	678.3 $\pm$ 40.6	867.5 $\pm$ 25.3	527.0 $\pm$ 16.5	413.3 $\pm$ 15.5	370.5 $\pm$ 14.7	530.3 $\pm$ 17.6	323.8 $\pm$ 16.3
	TinyImageNet	177.3 $\pm$ 8.3	164.3 $\pm$ 6.7	85.5 $\pm$ 5.7	75.3 $\pm$ 5.7	74.0 $\pm$ 6.6	157.0 $\pm$ 6.4	66.3 $\pm$ 4.4

Appendix

Table IX: Rounds to target for Adam instantiations under different

\beta_{2}

choices (mean and standard deviation over four trials).

Learning Rate	Dataset	FedAvg-M	SCAFFOLD-M	FedAdam	FedAMS	LocalAdam	PAdaMFed	FA-NT (Adam)	FA-GC (Adam)
$\eta_{g}$	CIFAR-10	1	1	$1\times 10^{-3}$	$1\times 10^{-3}$	1	1	1	1
	CIFAR-100	1	1	$1\times 10^{-3}$	$1\times 10^{-3}$	1	1	1	1
	TinyImageNet	$3\times 10^{-1}$	$3\times 10^{-1}$	$1\times 10^{-3}$	$1\times 10^{-3}$	$1\times 10^{-1}$	$1\times 10^{-1}$	$3\times 10^{-1}$	$3\times 10^{-1}$
$\eta_{\ell}$	CIFAR-10	$1\times 10^{-2}$	$1\times 10^{-2}$	$3\times 10^{-2}$	$3\times 10^{-2}$	$1\times 10^{-3}$	$1\times 10^{-3}$	$1\times 10^{-3}$	$1\times 10^{-3}$
	CIFAR-100	$1\times 10^{-2}$	$1\times 10^{-2}$	$3\times 10^{-2}$	$3\times 10^{-2}$	$1\times 10^{-3}$	$1\times 10^{-3}$	$1\times 10^{-3}$	$1\times 10^{-3}$
	TinyImageNet	$1\times 10^{-2}$	$1\times 10^{-2}$	$1\times 10^{-2}$	$1\times 10^{-2}$	$1\times 10^{-3}$	$1\times 10^{-3}$	$1\times 10^{-3}$	$1\times 10^{-3}$

Appendix

Table X: The hyperparameters for image classification tasks.

Learning Rate	Dataset	FedAvg-M	SCAFFOLD-M	FedAdam	FedAMS	LocalAdam	PAdaMFed	FA-NT (Adam)	FA-GC (Adam)
$\eta_{g}$	20NEWSGROUPS	$1\times 10^{-1}$	$1\times 10^{-1}$	$1\times 10^{-2}$	$1\times 10^{-2}$	$1\times 10^{-1}$	$1\times 10^{-1}$	$1\times 10^{-1}$	$1\times 10^{-1}$
	QQP	$1\times 10^{-1}$	$1\times 10^{-1}$	$1\times 10^{-3}$	$1\times 10^{-3}$	$3\times 10^{-1}$	$1\times 10^{-1}$	$3\times 10^{-1}$	$3\times 10^{-1}$
	QNLI	$1\times 10^{-1}$	$1\times 10^{-1}$	$1\times 10^{-3}$	$1\times 10^{-3}$	$3\times 10^{-1}$	$1\times 10^{-1}$	$3\times 10^{-1}$	$3\times 10^{-1}$
	SST-2	$1\times 10^{-1}$	$1\times 10^{-1}$	$1\times 10^{-3}$	$1\times 10^{-3}$	$3\times 10^{-1}$	$1\times 10^{-1}$	$3\times 10^{-1}$	$3\times 10^{-1}$
$\eta_{\ell}$	20NEWSGROUPS	$1\times 10^{-2}$	$1\times 10^{-2}$	$1\times 10^{-3}$	$1\times 10^{-3}$	$5\times 10^{-3}$	$5\times 10^{-3}$	$5\times 10^{-3}$	$5\times 10^{-3}$
	QQP	$1\times 10^{-2}$	$1\times 10^{-2}$	$3\times 10^{-4}$	$3\times 10^{-4}$	$1\times 10^{-4}$	$1\times 10^{-4}$	$1\times 10^{-4}$	$1\times 10^{-4}$

Appendix

Table XI: The hyperparameters for language tasks.

Task	FedAvg-M	SCAFFOLD-M	FedAdam	FedAMS	LocalAdam	FA-GC (Adam)
Image task (ResNet-18)	35.1M ( $\sim$ 0.14 GB)	46.8M ( $\sim$ 0.18 GB)	23.4M ( $\sim$ 0.09 GB)	23.4M ( $\sim$ 0.09 GB)	46.8M ( $\sim$ 0.18 GB)	58.5M ( $\sim$ 0.23 GB)
Language task	125.48M ( $\sim$ 0.50 GB)	126.22M ( $\sim$ 0.50 GB)	124.74M ( $\sim$ 0.49 GB)	124.74M ( $\sim$ 0.49 GB)	126.22M ( $\sim$ 0.50 GB)	126.96M ( $\sim$ 0.51 GB)

Appendix

Table XII: Model, gradient, optimizer, and tracker storage per client, excluding activations and temporary workspace. Counts are scalars, with approximate FP32 memory in parentheses at 4 bytes per scalar.

Dataset	FedAvg-M	SCAFFOLD-M	FedAdam	FedAMS	LocalAdam	PAdaMFed	FA-NT (Adam)	FA-GC (Adam)
Vectors/round	30	40	20	20	20	60	35	35
CIFAR-10	1424.36 $\pm$ 351.42	1018.93 $\pm$ 112.13	2370.42 $\pm$ 321.33	2235.92 $\pm$ 268.26	551.77 $\pm$ 69.26	1011.72 $\pm$ 58.41	646.68 $\pm$ 51.27	507.78 $\pm$ 27.52
CIFAR-100	1401.89 $\pm$ 124.25	1164.01 $\pm$ 103.71	1735.34 $\pm$ 173.44	1548.14 $\pm$ 146.39	634.89 $\pm$ 38.00	1434.89 $\pm$ 42.68	868.63 $\pm$ 28.83	530.38 $\pm$ 26.70
TinyImageNet	302.56 $\pm$ 14.60	453.40 $\pm$ 28.83	508.44 $\pm$ 42.31	433.84 $\pm$ 33.42	165.95 $\pm$ 7.77	245.70 $\pm$ 17.97	257.17 $\pm$ 10.48	108.60 $\pm$ 7.21

Appendix

Table XIII: Aggregate training-round communication to reach the target accuracy, in decimal GB. Mean and standard deviation are obtained from the round statistics in Table II using a common 46.8 MB ResNet-18 payload and the message counts in Appendix H-G .

Federated Learning

Federated learning trains shared models across data holders without requiring them to pool their raw data.

See all 951 papers

FedZMG: Efficient Client-Side Optimization in Federated Learning
Feb 20, 2026Fotios Zantalis, Evangelos Zervas, Grigorios KoulourasFederated Learning Non-IID Federated Learning
TelSiP Research Laboratory, Department of Electrical and Electronic Engineering, School of Engineering, University of West Attica, Ancient Olive Grove Campus, 250 Thivon Str., GR-12241 Athens, Greece
SpecGradFilter: A Spectral Gradient Filtering Framework for Taming Federated Heterogeneity
Jul 5, 2026Liyang Yuan, Yibo Yang, Dandan Guo +2Non-IID Federated Learning Federated Learning
School of Artificial Intelligence, Jilin University, Changchun 130012, China · King Abdullah University of Science and Technology (KAUST), Thuwal, Saudi Arabia · School of Intelligence Science and Technology, Peking University, China
FedSteer: Taming Extreme Gradient Staleness in Federated Learning with Corrective Projections and Caching
Jun 8, 2026Haoran Zhang, Cainã Figueiredo Pereira, Marie Siew +3Federated Learning Aggregation Federated Learning
Electrical and Computer Engineering Dept., The University of Texas at Austin, Austin, Texas, USA · CERI/LIA, University of Avignon, Avignon, France · Information Systems Technology and Design Pillar, Singapore University of Technology and Design, Singapore +2

Feb 20, 2026J/K move · Enter open · S save

FedZMG: Efficient Client-Side Optimization in Federated Learning

Fotios Zantalis, Evangelos Zervas, Grigorios Koulouras

TelSiP Research Laboratory, Department of Electrical and Electronic Engineering, School of Engineering, University of West Attica, Ancient Olive Grove Campus, 250 Thivon Str., GR-12241 Athens, Greece

Federated Learning Non-IID Federated Learning

Open paper AskPDF arXiv

Federated Learning (FL) enables distributed model training on edge devices while preserving data privacy. However, clients tend to have non-Independent and Identically Distributed (non-IID) data, which often leads to client-drift, and therefore diminishing convergence speed and model performance. While adaptive optimizers have been proposed to mitigate these effects, they frequently introduce computational complexity or communication overhead unsuitable for resource-constrained IoT environments. This paper introduces Federated Zero Mean Gradients (FedZMG), a novel, parameter-free, client-side optimization algorithm designed to tackle client-drift by structurally regularizing the optimization space. Advancing the idea of Gradient Centralization, FedZMG projects local gradients onto a zero-mean hyperplane, effectively neutralizing the "intensity" or "bias" shifts inherent in heterogeneous data distributions without requiring additional communication or hyperparameter tuning. A theoretical analysis is provided, proving that FedZMG reduces the effective gradient variance and guarantees tighter convergence bounds compared to standard FedAvg. Extensive empirical evaluations on EMNIST, CIFAR100, and Shakespeare datasets demonstrate that FedZMG achieves better convergence speed and final validation accuracy compared to the baseline FedAvg and the adaptive optimizer FedAdam, particularly in highly non-IID settings.

Abstract

Federated Learning

Explore similar work

FedZMG: Efficient Client-Side Optimization in Federated Learning

SpecGradFilter: A Spectral Gradient Filtering Framework for Taming Federated Heterogeneity

FedSteer: Taming Extreme Gradient Staleness in Federated Learning with Corrective Projections and Caching