Our new X account is live! Follow @wizwand_team for updates
WorkDL logo mark

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

About

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 \emph{preconditioner drift}: client-side second-order training induces heterogeneous \emph{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 \texttt{FedPAC}, a \emph{preconditioner alignment and correction} framework for reliable federated second-order optimization. \texttt{FedPAC} explicitly decouples parameter aggregation from geometry synchronization by: (i) \textbf{Alignment} (i.e.,aggregating local preconditioners into a global reference and warm-starting clients via global preconditioner); and (ii) \textbf{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, \texttt{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.

Junkang Liu, Fanhua Shang, Hongying Liu, Jin Liu, Weixin An, Yuanyuan Liu• 2026

Related benchmarks

TaskDatasetResultRank
Image ClassificationTiny ImageNet (test)
Accuracy87.92
265
Language ModelingC4--
73
Image ClassificationCIFAR-100 IID
Accuracy72.79
37
Image ClassificationTiny-ImageNet Dirichlet-0.05 (test)
Accuracy54
32
Image ClassificationTiny-ImageNet Dirichlet alpha=0.1 (test)
Test Accuracy57.95
30
Image ClassificationCIFAR-100 Dir-0.1
Accuracy71.85
28
Image ClassificationCIFAR-100 Dirichlet-0.1 (test)
Accuracy71.85
20
Image ClassificationCIFAR-100 Dirichlet-0.05 (test)
Accuracy65.56
20
Image ClassificationCIFAR-100 Dir-0.05
Accuracy65.56
12
Image ClassificationTiny-ImageNet Dir-0.1
Accuracy57.95
12
Showing 10 of 14 rows

Other info

Follow for update