Taming Preconditioner Drift: Unlocking the Potential of Second-Order Optimizers for Federated Learning on Non-IID Data
Not on the map. Why?
Abstract (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 textttFedPAC, a preconditioner alignment and correction framework for reliable federated second-order optimization. textttFedPAC 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, textttFedPAC 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.
Key facts. Cited by 4 papers in the map (data through 7 Oct 2026). Builds on 15 of its 24 references that are in the map. First posted on arXiv 22 Feb 2026.
Builds on
15 of 24 references.
References · 15
- 2015Deep Residual Learning for Image Recognitioncited by 51,809
- 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 29,915
- 2019Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformercited by 14,491
- 2023LLaMA: Open and Efficient Foundation Language Modelscited by 13,958
- 2021Swin Transformer: Hierarchical Vision Transformer using Shifted Windowscited by 12,082
- 2018Federated Optimization in Heterogeneous Networkscited by 2,971
- 2017Decoupled Weight Decay Regularizationcited by 1,722
- 2019Measuring the Effects of Non-Identical Data Distribution for Federated Visual Classificationcited by 836
- 2023Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-trainingcited by 260
- 2024SOAP: Improving and Stabilizing Shampoo using Adamcited by 179
- 2021FedCM: Federated Learning with Client-level Momentumcited by 77
- 2025FedSWA: Improving Generalization in Federated Learning with Highly Heterogeneous Data via Momentum-Based Stochastic Controlled Weight Averagingcited by 25
3 more
- 2025The Potential of Second-Order Optimization for LLMs: A Study with Full Gauss-Newtoncited by 17
- 2024Fed-Sophia: A Communication-Efficient Second-Order Federated Learning Algorithmcited by 5
- 2025FedPM: Federated Learning Using Second-order Optimization with Preconditioned Mixing of Local Parameterscited by 3
Built on by
Citing papers per year: 2026: 4.
4 papers.
- 2026FedNSAM:Consistency of Local and Global Flatness for Federated Learningcited by 7
- 2026CTP-FL: Common-Trajectory Gradient Prediction for Federated Learningcited by 0
- 2026FedMIX-P: Mixing Local and Global Preconditioners for Federated Vision and Language Model Trainingcited by 0
- 2026FedLore: Communication and Memory Efficient Federated Learning via Shared Gradient Low-Rank Projectioncited by 0
- 2026FedLore: Communication and Memory Efficient Federated Learning via Shared Gradient Low-Rank Projection1 Oct 2026
- 2026FedMIX-P: Mixing Local and Global Preconditioners for Federated Vision and Language Model Training1 Oct 2026
- 2026CTP-FL: Common-Trajectory Gradient Prediction for Federated Learning28 Sep 2026
- 2026FedNSAM:Consistency of Local and Global Flatness for Federated Learning27 Feb 2026