[Paper Review] FedADMM: A Federated Primal-Dual Algorithm Allowing Partial Participation

Han Wang, Siddartha Marella+1 more|arXiv (Cornell University)|Mar 28, 2022

Privacy-Preserving Technologies in Data4 citations

TL;DR

FedADMM is a novel federated primal-dual algorithm that enables partial participation in non-convex composite optimization, extending FedPD to allow flexible client participation while maintaining convergence. It achieves $Ó(\varepsilon^{-2})$ communication complexity and proves equivalence to FedDR, ensuring robustness to statistical and system heterogeneity with theoretical convergence guarantees under general sampling models.

ABSTRACT

Federated learning is a framework for distributed optimization that places emphasis on communication efficiency. In particular, it follows a client-server broadcast model and is particularly appealing because of its ability to accommodate heterogeneity in client compute and storage resources, non-i.i.d. data assumptions, and data privacy. Our contribution is to offer a new federated learning algorithm, FedADMM, for solving non-convex composite optimization problems with non-smooth regularizers. We prove converges of FedADMM for the case when not all clients are able to participate in a given communication round under a very general sampling model.

Motivation & Objective

• To address the challenge of partial client participation in federated learning, particularly in systems with stragglers or unreliable connectivity.
• To develop a primal-dual algorithm that supports non-convex, composite optimization problems with non-smooth regularizers, enabling broader applicability than existing methods.
• To extend FedPD, which requires full participation, to allow partial participation while preserving convergence and performance.
• To establish theoretical equivalence between FedADMM and FedDR, enabling shared convergence analysis and algorithmic insights.
• To provide convergence guarantees under general sampling models, including non-i.i.d. data and unbalanced client participation.

Proposed method

• FedADMM is derived by applying the dual formulation of the non-convex Douglas-Rachford splitting (DRS) to the federated optimization problem, leveraging the known equivalence between DRS and ADMM.
• The algorithm uses a client-server architecture where only a subset of clients participate per round, with the server aggregating updates via a dual consensus mechanism.
• It employs proximal operators for non-smooth regularizers and allows inexact local solves, similar to FedProx and FedPD, enhancing practicality.
• The method uses a uniform sampling scheme for client selection and maintains dual variables to enforce consensus across clients.
• Key components include the dual variable update, primal variable update via local optimization, and a global consensus step on the server.
• The algorithm is designed to inherit the convergence properties of FedDR by establishing a one-to-one, onto mapping between FedDR and FedADMM iterates.

Experimental results

Research questions

• RQ1 Can a primal-dual federated algorithm be designed to allow partial client participation while maintaining convergence for non-convex composite problems?
• RQ2 How does FedADMM relate to FedDR and FedPD in terms of algorithmic structure and convergence behavior?
• RQ3 What is the communication complexity of FedADMM under partial participation, and does it match known lower bounds?
• RQ4 Can FedADMM handle non-smooth regularizers and constraints, extending beyond the scope of FedPD?
• RQ5 Does FedADMM achieve the same convergence rate as FedDR when clients participate partially?

Key findings

• FedADMM achieves $Ó(\varepsilon^{-2})$ communication complexity to reach an $µ$-stationary solution, matching the known lower bound for non-convex federated optimization.
• FedADMM is theoretically equivalent to FedDR, with a one-to-one and onto mapping between their iterates, ensuring identical convergence behavior.
• When $g \equiv 0$, FedADMM reduces to FedPD but allows partial participation, overcoming a key limitation of FedPD.
• The algorithm maintains convergence under general sampling models, including non-i.i.d. data and unbalanced client participation.
• Numerical simulations on synthetic and FEMNIST datasets confirm that FedADMM and FedDR achieve identical training accuracy and loss, validating the theoretical equivalence.
• The convergence guarantee holds even when local proximal operators are solved inexact, enhancing practical robustness.