Distributed optimization: designed for federated learning

TL;DR

This paper introduces a flexible distributed optimization framework tailored for federated learning, capable of handling various communication structures and providing theoretical convergence guarantees, with demonstrated effectiveness in large-scale, heterogeneous environments.

Contribution

It develops a unified augmented Lagrangian-based optimization framework for federated learning, generalizing classical methods and ensuring convergence under diverse settings.

Findings

Strong convergence guarantees for the proposed algorithms.

Effective performance in large-scale, heterogeneous federated learning scenarios.

Generalization of classical optimization methods within the framework.

Abstract

Federated learning (FL), as a distributed collaborative machine learning (ML) framework under privacy-preserving constraints, has garnered increasing research attention in cross-organizational data collaboration scenarios. This paper proposes a class of distributed optimization algorithms based on the augmented Lagrangian technique, designed to accommodate diverse communication topologies in both centralized and decentralized FL settings. Furthermore, we develop multiple termination criteria and parameter update mechanisms to enhance computational efficiency, accompanied by rigorous theoretical guarantees of convergence. By generalizing the augmented Lagrangian relaxation through the incorporation of proximal relaxation and quadratic approximation, our framework systematically recovers a broad of classical unconstrained optimization methods, including proximal algorithm, classic gradient descent, and stochastic gradient descent, among others. Notably, the convergence properties of these methods can be naturally derived within the proposed theoretical framework. Numerical experiments demonstrate that the proposed algorithm exhibits strong performance in large-scale settings with significant statistical heterogeneity across clients.