FedPM: Federated Learning Using Second-order Optimization with Preconditioned Mixing of Local Parameters

AI summary

Overview

Research area: Federated Learning (FL), specifically distributed second-order optimization for training models across clients that never share their data.

Technical level: Advanced. The paper is built on matrix-level derivations of Newton-style update rules, a convergence theorem, and Hessian/Fisher-information approximations; readers will need comfort with gradients, Hessians, and preconditioning.

Scope: The paper derives a new second-order federated learning method (FedPM) that replaces the usual "simple averaging" of client parameters on the server with a curvature-weighted "preconditioned mixing" rule, and evaluates it theoretically and empirically.

What This Paper Is About

A family of federated learning methods (LocalNewton, LTDA, FedSophia) lets each client perform local second-order updates and then has the server simply average the resulting parameters. The paper shows analytically that this simple averaging produces a global update made of locally preconditioned local gradients, which is not the same as the ideal globally preconditioned global gradient — a mismatch that is especially damaging when client data are heterogeneous, because local curvature is a poor stand-in for global curvature.

The goal is to design a second-order parameter-mixing method whose client–server rules decompose exactly into the ideal global second-order update, so that local curvature drift no longer disrupts convergence.

Key Contributions

1. Derivation of a curvature-aware SOPM algorithm (FedPM). By algebraically decomposing the ideal global second-order update into a client step and a server step, the authors obtain a server operation they call preconditioned mixing — a curvature-weighted mixture of local parameters, rather than a plain average.
2. A theoretical convergence result. FedPM with a single local update (K = 1) attains a superlinear convergence rate for strongly convex objectives under strong-convexity and Hessian-smoothness assumptions, building on the proof strategy used for FedNL but using the most recent Hessian instead of a stale one.
3. A practical implementation for large models. Full Hessian computation is impractical for deep networks, so the authors adopt the FOOF approximation (Benzing 2022), which approximates the Fisher Information Matrix with block-diagonal, per-layer uncentered input covariance matrices.
4. Extensive empirical validation. Two test suites — a strongly convex logistic regression test and non-convex CNN/ResNet tests on image classification — show FedPM improving accuracy and convergence speed over first- and second-order baselines, particularly under strong data heterogeneity.

Main Findings

• Conventional SOPM is not global second-order optimization. Under K = 1, existing SOPM methods reduce to an average of locally preconditioned local gradients, whereas the ideal update uses a globally preconditioned global gradient. The paper identifies this discrepancy as the source of preconditioner drift under heterogeneity.
• FedPM is the only SOPM method listed that combines global second-order optimization with multiple local updates, as summarized in the paper's Table 1.
• Superlinear convergence is proven for a restricted setting: strongly convex objective, single local update (K = 1), full (non-stochastic) local gradients and Hessians, and an initial parameter implicitly assumed sufficiently close to the optimum. Multiple local updates are excluded from the analysis because preconditioner drift complicates a rigorous proof.
• FedPM requires extra communication. Unlike LocalNewton, LTDA, and FedSophia, clients must transmit preconditioners alongside parameters. With the FOOF approximation, communication drops from O(d²) (full preconditioner) to O(d), and construction cost from O(Md²) to O(Md).
• Test 1 (strongly convex, w8a and a9a): FedPM and FedNL significantly outperformed FedAvg/PSGD, FedAvgM, SCAFFOLD, FedAdam, FedNS, and LocalNewton, achieving smaller function-output gaps and smaller parameter distances; the accelerating decrease in parameter distance is consistent with the superlinear rate. FedNL and FedPM were noted to be equivalent in global updates at K = 1, with minor differences attributed to precision errors in their update procedures.
• Test 2 (CIFAR10 with a simple CNN, CIFAR100 with ResNet18 + GroupNorm), 5 local epochs, averaged over three seeds:
  • CIFAR10, α = 1.0: FedPM 74.8 ± 0.5 vs. LocalNewton 73.4 ± 1.4, FedAvgM 72.6 ± 1.5, SCAFFOLD 71.4 ± 0.5, FedAdam 70.4 ± 1.6, FedAvg 70.1 ± 1.3, FedProx 70.1 ± 1.1.
  • CIFAR10, α = 0.1: FedPM 68.6 ± 2.0 vs. SCAFFOLD 65.3 ± 2.7, FedAvgM 64.8 ± 3.0, FedAdam 64.6 ± 2.8, FedAvg 63.1 ± 2.1, FedProx 62.8 ± 1.3, LocalNewton 62.4 ± 2.3.
  • CIFAR100, α = 1.0: LocalNewton 70.2 ± 0.3, FedPM 68.4 ± 2.2, SCAFFOLD 64.7 ± 1.1, FedAvgM 61.2 ± 0.2, FedAvg 59.8 ± 0.7, FedAdam 58.3 ± 0.3, FedProx 56.6 ± 0.7.
  • CIFAR100, α = 0.1: FedPM 63.1 ± 0.8, LocalNewton 61.9 ± 0.6, SCAFFOLD 58.2 ± 0.6, FedAvgM 54.9 ± 0.6, FedAvg 52.4 ± 0.9, FedAdam 51.9 ± 0.4, FedProx 50.7 ± 0.9.
