Taming Preconditioner Drift: Unlocking the Potential of Second-Order Optimizers for Federated Learning on Non-IID Data

Junkang Liu; Fanhua Shang; Hongying Liu; Jin Liu; Weixin An; Yuanyuan Liu

cs.LGFeb 22, 2026

Taming Preconditioner Drift: Unlocking the Potential of Second-Order Optimizers for Federated Learning on Non-IID Data

Authors: Junkang Liu, Fanhua Shang, Hongying Liu, Jin Liu, Weixin An, Yuanyuan Liu

Organizations: College of Intelligence and Computing, Tianjin University, Tianjin, China · Medical School, Tianjin University, Tianjin, China · Peng Cheng Lab, Shenzhen, China · School of Cyber Engineering, Xidian University, Xi’an, China ·

Abstract

Second-order optimizers can significantly accelerate large-scale training, yet their naive federated variants are often unstable or even diverge on non-IID data. We show that a key culprit is preconditioner drift: client-side second-order training induces heterogeneous curvature-defined geometries (i.e., preconditioner coordinate systems), and server-side model averaging updates computed under incompatible metrics, corrupting the global descent direction. To address this geometric mismatch, we propose FedPAC, a preconditioner alignment and correction framework for reliable federated second-order optimization. FedPAC explicitly decouples parameter aggregation from geometry synchronization by: (i) Alignment (i.e.,aggregating local preconditioners into a global reference and warm-starting clients via global preconditioner); and (ii) Correction (i.e., steering local preconditioned updates using a global preconditioned direction to suppress long-term drift). We provide drift-coupled non-convex convergence guarantees with linear speedup under partial participation. Empirically, FedPAC consistently improves stability and accuracy across vision and language tasks, achieving up to $5.8\%$ absolute accuracy gain on CIFAR-100 with ViTs. Code is available at https://anonymous.4open.science/r/FedPAC-8B24.

Figures & tables

Figure 1: (a) In non-IID FL, first-order methods converge slowly, inducing little client drift. (b) Second-order methods converge faster locally and thus drift toward local optima, causing the aggregated global model to deviate from global optimum. (c) FedPAC corrects local second-order updates, yielding faster convergence and a global model closer to the global optimum. (d–f) FedPAC accelerates Sophia, Muon and SOAP to train ResNet-18 on CIFAR-100. The x-axis denotes the number of communication rounds, and the y-axis denotes test accuracy.

Panel of Figure 2 — Figure 2: The x-axis denotes communication rounds, and the y-axis is test accuracy. (a, c): In FL on IID data, second-order optimizers converge significantly faster than SGD and AdamW for training ResNet and Transformer. (b),(d): However, on non-IID data, second-order optimizer (e.g., Local Muon) converges much more slowly and can even underperform first-order methods such as Local SGD.

Panel of Figure 3 — Figure 3: Performance and Preconditioner Drift of Local SOAP, FedPAC_SOAP . For SOAP, we compute the preconditioner drift in a layer-wise manner by measuring the spectral norm of the difference between each client’s left/right preconditioner and the aggregated global preconditioner. Our FedPAC_SOAP substantially reduces the preconditioner drift of Local SOAP and accelerates convergence in (a,b), reaching 40% test accuracy in fewer rounds, as shown in (c) and (d).

Figure 4: An illustration of preconditioner drift, which corrects client drift through local-global alignment.

Panel of Figure 5 — Figure 5: Test accuracy versus communication rounds. The x-axis denotes communication rounds and the y-axis denotes the test accuracy. Panels (a–d) correspond to training with ResNet-18 , while panels (e–h) correspond to training with ViT-Tiny .

Method	ResNet-18				ViT-Tiny
	CIFAR-100		Tiny-Imagenet		CIFAR-100		Tiny-Imagenet
	Dir-0.1	Dir-0.05	Dir-0.1	Dir-0.05	Dir-0.1	Dir-0.05	Dir-0.1	Dir-0.05
FedAvg	60.17	56.75	47.48	43.80	27.24	23.42	15.68	14.05
SCAFFOLD	60.69	56.43	47.76	43.92	26.86	23.23	15.70	14.21
FedCM	66.61	62.65	41.16	36.00	16.95	14.74	8.88	8.15
Local AdamW	59.23	55.24	44.01	40.00	37.57	36.06	24.31	21.35

