FedMIX-P: Mixing Local and Global Preconditioners for Federated Vision and Language Model Training
cs.LG, cs.AI
Submitted: 2026-10-01
Updated: 2026-10-01
License: http://creativecommons.org/licenses/by/4.0/
The gist: Adaptive preconditioners accelerate model training, but heterogeneous client geometries can bias federated updates even when gradients are evaluated at the same model.
Terminology
Abstract
Adaptive preconditioners accelerate model training, but heterogeneous client geometries can bias federated updates even when gradients are evaluated at the same model. Round-start synchronization alone cannot prevent this mismatch from reappearing during local training. We propose FedMIX-P, which mixes shared and local preconditioners at every local step, retaining local adaptation while reducing mean-squared operator mismatch by a factor of λ squared. For smooth nonconvex objectives with stochastic gradients and partial participation, we establish an O(R-1/2) stationarity bound using suitable stepsizes and a horizon-dependent mixing weight, without requiring local preconditioners to converge to one another. A two-client counterexample shows that fixed positive mixing can preserve a nonstationary fixed point. The theory covers bounded linear symmetric positive-definite preconditioners. Experiments with SOAP, Sophia, and Muon variants across vision and language tasks show improvements over corresponding local optimizers, including accuracy gains of up to 19.47 percentage points and lower validation loss for 60M--350M language models. Full nonlinear and momentum-based updates require separate analysis.
Sources
- Measuring the Effects of Non-Identical Data Distribution for Federated Visual Classification
- FedMuon: Accelerating Federated Learning with Matrix Orthogonalization
- DP-FedPGN: Finding Global Flat Minima for Differentially Private Federated Learning via Penalizing Gradient Norm
- Taming Preconditioner Drift: Unlocking the Potential of Second-Order Optimizers for Federated Learning on Non-IID Data
- FedMuon: Federated Learning with Bias-corrected LMO-based Optimization
- LLaMA: Open and Efficient Foundation Language Models
- FedCM: Federated Learning with Client-level Momentum
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks