Linear Convergence and Generalization of FedAvg Under Constrained PL-Type Assumptions: A Single Hidden Layer Neural Network Analysis

TMLR Paper6597 Authors

Abstract:

In this work, we study the generalization performance of the widely adopted FedAvg algorithm for solving Federated Learning (FL) problems. FedAvg has been extensively studied from an optimization perspective under different settings; however, analyzing the generalization performance of FedAvg is particularly challenging under practical settings since it involves simultaneously bounding (1) the optimization error and (2) the Rademacher complexity of the model to be learned, which are often contradictory. Specifically, obtaining optimization guarantees for FedAvg relies on restrictive assumptions on the loss landscape, such as (strong) convexity or Polyak-{\L}ojasiewicz (PL) inequality to be satisfied over the entire parameter space. However, for an unbounded space, it is challenging to control the Rademacher complexity, leading to worse generalization guarantees. In this work, we address this challenge by proposing novel {\em constrained PL-type} conditions on the {objective function} that ensure the existence of a global optimal to which {FedAvg converges} linearly after $\mathcal{O}( \log ({1}/{\epsilon}))$ rounds of communication, where $\epsilon$ is the desired optimality gap. Importantly, we demonstrate that a class of single hidden layer neural networks satisfies the proposed {\em constrained PL-type} conditions
% required to establish the linear convergence of FedAvg
as long as $m > {nK}/{d}$, where $m$ is the width of the neural network, $K$ is the number of clients, $n$ is the number of samples at each client, and $d$ is the feature dimension. Finally, we bound the Rademacher complexity for this class of neural networks and establish that the generalization error of FedAvg diminishes at the rate of $\mathcal{O}({1}/{\sqrt{n}})$.

Submission Type: Regular submission (no more than 12 pages of main content)

Assigned Action Editor: ~Sebastian_U_Stich1

Submission Number: 6597