Table 1: Test accuracy of each method on CIFAR-100 and Tiny-Imagenet using ResNet-18 and ViT-Tiny over 300 communication rounds under Dir-0.1 and Dir-0.05 (100 clients, 10% participation, batch size 50,

K=50

Method	CIFAR-100		Tiny-ImageNet
Method	Test Acc	Loss	Test Acc	Loss
FedAvg	$89.28$	0.332	$86.97$	0.421
SCAFFOLD	$89.35$	0.328	$86.85$	0.432
FedCM	$87.44$	0.592	$84.24$	0.708
Local AdamW	$90.10$	0.281	$86.65$	0.252
Local Sophia	$90.21$	0.275	$86.16$	0.188

Table 2: Comparison of test accuracy and training loss for Vision Transformer (ViT-Base) under Dir-0.1 with 100 rounds (50 clients, 10% participation, batch size 16,

K=50

Method	C4 (Train Loss)
Method	LLaMA 60M	LLaMA 130M	LLaMA 350M
FedAvg	4.156	4.258	4.354
SCAFFOLD	4.028	4.238	4.365
FedCM	4.231	4.356	4.426
Local AdamW	3.213	3.418	3.519
Local Sophia	3.201	3.389	3.521

Table 3: The train loss of each method on C4 data using LLaMA 60M , LLaMA 130M , LLaMA 350M over 100 communication rounds (20 clients, 20% participation, batch size 16,

K=50

Model, Dataset, $\beta$	0.0	0.1	0.2	0.3	0.4	0.5	0.6	0.7	0.8	0.9
ResNet-18, CIFAR-100	62.05	62.35	63.10	63.57	63.75	64.16	63.75	61.75	61.70	58.86
ViT-Tiny, CIFAR-100	44.56	45.28	45.63	46.22	46.84	47.55	46.86	46.59	46.05	45.66

Table 4: Impact of

\beta

on FedPAC_SOAP with ViT-Tiny and ResNet-18 on CIFAR-100 under Dir-0.05.

Variant	ResNet-18	ViT-Tiny
Local SOAP	$58.16$	$41.68$
w/o preconditioner alignment ( $\bar{\Theta}$ )	$61.12$	$43.67$
w/o preconditioner correction ( $\Delta_{G}$ )	$62.05$	$44.56$
FedPAC_SOAP (full)	$\mathbf{64.16}$	$\mathbf{47.55}$

Table 5: Ablation on FedPAC_SOAP components under Dir-0.05 on CIFAR-100. We report test accuracy after 300 rounds.

Aggregation	Acc	CommCost	CompCost
Local SOAP	41.68	22.8 MB (1 $\times$ )	5.56 s
FedPAC_SOAP	47.55	68.4 MB (3 $\times$ )	5.85 s
FedPAC_SOAP_light	46.75	25.1 MB (1.1 $\times$ )	5.91s

Table 6: Communication-efficient preconditioner aggregation for SOAP on CIFAR-100, ViT-Tiny under Dir-0.05. CommCost denotes communication cost per round (MB), and CompCost denotes computation time per round (s).

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	#Classes	Image Size	Train	Val	Test	Total	Train / class
CIFAR-100	100	$3\times 32\times 32$	50,000	—	10,000	60,000	500
Tiny ImageNet	200	$3\times 64\times 64$	100,000	10,000	10,000	120,000	500

Appendix

Table 7: A detailed summary of 100 and Tiny-ImageNet: number of classes, image size, and dataset splits.

Method	Local Optimizer	Local LR	$\beta$	$\beta_{1}$	$\beta_{2}$	Weight Decay
FedAvg (Local SGD)	SGD	0.1	—	—	—	0.001
SCAFFOLD	SGD	0.1	—	—	—	0.001
FedCM	SGD	0.1	0.9	—	—	0.001
Local AdamW	AdamW	3e-4	—	0.9	0.999	0.01
Local Sophia	Sophia	3e-4	—	0.9	0.99	0.01
Local Muon	Muon	3e-2	—	0.9	0.95	0.01

Appendix

Table 8: Hyperparameter configuration of ResNet-18 and Vit-Tiny (CIFAR100, Tiny-ImageNet) across different algorithms.

Method	Local Optimizer	Local LR	$\beta$	$\beta_{1}$	$\beta_{2}$	Weight Decay
FedAvg (Local SGD)	SGD	0.1	—	—	—	0.001
SCAFFOLD	SGD	0.1	—	—	—	0.001
FedCM	SGD	0.1	0.9	—	—	0.001
Local AdamW	AdamW	1e-4	—	0.9	0.999	0.01
Local Sophia	Sophia	1e-4	—	0.9	0.99	0.01
Local Muon	Muon	1e-2	—	0.9	0.95	0.01

Appendix

Table 9: Hyperparameter configuration of ViT-Base fine-tuning across different algorithms.

Method	Local Optimizer	Local LR	$\beta$	$\beta_{1}$	$\beta_{2}$	Weight Decay
FedAvg (Local SGD)	SGD	0.1	—	—	—	0.001
SCAFFOLD	SGD	0.1	—	—	—	0.001
FedCM	SGD	0.1	0.9	—	—	0.001
Local AdamW	AdamW	3e-4	—	0.9	0.999	0.01
Local Sophia	Sophia	3e-4	—	0.9	0.99	0.01
Local Muon	Muon	3e-2	—	0.9	0.95	0.01

Appendix

Table 10: Hyperparameter configuration of LLAMA (C4) across different algorithms.

Method / Strategy	Communication	CommCost	Compute-Cost (s)	Acc(%)
FedAvg	$\|x\|$	22.8 MB	4.56 s	27.24
SCAFFOLD	$2\|x\|$	45.6 MB	5.22 s	26.86
FedCM	$\|x\|$	22.8 MB	4.68 s	16.95
Local AdamW	$\|x\|$	22.8 MB	4.89 s	37.57
Local Sophia	$\|x\|$	22.8 MB	4.92 s	34.05
Local Muon	$\|x\|$	22.8 MB	5.14 s	44.00

Appendix

Table 11: Per-round communication cost of different momentum aggregation strategies. Here

|x|

denotes the number of model parameters (in floats), and

\operatorname{CommCost}

is per-round communication cost, Compute-Cost is per-round computation time. (ViT-Tiny,

R=300

, Dir-0.1,

K=50

)

Method	CIFAR-100 (ResNet-18)				Tiny-Imagenet (ResNet-18)
Method	iid	Dir-0.5	Dir-0.1	Dir-0.05	iid	Dir-0.5	Dir-0.1	Dir-0.05
FedAvg	65.04	65.07	60.17	56.75	53.88	52.55	47.48	43.80
SCAFFOLD	66.04	65.51	60.69	56.43	54.49	53.36	47.76	43.92
FedCM	71.12	69.85	66.61	62.65	46.51	44.06	41.16	36.00
Local AdamW	64.37	63.58	59.23	55.24	51.06	49.97	44.01	40.00
Fed-Sophia	62.56	60.62	57.29	51.02	49.86	47.62	41.89	36.65

Appendix

Table 12: Test accuracy, training loss of each method on CIFAR-100 and Tiny-Imagenet using ResNet-18 over 300 communication rounds under IID Dir-0.5, Dir-0.1, Dir-0.05 (100 clients, 10% participation, batch size 50,

K=50

Method	CIFAR-100 (ViT-Tiny)				Tiny-Imagenet (ViT-Tiny)
Method	iid	Dir-0.5	Dir-0.1	Dir-0.05	iid	Dir-0.5	Dir-0.1	Dir-0.05
FedAvg	32.97	33.66	27.24	23.42	18.61	18.03	15.68	14.05
SCAFFOLD	32.29	33.01	26.86	23.23	18.32	18.36	15.70	14.21
FedCM	21.97	22.02	16.95	14.74	10.23	11.66	8.88	8.15
Local AdamW	41.27	41.21	37.57	36.06	27.94	26.82	24.31	21.35
Fed-Sophia	38.34	38.56	34.23	32.27	23.56	24.46	23.85	22.17

Appendix

Table 13: Test accuracy, training loss of each method on CIFAR-100 using ViT-Tiny over 300 communication rounds under IID, Dir-0.5, Dir-0.1, Dir-0.05 (100 clients, 10% participation, batch size 50,

K=50

Family	What is synchronized (server $\leftrightarrow$ clients)?	Extra cost (comm./state)	Main issue it targets
FedAvg / Local-SGD	Model parameters/updates $x$ only	$\approx\|x\|$ per round	General FL baseline; does not explicitly correct non-IID induced drift
FedOpt (server adaptive; e.g., FedAdam/FedYogi)	$x$ only; optimizer state (moments) kept on server (not shared as a common client geometry)	$\approx\|x\|$ (no extra client-side sync)	Improves global update scaling/adaptivity, but does not align heterogeneous client metrics
Control-variate / drift correction (e.g., SCAFFOLD)	$x$ + first-order control variates (gradient-level states)	typically larger (e.g., SCAFFOLD uses $\approx 2\|x\|$ in our comm accounting)	Corrects first-order client drift under non-IID via control variates; operates at gradient level, not metric level
Naive second-order FL (Local Sophia/SOAP/Muon + FedAvg agg.)	$x$ only; each client maintains local preconditioner state $\Theta$ but does not synchronize it	$\approx\|x\|$ per round	Fails under non-IID due to preconditioner drift (mismatched curvature-induced geometries)
FedPAC (ours)	$x$ + preconditioner state $\Theta$ (geometry / metric) : server aggregates $\Theta$ and broadcasts a global reference (Alignment); local steps are corrected using a global preconditioned direction (Correction)	$\approx\|x\|+c\|\Theta\|$ (e.g., $c=1$ for Muon/Sophia, $c=2$ for SOAP; light variants use low-rank upload such as $0.05\|\Theta\|$ – $0.1\|\Theta\|$ )	Targets preconditioner drift directly by synchronizing the metric (Alignment) and suppressing long-term drift (Correction)

Appendix

Table 14: Positioning of FedPAC. Unlike FedOpt and control-variate methods that synchronize only parameter/first-order states, FedPAC explicitly synchronizes the curvature-defined geometry (preconditioner state

\Theta

) to mitigate preconditioner drift under non-IID data. Here

|x|

denotes per-round model communication in our accounting and

|\Theta|

denotes the size of the optimizer preconditioner state.

Federated Learning

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

See all 923 papers

Taming the Instability: A Robust Second-Order Optimizer for Federated Learning over Non-IID Data
Mar 30, 2026Yuanqiao Zhang, Tiantian He, Yuan Gao +5Federated Learning Second-Order Optimization
Key Laboratory of Collaborative Intelligence Systems of Ministry of Education, Xidian University, Xi’an, Shaanxi, China · Centre for Frontier AI Research, Institute of High Performance Computing, Agency for Science, Technology and Research, Singapore · State Key Laboratory of ISN, School of Cyber Engineering, Xidian University, Xi’an, Shaanxi, China +2
FedMIX-P: Mixing Local and Global Preconditioners for Federated Vision and Language Model Training
Oct 1, 2026Junkang LiuFederated Learning Heterogeneous Federated Learning
Tianjin University
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

Mar 30, 2026J/K move · Enter open · S save

Taming the Instability: A Robust Second-Order Optimizer for Federated Learning over Non-IID Data

Yuanqiao Zhang, Tiantian He, Yuan Gao, Yixin Wang+4

Key Laboratory of Collaborative Intelligence Systems of Ministry of Education, Xidian University, Xi’an, Shaanxi, China · Centre for Frontier AI Research, Institute of High Performance Computing, Agency for Science, Technology and Research, Singapore · State Key Laboratory of ISN, School of Cyber Engineering, Xidian University, Xi’an, Shaanxi, China+2

Federated Learning Second-Order Optimization Non-IID Federated Learning Federated Learning Aggregation

Open paper AskPDF arXiv

In this paper, we present Federated Robust Curvature Optimization (FedRCO), a novel second-order optimization framework designed to improve convergence speed and reduce communication cost in Federated Learning systems under statistical heterogeneity. Existing second-order optimization methods are often computationally expensive and numerically unstable in distributed settings. In contrast, FedRCO addresses these challenges by integrating an efficient approximate curvature optimizer with a provable stability mechanism. Specifically, FedRCO incorporates three key components: (1) a Gradient Anomaly Monitor that detects and mitigates exploding gradients in real-time, (2) a Fail-Safe Resilience protocol that resets optimization states upon numerical instability, and (3) a Curvature-Preserving Adaptive Aggregation strategy that safely integrates global knowledge without erasing the local curvature geometry. Theoretical analysis shows that FedRCO can effectively mitigate instability and prevent unbounded updates while preserving optimization efficiency. Extensive experiments show that FedRCO achieves superior robustness against diverse non-IID scenarios while achieving higher accuracy and faster convergence than both state-of-the-art first-order and second-order methods.

Abstract

Federated Learning

Explore similar work

Taming the Instability: A Robust Second-Order Optimizer for Federated Learning over Non-IID Data

FedMIX-P: Mixing Local and Global Preconditioners for Federated Vision and Language Model Training

FedZMG: Efficient Client-Side Optimization in Federated Learning