Federated Optimization of Smooth Loss Functions

In this work, we study empirical risk minimization (ERM) within a federated learning framework, where a central server seeks to minimize an ERM objective function using <inline-formula> <tex-math notation="LaTeX">$n$ </tex-math></inline-formula> samples of training data that is stored across <inline-formula> <tex-math notation="LaTeX">$m$ </tex-math></inline-formula> clients and the server. The recent flurry of research in this area has identified the Federated Averaging (<inline-formula> <tex-math notation="LaTeX">$\mathtt{FedAve} $ </tex-math></inline-formula>) algorithm as the staple for determining <inline-formula> <tex-math notation="LaTeX">$\epsilon $ </tex-math></inline-formula>-approximate solutions to the ERM problem. Similar to standard optimization algorithms, e.g., stochastic gradient descent, the convergence analysis of <inline-formula> <tex-math notation="LaTeX">$\mathtt{FedAve} $ </tex-math></inline-formula> and its variants only relies on smoothness of the loss function in the optimization parameter. However, loss functions are often very smooth in the training data too. To exploit this additional smoothness in data in a federated learning context, we propose the Federated Low Rank Gradient Descent (FedLRGD) algorithm. Since smoothness in data induces an approximate low rank structure on the gradient of the loss function, our algorithm first performs a few rounds of communication between the server and clients to learn weights that the server can use to approximate clients’ gradients using its own gradients. Then, our algorithm solves the ERM problem at the server using an inexact gradient descent method. To theoretically demonstrate that FedLRGD can have superior performance to <inline-formula> <tex-math notation="LaTeX">$\mathtt{FedAve} $ </tex-math></inline-formula>, we present a notion of federated oracle complexity as a counterpart to canonical oracle complexity in the optimization literature. Under some assumptions on the loss function, e.g., strong convexity and smoothness in the parameter, <inline-formula> <tex-math notation="LaTeX">$\eta $ </tex-math></inline-formula>-Hölder class smoothness in the data, etc., we prove that the federated oracle complexity of <inline-formula> <tex-math notation="LaTeX">$LRGD $ </tex-math></inline-formula> scales like <inline-formula> <tex-math notation="LaTeX">$\phi m (p/\epsilon)^{\Theta (d/\eta)}$ </tex-math></inline-formula> and that of <inline-formula> <tex-math notation="LaTeX">$\mathtt{FedAve} $ </tex-math></inline-formula> scales like <inline-formula> <tex-math notation="LaTeX">$\phi m (p / \epsilon)^{3/4}$ </tex-math></inline-formula> (neglecting typically sub-dominant factors), where <inline-formula> <tex-math notation="LaTeX">$\phi \gg 1$ </tex-math></inline-formula> is the ratio of client-to-server communication time to gradient computation time, <inline-formula> <tex-math notation="LaTeX">$p$ </tex-math></inline-formula> is the parameter dimension, and <inline-formula> <tex-math notation="LaTeX">$d$ </tex-math></inline-formula> is the data dimension. Then, we show that when <inline-formula> <tex-math notation="LaTeX">$d$ </tex-math></inline-formula> is small compared to <inline-formula> <tex-math notation="LaTeX">$n$ </tex-math></inline-formula> and the loss function is sufficiently smooth in the data, i.e., <inline-formula> <tex-math notation="LaTeX">$\eta = \Theta (d)$ </tex-math></inline-formula>, FedLRGD beats <inline-formula> <tex-math notation="LaTeX">$\mathtt{FedAve} $ </tex-math></inline-formula> in federated oracle complexity. Finally, in the course of analyzing FedLRGD, we also establish a general result on low rank approximation of smooth latent variable models.