Generalization of FedAvg Under Constrained Polyak-Lojasiewicz Type Conditions: A Single Hidden Layer Neural Network Analysis

Sumit Sah, Shruti P Maralappanavar, Bharath B N, Prashant Khanduri

Keywords: FedAvg, Linear Convergence, Generalization, Neural Network

TL;DR: We analyze the optimization and generalization performance of FedAvg in Federated Learning, showing linear convergence under novel conditions with generalization guarantee.

Abstract: In this work, we study the optimization and the generalization performance of the widely used FedAvg algorithm for solving Federated Learning (FL) problems. We analyze the generalization performance of FedAvg by handling the optimization error and the Rademacher complexity. Towards handling optimization error, we propose novel constrained Polyak-Lojasiewicz (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 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. We then bound the Rademacher complexity for this class of neural networks and establish that both Rademacher complexity and the generalization error of FedAvg decrease at an optimal rate of $\mathcal{O}({1}/{\sqrt{n}})$. We further show that increasing the number of clients $K$ decreases the generalization error at the rate of $\mathcal{O}({1}/{\sqrt{n}} + {1}/{\sqrt{nK}})$.

Primary Area: learning theory

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Reciprocal Reviewing: I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 11472