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

Yuanqiao Zhang; Tiantian He; Yuan Gao; Yixin Wang; Yew-Soon Ong; Maoguo Gong; A. K. Qin; Hui Li

cs.LGMar 30, 2026

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

Authors: Yuanqiao Zhang, Tiantian He, Yuan Gao, Yixin Wang, Yew-Soon Ong, Maoguo Gong, A. K. Qin, Hui Li

Organizations: 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 · College of Computing and Data Science, Nanyang Technological University, Singapore ·

Abstract

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.

Figures & tables

	CIFAR-10						EMNIST
Method	Dirichlet			Pathological		IID	Dirichlet			Pathological		IID
Method	$\alpha=0.1$	$\alpha=0.5$	$\alpha=1$	2	5	\	$\alpha=0.1$	$\alpha=0.5$	$\alpha=1$	10	30	\
FedAvg	0.563	0.620	0.632	0.537	0.601	0.650	0.818	0.835	0.836	0.809	0.826	0.838
FedProx	0.553	0.628	0.637	0.530	0.616	0.631	0.817	0.831	0.833	0.801	0.822	0.839
FedAdam	0.557	0.639	0.610	0.439	0.623	0.619	0.845	0.856	0.859	0.804	0.837	0.864
FedAvgM	0.556	0.614	0.635	0.530	0.607	0.658	0.814	0.832	0.832	0.800	0.823	0.836

Table 1: The experimental results on the Dirichlet-non, Pathological-non, and IID settings with client number 100, and the party ratio 0.8.

	CIFAR-10						EMNIST
Method	Client Number			Party Ratio			Client Number			Party Ratio
Method	10	50	100	0.1	0.5	1	10	50	100	0.1	0.5	1
FedAvg	0.539	0.521	0.563	0.538	0.576	0.565	0.779	0.814	0.818	0.806	0.814	0.815
FedProx	0.557	0.573	0.553	0.540	0.575	0.554	0.781	0.810	0.817	0.807	0.819	0.815
FedAdam	0.456	0.539	0.557	0.397	0.549	0.525	0.795	0.834	0.845	0.840	0.847	0.846
FedAvgM	0.517	0.570	0.556	0.555	0.559	0.551	0.781	0.816	0.814	0.811	0.819	0.812

Table 2: The ablation experimental results about client number and participation ratio on the Dirichlet distribution with

Dir(\alpha)=0.1

. For the client number part, the party ratio is set as 0.8, and for the party ratio part, the client number is set as 100.

Figure 4: Visualization of gradient stability. The plot compares the gradient trajectories with (red) and without (blue/orange) the Gradient Monitor.

Appendix figures & tables23 assets

Supplementary material from the paper’s appendix.

Appendix

Notation	Description mainly used in the text
$\bm{\theta}$	Model parameter, $\bm{\theta}\in\mathbb{R}^{d}$ .
$\mathcal{L}(\bm{\theta}),\mathcal{L}_{c}(\bm{\theta})$	Global loss function, local loss function of client $c$ .
$C,c$	Total number of clients, client index.
$T,t$	Total number of communication rounds, communication round index.
$K,k$	Total number of local epochs, local epoch index.
$n,n_{c}$	Total data volume, local data volume, $n=\sum_{c=1}^{C}n_{c}$ .

Appendix

Table 3: Notations used in this paper

Figure 5: The distribution example. The first figure is the Dirichlet distribution on 100 clients with $Dir(\alpha)=0.1$ , and the second figure is the Pathological distribution on 100 clients with 2 classes per client.

Figure 6: Dataset: EMNIST, Data distribution: Dirichlet 0.1, Party ratio: 0.1, Number clients: 100.

Figure 7: Dataset: EMNIST, Data distribution: Dirichlet 0.1, Party ratio: 0.5, Number clients: 100.

Figure 8: Dataset: EMNIST, Data distribution: Dirichlet 0.1, Party ratio: 0.8, Number clients: 100.

Figure 9: Dataset: EMNIST, Data distribution: Dirichlet 0.1, Party ratio: 1, Number clients: 100.

Figure 10: Dataset: EMNIST, Data distribution: Dirichlet 0.1, Party ratio: 0.8, Number clients: 10.

Figure 11: Dataset: EMNIST, Data distribution: Dirichlet 0.1, Party ratio: 0.8, Number clients: 50.

Figure 12: Dataset: EMNIST, Data distribution: Dirichlet 0.5, Party ratio: 0.8, Number clients: 100.

Figure 13: Dataset: EMNIST, Data distribution: Dirichlet 1, Party ratio: 0.8, Number clients: 100.

Figure 14: Dataset: EMNIST, Data distribution: Pathological 10, Party ratio: 0.8, Number clients: 100.

Figure 15: Dataset: EMNIST, Data distribution: Pathological 30, Party ratio: 0.8, Number clients: 100.

Figure 16: Dataset: EMNIST, Data distribution: iid, Party ratio: 0.8, Number clients: 100.

Figure 17: Dataset: CIFAR-10, Data distribution: Dirichlet 0.1, Party ratio: 0.1, Number clients: 100.

Figure 18: Dataset: CIFAR-10, Data distribution: Dirichlet 0.1, Party ratio: 0.5, Number clients: 100.

Figure 19: Dataset: CIFAR-10, Data distribution: Dirichlet 0.1, Party ratio: 1, Number clients: 100.

Figure 20: Dataset: CIFAR-10, Data distribution: Dirichlet 0.1, Party ratio: 0.8, Number clients: 10.

Figure 21: Dataset: CIFAR-10, Data distribution: Dirichlet 0.1, Party ratio: 0.8, Number clients: 50.

Figure 22: Dataset: CIFAR-10, Data distribution: Dirichlet 0.5, Party ratio: 0.8, Number clients: 100.

Figure 23: Dataset: CIFAR-10, Data distribution: Dirichlet 1, Party ratio: 0.8, Number clients: 100.

Figure 24: Dataset: CIFAR-10, Data distribution: Pathological 2, Party ratio: 0.8, Number clients: 100.

Figure 25: Dataset: CIFAR-10, Data distribution: Pathological 5, Party ratio: 0.8, Number clients: 100.

Figure 26: Dataset: CIFAR-10, Data distribution: iid, Party ratio: 0.8, Number clients: 100.

Federated Learning

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

See all 923 papers

Taming Preconditioner Drift: Unlocking the Potential of Second-Order Optimizers for Federated Learning on Non-IID Data
Feb 22, 2026Junkang Liu, Fanhua Shang, Hongying Liu +3Federated Learning Second-Order Optimization
College of Intelligence and Computing, Tianjin University, Tianjin, China · Medical School, Tianjin University, Tianjin, China · Peng Cheng Lab, Shenzhen, China +2
Fractional-Order Federated Learning
Feb 17, 2026Mohammad Partohaghighi, Roummel Marcia, YangQuan ChenFederated Learning Federated Learning Aggregation
Electrical Engineering and Computer Science, University of California at Merced, USA · Department of Applied Mathematics, University of California Merced, Merced, CA, USA · Mechatronics, Embedded Systems and Automation (MESA) Lab, Department of Mechanical Engineering, School of Engineering, University of California, Merced, CA 95343, USA
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

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

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+2

College of Intelligence and Computing, Tianjin University, Tianjin, China · Medical School, Tianjin University, Tianjin, China · Peng Cheng Lab, Shenzhen, China+2

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

Open paper AskPDF arXiv

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.

Abstract

Federated Learning

Explore similar work

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

Fractional-Order Federated Learning

FedZMG: Efficient Client-Side Optimization in Federated Learning