Mixing local and global preconditioners improves federated model training
FedMIX-P: Mixing Local and Global Preconditioners for Federated Vision and Language Model Training
Machine LearningArtificial Intelligence
Summary
Training large AI models on many devices at once can be tricky because each device learns a bit differently, which can slow things down or lead to poorer results. The authors propose a new way to combine shared and local information during training steps to keep the benefits of learning locally while reducing mismatches between devices. They prove this method converges under certain assumptions and show it improves performance on vision and language tasks, sometimes very significantly. This approach helps train big models more efficiently across different devices without requiring all devices to behave the same way.
What this means in practice
- •For distributed ai engineers: Improve federated training efficiency and accuracy by adaptively mixing local and shared optimization steps to reduce update mismatches across devices.
- •For natural language processing teams: Train large language models over distributed datasets with better stability and faster convergence by using mixed local and global preconditioners.
Authors
Junkang Liu
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 \texttt{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 $λ^2$. 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.