FedMIX-P: Mixing Local and Global Preconditioners for Federated Vision and Language Model Training

Motivation

• At a shared model, preconditioner mismatch coupled with gradient heterogeneity can bias the aggregate direction, even when the global gradient is zero.Round-start synchronization does not prevent local adaptation from recreating mismatch during a round.Source: §1 · p. 2

Theory

• With suitable stepsizes and horizon-dependent mixing, the analysis gives an O(R^-1/2) stationarity bound without requiring local preconditioners to converge to one another.The bound accounts for stochastic gradients, partial participation, and a mismatch contribution scaling with λ^2 times bounded local preconditioner drift.Source: §1 · p. 2

Experiments

• Coordinating local and shared geometry improves on corresponding local optimizers across vision and language tasks, with particularly large gains under severe data heterogeneity.The theory covers bounded linear symmetric positive-definite preconditioners; full nonlinear and momentum-based updates require separate analysis.Source: §1 · p. 2

Contributions

• The paper identifies geometry–heterogeneity coupling as a source of aggregation bias distinct from trajectory drift caused by multiple local steps.Its exact decomposition separates the two effects.Source: §1 · pp. 2–3

Method

• FedMIX-P mixes a shared reference with each client’s adaptive preconditioner at every local step, retaining local adaptation while controlling applied operator mismatch.This targets the geometry actually used for updates, even when local preconditioner states do not converge to one another.Source: §1 · pp. 2–3

Heterogeneity-aware federated optimization

• FedAvg, FedProx, SCAFFOLD, and FedCM primarily regulate local trajectories or gradient-based directions, whereas differing preconditioners can create mismatch even at a shared model.FedMIX-P addresses this operator–gradient coupling by controlling the preconditioners applied during local training.Source: §2 · p. 3

Structure-aware optimizers

• SOAP, Sophia, and Muon use distinct geometric transformations, and Muon’s nonlinear orthogonalization should not be identified with a fixed linear positive-definite preconditioner.These optimizer-specific rules delimit which implementations are covered by a unified linear-preconditioner analysis.Source: §2 · p. 3

Preconditioning and matrix optimization in FL

• Fed-Sophia, FedPM, and federated Muon studies establish client-specific update geometry as a federated optimization challenge, with FedPM’s theory focused on strongly convex single-step settings.The cited federated Muon studies investigate matrix orthogonalization and bias of locally aggregated directions.Source: §2 · p. 3

FedMIX-P

• FedMIX-P mixes shared and local preconditioners at every client step and obtains a vanishing stationarity bound for multiple local steps, stochastic gradients, and partial participation under bounded local-state drift.The guarantee is for linear positive-definite operators and does not automatically extend to nonlinear Muon orthogonalization or complete SOAP and Sophia updates.Source: §2 · p. 3

3.1 PROBLEM SETUP

• The setup frames federated learning as collaborative minimization of population risk over clients, with client losses and stochastic samples specified in the procedure.The supplied setup passages also identify the global model, clients, rounds, and local steps as algorithm inputs.Source: §3.1 · pp. 3–4

3.2 FEDERATED PRECONDITIONED OPTIMIZATION

• FedPO directly applies client-side preconditioned optimizers and averages model changes, while clients retain optimizer states that are not synchronized.Consequently, the global direction combines gradients transformed by potentially different client geometries.Source: §3.2 · p. 4

3.3 PRECONDITIONER DRIFT AND AGGREGATION BIAS

• FedMIX-P starts clients from a shared reference and mixes local and shared preconditioners so the applied-operator drift is controlled by the mixing weight.The aggregation decomposition distinguishes geometry–gradient coupling from trajectory error, which arises when multiple local steps evaluate gradients at different client models.Source: §3.3 · pp. 4–5

4 FEDERATED MIXING OF LOCAL AND GLOBAL PRECONDITIONERS

• FedMIX-P combines a shared preconditioner reference with locally updated preconditioners through coupled operations.The server broadcasts the shared reference; clients begin aligned while retaining the ability to adapt their states to local data.Source: §4 · p. 6
• Clients upload model differences and terminal preconditioner states, which the server averages to form the next shared reference.With γr = Kηl, direct averaging of returned client models is recovered.Source: §4 · p. 6
• At each local step, clients construct the operator for the current gradient from shared and local preconditioners, controlling mismatch without requiring local states to converge.The current gradient updates the local state for the next step; in the predictable-operator case, the direction is conditionally unbiased with respect to the current mini-batch.Source: §4 · p. 6
• With ηl,r = Θ(1/(K√R)) and λr = Θ(R^-1/2), FedMIX-P attains an O(R^-1/2) bound on average expected squared global-gradient norm under bounded local-state drift.The guarantee is stated for an R-round run.Source: §4 · p. 6
• The guarantee covers linear, symmetric positive-definite preconditioners; nonlinear Muon orthogonalization and complete SOAP and Sophia updates require separate analysis.Source: §4 · p. 6

