Federated Optimization of Smooth Loss Functions

Abstract

In this work, we study empirical risk minimization (ERM) within a federated learning framework, where a central server minimizes an ERM objective function using training data that is stored across $m$ clients. In this setting, the Federated Averaging (FedAve) algorithm is the staple for determining $\epsilon$-approximate solutions to the ERM problem. Similar to standard optimization algorithms, the convergence analysis of FedAve 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, we propose the Federated Low Rank Gradient Descent (FedLRGD) algorithm. Since smoothness in data induces an approximate low rank structure on the loss function, our method first performs a few rounds of communication between the server and clients to learn weights that the server can use to approximate clients' gradients. Then, our method solves the ERM problem at the server using inexact gradient descent. To show that FedLRGD can have superior performance to FedAve, we present a notion of federated oracle complexity as a counterpart to canonical oracle complexity. Under some assumptions on the loss function, e.g., strong convexity in parameter, $\eta$-Hölder smoothness in data, etc., we prove that the federated oracle complexity of FedLRGD scales like $\varphi m(p/\epsilon {)}^{\Theta (d/\eta )}$ and that of FedAve scales like $\varphi m(p/\epsilon {)}^{3/4}$ (neglecting sub-dominant factors), where $\varphi \gg 1$ is a "communication-to-computation ratio," $p$ is the parameter dimension, and $d$ is the data dimension. Then, we show that when $d$ is small and the loss function is sufficiently smooth in the data, FedLRGD beats FedAve in federated oracle complexity. Finally, in the course of analyzing FedLRGD, we also establish a result on low rank approximation of latent variable models.

31 pages, double column format, 2 figures

References in corpus (10)

• A Stochastic Gradient Method with an Exponential Convergence Rate for Finite Training Sets
• Prevalence of Neural Collapse during the terminal phase of deep learning training
• Non-convex Optimization for Machine Learning
• One-Shot Federated Learning
• Gradient Descent Happens in a Tiny Subspace
• Faster On-Device Training Using New Federated Momentum Algorithm
• Sample Efficient Reinforcement Learning via Low-Rank Matrix Estimation
• Limitations of Neural Collapse for Understanding Generalization in Deep Learning
• Gradient-Based Empirical Risk Minimization using Local Polynomial Regression
• Federated Learning with Server Learning: Enhancing Performance for Non-IID Data

Cited by in corpus (1)

• Unified Analysis of Decentralized Gradient Descent: a Contraction Mapping Framework