Federated Optimization of Smooth Loss Functions

TL;DR

The paper tackles ERM in federated settings and introduces FedLRGD, a gradient-descent method that exploits data-smoothness to induce an approximate low-rank structure in client gradients. By learning rank-revealing weights through limited communication and performing inexact server GD, FedLRGD achieves federated oracle complexity scaling as $\phi m (p/\epsilon)^{\Theta(d/\eta)}$, potentially beating the FedAve benchmark, which scales as $\phi m (p/\epsilon)^{3/4}$ when the data dimension $d$ is small and the gradient is highly smooth in the data. A novel latent-variable approximation result shows that smoothness in the data yields effective low-rank representations for latent models, which underpins the FedLRGD analysis. The work provides a formal framework for comparing federated optimization algorithms via federated oracle complexity and demonstrates regimes where exploiting data smoothness yields practical gains for distributed learning tasks.

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 $ε$-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, $η$-Hölder smoothness in data, etc., we prove that the federated oracle complexity of FedLRGD scales like $φm(p/ε)^{Θ(d/η)}$ and that of FedAve scales like $φm(p/ε)^{3/4}$ (neglecting sub-dominant factors), where $φ\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.

Figures (2)

• Figure 1: Histograms of approximate ranks for MNIST neural networks with different choices of hyper-parameters. The different values of $K$ and $k$ are stated under each plot. Unless stated otherwise, the neural network architectures for all the plots have one hidden layer, use logistic activation functions in all hidden layers, and use the Adam algorithm for training. Moreover, unless stated otherwise, the plots only consider approximate ranks of ${\mathcal{M}}_q$'s corresponding to non-trivial weights $\theta^*_q$. Note that Figure \ref{['Fig: 784, logistic, 60, all weights']} displays approximate ranks for ${\mathcal{M}}_q$'s corresponding to all weights $\theta^*_q$, Figure \ref{['Fig: 784, logistic, 60, RMSprop']} uses the RMSprop algorithm for training, Figure \ref{['Fig: 784, logistic, 60, two hidden layers']} has two hidden layers of size $K = 784$ each, and Figure \ref{['Fig: 784, relu, 60']} uses ReLU activation functions in the hidden layer.
• Figure 2: Histograms of approximate ranks for CIFAR-10 neural networks with different choices of hyper-parameters. The different values of $k$ are stated under each plot. Unless stated otherwise, the neural network architectures for all the plots have two hidden convolutional layers, a fourth max-pooling layer, and two final fully connected layers, and they use logistic activation functions in all non-sub-sampling hidden layers and the Adam algorithm for training. Moreover, unless stated otherwise, the plots only consider approximate ranks of ${\mathcal{M}}_q$'s corresponding to non-trivial weights $\theta^*_q$. Note that Figure \ref{['Fig: cifar10, RMSprop']} uses the RMSprop algorithm for training and Figure \ref{['Fig: cifar10, relu']} uses ReLU activation functions in all non-sub-sampling hidden layers.

Theorems & Definitions (17)

• Definition 1: Federated Oracle Complexity
• Lemma 1: Uniform Piecewise Polynomial Approximation
• Theorem 1: Low Rank Approximation
• Theorem 2: Federated Oracle Complexity of ${\mathsf{FedLRGD}}$
• Theorem 3: Federated Oracle Complexity of ${\mathsf{FedAve}}$
• Proposition 1: Comparison between ${\mathsf{FedLRGD}}$ and ${\mathsf{FedAve}}$
• proof
• proof
• Lemma 2: Inexact GD Bound FriedlanderSchmidt2012
• proof : Proof of Theorem \ref{['Thm: Federated Oracle Complexity of FedLRGD']}