5 THEORETICAL ANALYSIS

• The analysis covers nonconvex objectives, stochastic local gradients, uniform partial participation, and bounded linear symmetric positive-definite preconditioners.The objective is smooth and lower bounded; the analysis also uses bounded local preconditioner drift.Source: §5 · p. 7
• With u = Θ(R^-1/2) and λ = Θ(R^-1/2), FedMIX-P obtains GR = O(R^-1/2) when local-state drift is bounded.The mixing schedule makes the systematic preconditioner term O(R^-1).Source: §5 · pp. 7–8
• FedPO’s preconditioning-bias term need not vanish as the learning rate decreases, whereas FedMIX-P controls the term through λ² times bounded local-state drift.Round-start alignment removes initial mismatch, and stepwise mixing controls the mismatch of the operators applied to gradients.Source: §5 · p. 8
• A fixed positive mixing weight does not generally provide the same guarantee and can admit a nonstationary fixed point.Source: §5 · p. 8
• The result establishes a stationarity-guarantee advantage, not uniform gains in convergence speed or wall-clock time; communication costs and shared-reference quality need empirical assessment.Nonlinear orthogonalization and full optimizer-specific state updates require separate analysis.Source: §5 · p. 8

6 EXPERIMENTS

• Experiments span CIFAR-100, Tiny-ImageNet, and C4, using Dirichlet partitions where smaller α means greater data heterogeneity.The evaluated models include ResNet-18, ViT-Base, ViT-Tiny, and LLaMA, with FedMIX-P variants of SOAP, Sophia, and Muon.Source: §6 · pp. 8–9

6.3 C4 FEDERATED PRE-TRAINING WITH LLAMA

• FedMIX-P SOAP achieves the lowest validation loss at all three LLaMA sizes (3.314/3.306/3.228) and lowest validation perplexity (27.49/27.28/25.23).Results are after 100 communication rounds with 20 clients and 20% participation per round.Source: Table 2 · p. 9, §6.3 · p. 10

6.1 RESULTS ON CONVOLUTIONAL NEURAL NETWORKS

• FedMIX-P improves all eight ResNet-18 comparisons; on Tiny-ImageNet, Muon gains 5.84/19.47 points over Local Muon under Dir-0.1/Dir-0.05.On CIFAR-100, FedMIX-P Muon’s advantage over Local Muon grows from 5.39 to 17.01 points across the same partitions.Source: §6.1 · pp. 9–10

6.2 RESULTS ON VISION TRANSFORMER

• On ViT-Tiny, FedMIX-P beats local counterparts in all 12 comparisons; on pretrained ViT-Base, all three variants improve mean accuracy on both datasets under Dir-0.1.For ViT-Tiny, SOAP’s CIFAR-100 advantage grows from 2.15 to 6.77 points across partitions; ViT-Base gains are described as modest.Source: §6.2 · p. 10

6.4 ABLATION STUDY

• The mixing weight controls reliance on evolving local preconditioners, with λ = 0 and λ = 1 representing shared-only and local-only preconditioning after synchronization.The reported state-drift metric does not by itself determine predictive accuracy.Source: Table 3 · p. 9, §6.4 · p. 10
• On CIFAR-100 under Dir-0.05, mixed SOAP reaches 65.18%/48.45% on ResNet-18/ViT-Tiny, exceeding shared-only and synchronized local-only endpoints on both.Synchronized local-only SOAP also exceeds Local SOAP, which retains client-specific state across rounds.Source: Table 4 · p. 11, §6.4 · p. 11
• Compressed FedMIX-P SOAP-light reaches 47.89% ViT-Tiny accuracy using 25.08 MB per round, retaining all but 0.56 points of full FedMIX-P SOAP’s gain.Full FedMIX-P SOAP uses 45.6 MB per round versus 22.8 MB for Local SOAP; both FedMIX-P variants take 5.67 s per round versus 5.56 s for Local SOAP.Source: Table 4 · p. 11, §6.4 · p. 11

• FedMIX-P mixes a shared reference with each client’s adaptive preconditioner at every local step, retaining local adaptation while controlling mismatch between applied operators.The mismatch can bias aggregation even when gradients are evaluated at the same model.Source: §7 · p. 11
• For bounded linear symmetric positive-definite preconditioners, horizon-dependent mixing yields a vanishing stationarity bound without requiring local preconditioners to converge to one another.The analysis bounds the geometric-mismatch contribution by a term proportional to λ^2 times bounded preconditioner drift.Source: §7 · p. 11
• Vision and language experiments show improvements over corresponding local optimizers, particularly under severe data heterogeneity.Extending the analysis to full nonlinear and momentum-based optimizer updates remains open.Source: §7 · p. 11

8 CONVERGENCE OF FEDPO AND FEDMIX-P

