Order Optimal Bounds for One-Shot Federated Learning over non-Convex Loss Functions

TL;DR

This paper establishes fundamental limits and proposes an optimal algorithm for one-shot federated learning with non-convex loss functions, balancing communication constraints and sample size to minimize expected loss.

Contribution

It derives the first order-optimal bounds for one-shot federated learning with non-convex losses and introduces the MRE-NC algorithm matching these bounds.

Findings

Lower bound on expected loss: (rac{1}{\u221a{n}(mB)^{1/d}}, rac{1}{{mn}})

Proposed MRE-NC algorithm achieves near-optimal expected loss in large {mn} regime

Results highlight the trade-off between communication budget and sample size in federated learning

Abstract

We consider the problem of federated learning in a one-shot setting in which there are $m$ machines, each observing $n$ sample functions from an unknown distribution on non-convex loss functions. Let $F:[-1,1]^d\to\mathbb{R}$ be the expected loss function with respect to this unknown distribution. The goal is to find an estimate of the minimizer of $F$. Based on its observations, each machine generates a signal of bounded length $B$ and sends it to a server. The server collects signals of all machines and outputs an estimate of the minimizer of $F$. We show that the expected loss of any algorithm is lower bounded by $\max\big(1/(\sqrt{n}(mB)^{1/d}), 1/\sqrt{mn}\big)$, up to a logarithmic factor. We then prove that this lower bound is order optimal in $m$ and $n$ by presenting a distributed learning algorithm, called Multi-Resolution Estimator for Non-Convex loss function (MRE-NC), whose expected loss matches the lower bound for large $mn$ up to polylogarithmic factors.