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

TL;DR

FedPM introduces a second-order federated learning method that uses preconditioned mixing of local parameters to improve convergence and accuracy, especially in heterogeneous data environments.

Contribution

It proposes a novel second-order FL algorithm with preconditioned mixing, addressing preconditioner drift and enhancing convergence over prior methods.

Findings

FedPM achieves higher test accuracy than existing methods.

Theoretical convergence rate is superlinear for strongly convex objectives.

Extensive experiments validate practical improvements in heterogeneous settings.

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.