• Under A.1–A.5 and predictable applied operators, Theorem 8.4 gives both FedPO and FedMIX-P a partial-participation stationarity bound.The theorem’s conditions include the stated stepsize restrictions and its parameters a, b, and V.Source: §8 · p. 16
• For FedMIX-P, with all problem parameters, including K and S, fixed independently of R, the three stationarity-bound terms scale as O(R^-1/2), O(R^-1), and O(R^-1/2) for sufficiently large R.Source: §8 · p. 17

8.1 CURRENT-SAMPLE-DEPENDENT PRECONDITIONERS

• For FedMIX-P, the schedule also yields O(R^-1/2) expected stationarity under A.1–A.5 alone, without requiring operator predictability.This result permits applied operators to depend on the current mini-batch, whose correlation with the gradient noise is retained in the analysis.Source: §8.1 · pp. 17–18
• For FedPO, the predictable-operator comparison allows arbitrary bounded preconditioner drift, and its stepsize restrictions do not depend on D0.The comparison is also extended to current-sample-dependent operators under A.1–A.5 using separate bias and noise bounds.Source: §8.1 · pp. 18–19
• For every fixed λ > 0, a deterministic full-participation two-client FedMIX-P instance satisfying A.1–A.5 has a nonstationary fixed point; λ = O(R^-1/2) removes this particular obstruction at that rate.At the constructed fixed point, the update direction is zero although the gradient is nonzero.Source: §8.1 · p. 19
• Setting γr = Kηr recovers direct model averaging, and the operator update must preserve the theorem’s spectral bounds.Source: §8.1 · p. 21

8.2 FEDMIX-P AND OPTIMIZER-SPECIFIC IMPLEMENTATIONS

• Practical implementations initialize participating clients from a shared server state, while local baselines retain their own states and omit shared-state communication.The specified protocol uses complete state exchange; compressed variants require an additional explicit rule.Source: §8.2 · p. 21

8.2.1 SOAP: MIXING INDUCED LINEAR GEOMETRY

• SOAP mixes reference and local preconditioners on the same current momentum matrix, then updates geometry; second moments must be transported before common-coordinate aggregation.Averaging second moments without transporting their coordinates is not valid common-coordinate aggregation.Source: §8.2.1 · pp. 21–22

8.2.2 SOPHIA: MIXING INVERSE CURVATURE BEFORE CLIPPING

• Sophia mixes inverse-curvature scales before applying clipping once, rather than inverting an averaged curvature estimate.Source: §8.2.2 · p. 22
• Sophia-G requires specified predictive-label sampling and loss/batch normalization for its Gauss–Newton–Bartlett estimator; empirical squared gradients on observed labels are not a substitute.Source: §8.2.2 · p. 23

8.2.3 MUON: MIXING CURRENT MOMENTUM-CONDITIONED MAPS

• Muon forms shared and local momentum-conditioned maps from the same current gradient, rather than reusing an old orthogonalized direction.The practical Muon map is nonlinear and momentum-conditioned, so it is among the variants requiring additional convergence analysis.Source: §8.2.3 · pp. 23–24

8.2 FEDMIX-P AND OPTIMIZER-SPECIFIC IMPLEMENTATIONS

• Algorithms 4–7 are practical extensions, not covered by the linear-SPD theorem; their convergence requires additional analysis because their state aggregation generally does not reproduce exact terminal-operator aggregation.Theorem-matched specialization requires current gradients, no clipping or decoupled weight decay, uniform spectral bounds, and exact aggregation of terminal applied operators.Source: §8.2.3 · p. 24

Communication-efficient federated optimization

• CTP-FL communicates one model-sized vector per participating client and uses a shared predictive trajectory, unlike FedBCGD’s parameter-block communication approach.FedBCGD also incorporates drift control and variance reduction in its accelerated variant (Liu et al., 2024b).Source: §9 · p. 25

Alignment under heterogeneous data

• FedSWA and FedMoSWA target generalization, while FedNSAM targets local–global flatness consistency; this work focuses on gradient-evaluation geometry at common points.At common evaluation points, averaging client gradients estimates the global-objective gradient there, regardless of differences among individual client gradients (Liu et al., 2025a; Liu et al., 2025b).Source: §9 · p. 25

Federated adaptive and structured optimizers

• FedAdamW, FedMuon, and FedPAC use optimizer-specific mechanisms, whereas CTP-FL aligns gradient-evaluation locations without transmitting moments or preconditioners.The cited mechanisms include local correction and second-moment aggregation, matrix orthogonalization, and preconditioner alignment with update correction (Liu et al., 2026b; Liu et al., 2025c; Liu et al., 2026a).Source: §9 · p. 25

Global flatness and privacy

• DP-FedPGN encourages globally flat solutions in client-level differentially private federated learning through a global gradient-norm penalty, but uses a different objective and privacy accounting (Liu et al., 2025d).It illustrates that a global-objective quantity need not be faithfully represented by independently optimized local objectives.Source: §9 · p. 25