• FedPM leads in most settings, but not all: it wins in three of four configurations; LocalNewton is higher on CIFAR100 at α = 1.0 (70.2 vs. 68.4). LocalNewton's accuracy degrades sharply as heterogeneity increases, which the authors attribute to preconditioned mixing better aligning local and global updates.
• Convergence speed: on CIFAR10 at α = 0.1 with 5 local epochs, FedPM converged faster than other methods in both communication rounds and runtime.
• Robustness to local-epoch count: with a fixed total of 500 local epochs split as 1 epoch × 500 rounds, 5 epochs × 100 rounds, and 10 epochs × 50 rounds on CIFAR10 at α = 0.1, FedPM consistently outperformed the other methods.

Methodology in Plain English

The authors start from the update rule an ideal second-order method would use if all data were pooled on one machine: step in the direction of the global gradient, scaled by the inverse of the average Hessian. They then algebraically rewrite that single expression as an average over clients of a term that involves each client's own Hessian, the global (averaged) Hessian, and each client's locally updated parameter. That rewrite directly yields the algorithm: clients compute their local Hessian-based update, send both the updated parameter and the preconditioner, and the server weights each client's parameter by the interaction between the client's curvature and the global average curvature before summing.

For deep networks, computing and inverting a full Hessian is infeasible, so per-layer input covariance matrices (FOOF) stand in for the Hessian blocks. The server's mixing then becomes a per-layer operation using these matrices. The authors also compute the FOOF matrices only at the end of each round, just before communication, which they observed did not degrade performance versus updating them during local training.

To test the theory, they ran a deliberately simple strongly convex logistic-regression-with-L2 problem with full gradients, homogeneous data splits, and K = 1, so that the theorem's assumptions hold, and compared against seven baselines. Then they ran realistic non-convex image-classification tests with Dirichlet-distributed data across 10 clients to stress heterogeneity.

Why This Matters

Impact on research: The paper provides a clean diagnosis of why second-order parameter-mixing federated methods fail under heterogeneity — a mismatch between local and global preconditioning — and shows that the fix follows from an exact decomposition rather than a heuristic. It also clarifies the equivalence structure between gradient-mixing and parameter-mixing families, which for first-order methods holds at K = 1 but demonstrably does not for second-order methods.

Real-world applications:

• Cross-device model training on heterogeneous user data, such as keyboard prediction or photo classification where each user's data distribution differs sharply.
• Cross-silo settings such as hospitals training a shared diagnostic model without pooling patient records.
• Edge and IoT deployments where communication rounds are expensive and faster convergence in rounds translates directly into bandwidth and energy savings.
• Any federated deployment where data are non-IID by nature (regional or organizational splits), where the paper shows the largest gains.

Industry relevance: Communication rounds are a primary cost driver in production federated systems. FedPM's reported ability to reach target accuracy in fewer rounds, even accounting for the extra bytes of transmitting preconditioners, is directly relevant to operators. The released code (at the project's GitHub repository, rioyokotalab/fedpm) lowers the barrier to reproduction and adoption.

Future Directions

• Extending the convergence analysis beyond K = 1. The theorem covers only a single local update because preconditioner drift across multiple updates makes the proof intractable; a result for K > 1 remains open.
• Handling repeated preconditioner transmission. FedPM requires clients to send preconditioners in addition to parameters. Reducing or amortizing that overhead — for example by transmitting them less often — is not resolved here.
• Sampling and partial participation. Client sampling was not used in the main experiments, and the paper states its effects were explored only in the appendix, leaving the interaction with real-world client availability open.
• Scaling beyond the tested model sizes. SOGM baselines such as FedNL and FedNew were excluded from the DNN experiments because they were computationally prohibitive, so how FedPM behaves at substantially larger model scales is not established.

Target Audience

This paper is best suited to machine learning researchers and graduate students working on federated optimization, distributed optimization, or second-order and preconditioned methods, as well as to practitioners building federated systems on non-IID data who are willing to trade additional uplink communication for fewer training rounds.

Authors’ abstract

We propose Federated Preconditioned Mixing (FedPM), a novel Federated Learning (FL) method that leverages second-order optimization. Prior methods--such as LocalNewton, LTDA, and FedSophia--have incorporated second-order optimization in FL by performing iterative local updates on clients and applying simple mixing of local parameters on the server. However, these methods often suffer from drift in local preconditioners, which significantly disrupts the convergence of parameter training, particularly in heterogeneous data settings. To overcome this issue, we refine the update rules by decomposing the ideal second-order update--computed using globally preconditioned global gradients--into parameter mixing on the server and local parameter updates on clients. As a result, our FedPM introduces preconditioned mixing of local parameters on the server, effectively mitigating drift in local preconditioners. We provide a theoretical convergence analysis demonstrating a superlinear rate for strongly convex objectives in scenarios involving a single local update. To demonstrate the practical benefits of FedPM, we conducted extensive experiments. The results showed significant improvements with FedPM in the test accuracy compared to conventional methods incorporating simple mixing, fully leveraging the potential of second-order optimization.