Taming Preconditioner Drift: Unlocking the Potential of Second-Order Optimizers for Federated Learning on Non-IID Data
Abstract
Domain fit: AI-adjacent · Paper appears method- or tooling-adjacent to AI workflows with partial ecosystem coverage.
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.
Results and benchmarks
Second-order optimizers can significantly accelerate large-scale training, yet their naive federated variants are often unstable or even diverge on non-IID data.
| Task | Dataset | Metric | Value | Source |
|---|---|---|---|---|
| Computer vision | FedAvg | ResNet-18 | 60.17 | paper-derived |
| Computer vision | 47.48 | ResNet-18 | 43.80 | paper-derived |
| Computer vision | 23.42 | ResNet-18 | 15.68 | paper-derived |
Audit each benchmark finding before selecting an implementation path. Evidence refs map to the disclosure below.
Evidence graph: 2 refs, 1 links.
Utility signals: depth 90/100, grounding 68/100, status medium.
Implementation
No direct implementation yet
Maintained implementation evidence is not confirmed for this paper yet.
Use the implementation status and reproduction sections for the current action plan.
No verified maintained repo yet
There is no verified maintained implementation yet. Use this baseline plan to decide whether to prototype now or defer.
- This is primarily a method paper. Reproduce it within a maintained framework baseline instead of chasing paper-specific repos.
- Start with framework-native implementations (e.g. PyTorch optimizer module, Optax, or Transformers training loops).
- Replicate the paper ablation settings first, then compare against modern baselines.
Time to first repro: a few hours
This is primarily a method paper. Reproduce it within a maintained framework baseline instead of chasing paper-specific repos.
- No maintained paper-verified implementation is currently available
Reproduction readiness
No repo
No verified implementation available
- No maintained repository has been identified for this paper. Check adjacent implementations or HF artifacts below.
Hugging Face artifacts
No trustworthy direct or curated related Hugging Face artifacts were found yet. Use targeted searches to quickly locate candidate models, datasets, and demos.
Models
Datasets
Tip: start with models, then check datasets and spaces if you need evaluation data or demos.
Research context
Tasks
Computer vision
Methods
None detected
Domains
Computer vision
Open this paper in HFEPX to review benchmark signals, evaluation modes, and human-feedback protocol context.
Open in HFEPXJump to Paper2Code search queries derived from this paper's research context.