Federated Optimization of Smooth Loss Functions

The recent escalation of mobile phone usage and emergence of the Internet of Things (IoT) has engendered unprecedented volumes of collected data on such devices. While it is clearly of utility to train machine models based on this data, various new challenges pertaining to privacy, scalability, distributed computation, etc. must be addressed before doing so. For instance, data collected on a user’s mobile phone must be kept private, and even if privacy is not a concern, it is infeasible to communicate large quantities of data from millions of users to a central server for computation. To deal with such issues, a new paradigm known as federated learning has been formulated in the literature [28]. There are several classes of federated learning architectures (see, e.g., [29]): those with a synchronous central server, fully decentralized networks with synchronized clocks, and asynchronous versions of these topologies. We consider synchronous architectures with a central server in this paper.

In this synchronous setting, the architecture consists of a single central server and a collection of several clients (that all have synchronized clocks), cf. [28, 30, 31], and the references therein. Typically, the server and clients hold their own private data. So, to train machine learning models, federated optimization is carried out by clients training possibly inaccurate models based on their private data, and the server periodically aggregating client models to produce a more accurate model at the server. Several key dimensions are considered when studying such federated learning and optimization systems:

1. 1.

   Privacy: The training data resides with the clients and is not shared with the server (and vice versa) [28]. However, the server must leverage the data stored at the clients via several communication rounds to improve the statistical accuracy of the learnt model.

2. 2.

   Large client population with possibly inactive clients: Federated learning systems usually have a very large number of clients so that the client population size often dominates the number of training samples per client [28]. On the other hand, a nontrivial fraction of clients could be inactive during the learning process, e.g., mobile phones only send updates when they are online and being charged [30]. So, federated optimization algorithms must be robust to participation from smaller subsets of clients.

3. 3.

   Communication cost: There is significant cost associated with clients communicating with the server. Typically, clients have limited upload bandwidth and servers have limited download bandwidth. In particular, if a large number of clients seek to simultaneously communicate with the server, the server will put the clients in a queue to receive their updates. So, federated optimization schemes often try to increase the amount of (cheap) computation at each client in order to reduce the frequency of communication with the server (e.g., by employing several local updates at the clients), and utilize smaller fractions of active clients in certain rounds [28, 30]. (Moreover, such schemes may also utilize quantization to compress updates from clients to servers [30]; we do not consider data compression in this work.)

4. 4.

   Latency: Beyond the communication costs delineated above, the duration of client-to-server updates is determined by the slowest client among the set of active clients. Hence, using smaller fractions of active clients in certain rounds is also advantageous to eliminate stragglers. In fact, recent algorithms and analysis suggest that stragglers should be used at later stages of the learning process in order to reduce latency [31].

5. 5.

   Data heterogeneity: Different clients may have data that is drawn from different population distributions [28, 29]. Although some works in the literature, e.g., [30, 31], assume for simplicity of analysis that clients are homogeneous, i.e., the training data at every client is drawn from a common (unknown) probability distribution, we do not need to impose such a stringent assumption in this work. In particular, our new algorithm in Algorithm 1 and its theoretical analysis in Theorem 2 and Section V will hold in the general heterogeneous setting. Usually, data homogeneity simplifies the analysis of stochastic federated optimization methods since clients’ gradients become unbiased estimators of true full gradients [30]. However, since we analyze the (sample-version of the) ERM problem given below in (3), which has the same form in both homogeneous and heterogeneous settings, and we will not use randomized estimators of true full gradients at the clients, our analysis is agnostic to whether the server and client data is homogeneous or heterogeneous.

Formally, we model synchronous federated learning with a single central server as follows. We assume that there are $m\in{\mathbb{N}}$ clients, indexed by $[m]$. The server is endowed with $r\in{\mathbb{Z}}_{+}$ training data samples $x^{(0,1)},\dots,x^{(0,r)}\in[0,1]^{d}$, and each client $i\in[m]$ is given $s\in{\mathbb{N}}$ training data samples $x^{(i,1)},\dots,x^{(i,s)}\in[0,1]^{d}$, where $n=ms+r$ so that $\{x^{(i,j)}:i\in\{0\}\cup[m],\,j\in[s]\}=\{x^{(k)}:k\in[n]\}$.¹¹1It is assumed for simplicity that every client has the same number of training data samples. Under the assumptions outlined in Section II-B, the goal of the server is to utilize the clients to compute an $\epsilon$-approximate solution (in the sense of (11) or (12)) to the ERM problem (1), which may be rewritten as

To achieve this goal, we consider a specific family of federated optimization algorithms. In particular, any such algorithm runs over $T\in{\mathbb{N}}$ synchronized epochs (or rounds). During each epoch, the server may first communicate with $\tau m\in{\mathbb{N}}$ clients, where $\tau\in[0,1]$ denotes the proportion of active clients that participate in computation. Then, the server and the active clients perform gradient computations, or other less expensive computations, towards finding an $\epsilon$-approximate solution to (3), where we typically assume that all active clients perform the same number of gradient computations. Finally, each active client may communicate only one vector in ${\mathbb{R}}^{p}$ to the server (while taking care not to reveal their private data by transmitting vectors from which it could be easily extracted). This client-to-server communication has some nontrivial cost associated with it. On the other hand, it is usually assumed that server-to-client communication has negligible cost, because there are often no critical constraints on client download bandwidths and the server broadcasting bandwidth. So, the server broadcasts vectors in ${\mathbb{R}}^{p}$ to active clients during epochs at essentially no cost. We also note that every epoch need not contain the three aforementioned phases, i.e., the server need not always communicate with active clients or perform gradient computations, and clients need not always perform gradient computations and communicate with the server.

Several federated optimization algorithms for the synchronous central server setting have been proposed in the literature. The benchmark method in the field is known as Federated Averaging (${\mathsf{FedAve}}$) [28], where clients perform several iterations of (mini-batch) SGD using their private data, and the server aggregates client updates by periodically averaging them over several epochs. Needless to say, ${\mathsf{FedAve}}$ belongs to the aforementioned family of algorithms, and has been extended in several directions. For example, the authors of [32] introduce an accelerated version of the algorithm by introducing momentum, and the authors of [30] include a compression (or quantization) step before any client-to-server communication in the algorithm to reduce communication cost. In particular, the iteration complexity and convergence of ${\mathsf{FedAve}}$ are carefully analyzed in [30]. More generally, a unified analysis of the class of communication-efficient SGD algorithms is presented in [33]. Various other federated optimization methods have also been proposed that address different drawbacks of ${\mathsf{FedAve}}$. For instance, federated optimization methods with better fixed point properties are developed in [34], latency of federated optimization algorithms is improved via adaptive node participation in [31], a federated optimization algorithm for training with only positive labels is presented in [35], and adaptive algorithms that incorporate more sophisticated gradient-based iterative methods, such as Adam [19], into the federated learning framework are proposed in [36]. In a different vein, one-shot federated learning is considered in [37], where only one round of client-to-server communication is used in the optimization process (also see follow-up work). We refer readers to [29] and the references therein for an in-depth further discussion of federated learning, distributed optimization, and related ideas.

Finally, we note that in the traditional federated learning literature, only the clients possess training data which the server cannot access. However, as mentioned above, our work assumes that the server also possesses some training data samples. As explained in [38], this is a reasonable assumption in many federated learning settings because the server could also function as a client and have access to its own private training data, it may have to collect some training data from clients (despite privacy concerns) for system monitoring purposes, it may collect data from test devices (or test clients) which do not participate in training, and it could even possess simulated training data from models. For these reasons, there is a growing recent literature on federated learning with server data, cf. [39, 40, 38] and the references therein. Furthermore, as remarked earlier, the server and client datasets can be heterogeneous in our work since the smoothness (or approximate low rankness) of loss functions will circumvent the need for data homogeneity in our analysis.

As mentioned earlier, the vast literature on federated optimization, including the aforementioned works, deals with the ERM problem representation in (2). Thus, these works do not exploit smoothness of loss functions in the training data for federated learning. In this paper, we fill this void by making the following key theoretical contributions:

1. 1.

   We define a notion of federated oracle complexity in Definition 1 to capture the running time of federated optimization algorithms, and argue its suitability using some simple probabilistic analysis. This formally generalizes classical oracle complexity of iterative optimization methods (see [2, Section 1.1]) to a federated learning setting. Our definition allows us to theoretically compare our proposed federated optimization algorithm (mentioned next) with the benchmark ${\mathsf{FedAve}}$ method from the literature. We believe that this definition will be of further utility in future research in the area.

2. 2.

   We propose a new Federated Low Rank Gradient Descent (${\mathsf{FedLRGD}}$) method in Algorithm 1 to solve the ERM problem (1), or equivalently, (3), in the federated learning setting. Our method exploits the approximate low rankness of the gradient of the loss function, which is a consequence of its smoothness in data, by using a few rounds of communication to first learn weights that allow the server to approximate clients’ gradients by computing weighted sums of its own gradients. Then, the server can solve (1) by running an inexact GD method on its own, where the server estimates full gradients at each iteration from gradients corresponding to its private data. Moreover, like other federated optimization algorithms in the literature, our method preserves the privacy of the server and clients by never exchanging data samples, and it works in settings where server and client data is heterogeneous.

3. 3.

   Under the smoothness, strong convexity, and non-singularity assumptions outlined in Section II-B and Section III-C, we analyze the ${\mathsf{FedLRGD}}$ algorithm and show in Theorem 2 that its federated oracle complexity scales like $\phi m(p/\epsilon)^{\Theta(d/\eta)}$ (neglecting typically sub-dominant factors), where $\phi\gg 1$ is the ratio of communication cost to computation cost (see Section II-D), $\epsilon>0$ is the desired approximation accuracy of the solution (see Section II-B), and $\eta>0$ denotes a Hölder class parameter which determines the level of smoothness of the loss function in the data (see Section II-B).

4. 4.

   We also calculate the federated oracle complexity of the ${\mathsf{FedAve}}$ algorithm in Theorem 3 by building on known results from the literature, cf. [30], and show that it scales like $\phi m(p/\epsilon)^{3/4}$. Using this calculation, we demonstrate in Proposition 1 that when the data dimension is small, e.g., $d=O(\log(n)^{0.99})$, and the gradient of the loss function is sufficiently smooth in the training data so that it is approximately low rank, i.e., $\eta=\Theta(d)$, the federated oracle complexity of ${\mathsf{FedLRGD}}$ is much smaller than that of the benchmark ${\mathsf{FedAve}}$ algorithm.

5. 5.

   In the process of analyzing the ${\mathsf{FedLRGD}}$ algorithm, we also derive a new statistical result in Theorem 1 that utilizes the smoothness of loss functions in the training data to bound the Frobenius norm error in estimating a latent variable model (or graphon) corresponding to the loss function with its low rank approximation. This result generalizes the graphon approximation results in [41, 42, 43], and could be of independent utility for future research in high-dimensional statistics.

We briefly delineate the remainder of our discussion here. In Section II, we precisely state and explain the notation and assumptions used in our subsequent analysis, and also introduce and expound upon the concept of federated oracle complexity. In Section III, we present and discuss our main technical results and the ${\mathsf{FedLRGD}}$ algorithm mentioned above. Then, we provide proofs of our results on latent variable model approximation and the federated oracle complexity of ${\mathsf{FedLRGD}}$ in Section IV and Section V, respectively. The proofs of all other results are deferred to the appendices. Although the focus of this work is to theoretically analyze ${\mathsf{FedLRGD}}$ and demonstrate the resulting gain in federated oracle complexity, we also present some pertinent simulation results in Appendix C. These simulations illustrate that gradients of loss functions corresponding to neural networks trained on the MNIST database for handwritten digit recognition (see [44]) and CIFAR-10 dataset of tiny images (see [45]) exhibit approximate low rankness. This suggests that algorithms which exploit low rank structure, such as ${\mathsf{FedLRGD}}$, could have potential impact in applications. Finally, we conclude our discussion in Section VI.

In this paper, we let ${\mathbb{N}}\triangleq\{1,2,3,\dots\}$, ${\mathbb{Z}}_{+}={\mathbb{N}}\cup\!\{0\}$, and ${\mathbb{R}}_{+}$ denote the sets of natural numbers, non-negative integers, and non-negative real numbers, respectively, and for $m\in{\mathbb{N}}$, we let $[m]\triangleq\{1,\dots,m\}$. We let ${\mathbb{P}}(\cdot)$ and ${\mathbb{E}}[\cdot]$ denote the probability and expectation operators, respectively, where the underlying probability laws will be clear from context. Furthermore, $\exp(\cdot)$ and $\log(\cdot)$ denote the natural exponential and logarithm functions with base $e$, respectively, ${\mathds{1}}\{\cdot\}$ denotes the indicator function that equals $1$ if its input proposition is true and equals $0$ otherwise, and $\lceil\cdot\rceil$ denotes the ceiling function. Lastly, we will utilize Bachmann-Landau asymptotic notation throughout this paper, e.g., $O(\cdot)$, $\Omega(\cdot)$, $\Theta(\cdot)$, etc., where the limiting variable will be clear from context.

In addition, for any matrix $A\in{\mathbb{R}}^{m\times k}$ with $m,k\in{\mathbb{N}}$, $A_{i,j}$ denotes the $(i,j)$th entry of $A$ for $i\in[m]$ and $j\in[k]$, $A^{\mathrm{T}}$ denotes the transpose of $A$, $A^{-1}$ denotes the inverse of $A$ (when $m=k$ and $A$ is non-singular), $\operatorname{{rank}}(A)$ denotes the rank of $A$, $\|A\|_{\mathrm{F}}$ denotes the Frobenius norm of $A$, and $\|A\|_{\mathrm{op}}$ denotes the spectral or (induced $\ell^{2}$) operator norm of $A$. Similarly, for any (column) vector $y\in{\mathbb{R}}^{m}$, $y_{i}$ denotes the $i$th entry of $y$ for $i\in[m]$, and $\|y\|_{q}$ denotes the $\ell^{q}$-norm of $y$ for $q\in[1,+\infty]$. Moreover, for any continuously differentiable function $h:{\mathbb{R}}^{m}\rightarrow{\mathbb{R}}$, $h(y)$, we define $\partial h/\partial y_{i}:{\mathbb{R}}^{m}\rightarrow{\mathbb{R}}$ to be its partial derivative with respect to $y_{i}$ for $i\in[m]$, and $\nabla_{y}h:{\mathbb{R}}^{m}\rightarrow{\mathbb{R}}^{m}$, $\nabla_{y}h=\left[\partial h/\partial y_{1}\,\cdots\,\partial h/\partial y_{m}\right]^{\mathrm{T}}$ to be its gradient. As in [23], we will also require some multi-index notation. So, for any $m$-tuple $s=(s_{1},\dots,s_{m})\in{\mathbb{Z}}_{+}^{m}$, we let $|s|=\|s\|_{1}=s_{1}+\cdots+s_{m}$, $s!=s_{1}!\cdots s_{m}!$, $y^{s}=y_{1}^{s_{1}}\cdots y_{m}^{s_{m}}$ for $y\in{\mathbb{R}}^{m}$, and $\nabla^{s}h=\nabla_{y}^{s}h=\partial^{|s|}h/(\partial y_{1}^{s_{1}}\cdots\partial y_{m}^{s_{m}})$ be the $s$th partial derivative of $h$ with respect to $y$.

As mentioned earlier, for any $n,d,p\in{\mathbb{N}}$, assume that we are given $n$ samples of training data $x^{(1)},\dots,x^{(n)}\in[0,1]^{d}$ and a loss function $f:[0,1]^{d}\times{\mathbb{R}}^{p}\rightarrow{\mathbb{R}}$, $f(x;\theta)$, where $\theta\in{\mathbb{R}}^{p}$ is the parameter vector. In our ensuing analysis, we will impose the following assumptions on the loss function (cf. [2, 4, 23]):

1. 1.

   Smoothness in parameter $\theta$: There exists $L_{1}>0$ such that for all $x\in[0,1]^{d}$, $\nabla_{\theta}f(x;\cdot):{\mathbb{R}}^{p}\rightarrow{\mathbb{R}}^{p}$ exists and is $L_{1}$-Lipschitz continuous:

   	$\displaystyle\forall\theta^{(1)},\theta^{(2)}\in{\mathbb{R}}^{p},\,$	$\displaystyle\left\|\nabla_{\theta}f(x;\theta^{(1)})-\nabla_{\theta}f(x;\theta^{(2)})\right\|_{2}$
   		$\displaystyle\qquad\qquad\enspace\leq L_{1}\left\|\theta^{(1)}-\theta^{(2)}\right\|_{2}.$		(4)

2. 2.

   Strong convexity: There exists $\mu>0$ such that for all $x\in[0,1]^{d}$, $f(x;\cdot):{\mathbb{R}}^{p}\rightarrow{\mathbb{R}}$ is $\mu$-strongly convex:

   	$\displaystyle\forall\theta^{(1)},\theta^{(2)}$	$\displaystyle\in{\mathbb{R}}^{p},$
   	$\displaystyle f(x;\theta^{(1)})$	$\displaystyle\geq f(x;\theta^{(2)})+\nabla_{\theta}f(x;\theta^{(2)})^{\mathrm{T}}(\theta^{(1)}-\theta^{(2)})$
   		$\displaystyle\quad\,+\frac{\mu}{2}\left\|\theta^{(1)}-\theta^{(2)}\right\|_{2}^{2}.$		(5)

   Furthermore, we let $\kappa\triangleq L_{1}/\mu\geq 1$ be the condition number of the loss function.

3. 3.

   Smoothness in data $x$: There exist $\eta>0$ and $L_{2}>0$ such that for all fixed $\vartheta\in{\mathbb{R}}^{p}$ and for all $i\in[p]$, the $i$th partial derivative of $f$ at $\vartheta$, denoted $g_{i}(\cdot;\vartheta)\triangleq\frac{\partial f}{\partial\theta_{i}}(\cdot;\vartheta):[0,1]^{d}\rightarrow{\mathbb{R}}$, belongs to the $(\eta,L_{2})$-Hölder class (cf. [27, Definition 1.2]), i.e., $g_{i}(\cdot;\vartheta):[0,1]^{d}\rightarrow{\mathbb{R}}$ is $l=\lceil\eta\rceil-1$ times differentiable, and for every $s=(s_{1},\dots,s_{d})\in{\mathbb{Z}}_{+}^{d}$ such that $|s|=l$, we have

   	$\displaystyle\forall y^{(1)},y^{(2)}\in[0,1]^{d},$	$\displaystyle\enspace\left|\nabla_{x}^{s}g_{i}(y^{(1)};\vartheta)-\nabla_{x}^{s}g_{i}(y^{(2)};\vartheta)\right|$
   		$\displaystyle\quad\quad\leq L_{2}\left\|y^{(1)}-y^{(2)}\right\|_{1}^{\eta-l}.$		(6)

4. 4.

   Non-singularity: For any (large) size $r\in{\mathbb{N}}$, any index $i\in[p]$, any sequence of distinct data samples $y^{(1)},\dots,y^{(r)}\in[0,1]^{d}$, and any sequence of distinct parameter values $\vartheta^{(1)},\dots,\vartheta^{(r)}\in{\mathbb{R}}^{p}$,²²2More formally, this assumption only holds for $\vartheta^{(1)},\dots,\vartheta^{(r)}$ almost everywhere, but we neglect this detail here and in the sequel for simplicity. the matrix of partial derivatives $G^{(i)}\in{\mathbb{R}}^{r\times r}$ given by:

   $$\forall j,k\in[r],\enspace G^{(i)}_{j,k}=g_{i}(y^{(j)};\vartheta^{(k)})\,,$$

   (7)

   is invertible. Note that we suppress the dependence of $G^{(i)}$ on $r$ and the sequences of data and parameters for convenience.

5. 5.

   Conditioning: There exists an absolute constant $\alpha\geq 1$ such that for any (large) $r\in{\mathbb{N}}$, any $i\in[p]$, any sequence of distinct data samples $y^{(1)},\dots,y^{(r)}\in[0,1]^{d}$, and any sequence of distinct parameter values $\vartheta^{(1)},\dots,\vartheta^{(r)}\in{\mathbb{R}}^{p}$,³³3More formally, one would need to impose additional distributional assumptions to prove that this assumption holds for random $\vartheta^{(1)},\dots,\vartheta^{(r)}$ with high probability, but we also neglect this detail here and in the sequel for simplicity. the matrix $G^{(i)}\in{\mathbb{R}}^{r\times r}$ given in (7) satisfies

   $$\left\|(G^{(i)})^{-1}\right\|_{\mathrm{op}}\leq r^{\alpha}\,,$$

   (8)

   i.e., the minimum singular value of $G^{(i)}$ decreases to $0$ at most polynomially fast in the dimension $r$.

6. 6.

   Boundedness: There exists an absolute constant $B>0$ such that for all $i\in[p]$, we have

   $$\sup_{x\in[0,1]^{d}}\sup_{\vartheta\in{\mathbb{R}}^{p}}\left|g_{i}(x;\vartheta)\right|\leq B\,.$$

   (9)

Next, let $\epsilon>0$ be the desired approximation accuracy of the solution, $\theta^{\star}\in{\mathbb{R}}^{p}$ be the unique global minimizer of the objective function $F$ in (1), and

(Note that $\theta^{\star}$ exists and $F_{*}$ is finite due to the strong convexity assumption.) In the sequel, we consider federated optimization algorithms that output an $\epsilon$-approximate solution $\theta^{*}\in{\mathbb{R}}^{p}$ to (1) such that:

1. 1.

   If the algorithm generates a deterministic solution $\theta^{*}$, we have

   $$F(\theta^{*})-F_{*}\leq\epsilon\,.$$

   (11)

2. 2.

   If the (stochastic) algorithm produces a random solution $\theta^{*}$, we have

   $${\mathbb{E}}\!\left[F(\theta^{*})\right]-F_{*}\leq\epsilon\,,$$

   (12)

   where the expectation is computed with respect to the law of $\theta^{*}$.

We briefly explain some of the important assumptions in Section II-B and contextualize them within the optimization and statistics literatures.

Firstly, smoothness in parameter is a standard assumption in the optimization for machine learning literature (see, e.g.,[2, 4], and the references therein). This assumption implies that the empirical risk $F(\theta)$ has $L_{1}$-Lipschitz continuous gradient (which is used in the proof of Theorem 2 in Section V, cf. [23, Section 4.1]). In particular, when the gradient of the objective function $F(\theta)$ is $L_{1}$-Lipschitz continuous, it can be shown that each iteration of GD reduces the objective value:

for all $\theta\in{\mathbb{R}}^{p}$ [4, Equation (3.5)]. The inequality (13) illustrates the greedy nature in which GD decreases its objective value and is fundamental in the analysis of iterative descent methods. For this reason, this assumption has been used extensively in the optimization literature to prove convergence guarantees in a variety of settings, e.g., [9, 8, 5, 12, 13, 16, 6, 14, 21], etc. We also impose this assumption since the analysis to obtain Lemma 2 in Section V, which is a restatement of [5, Lemma 2.1], requires it.

Similarly, strong convexity (in parameter) is another standard assumption in the optimization for machine learning literature (see, e.g.,[2, 4], and the references therein). As before, this assumption implies that the empirical risk $F(\theta)$ is strongly convex (which is used in the proof of Theorem 2 in Section V, cf. [2, Lemma 2.1.4], [23, Section 4.1]). As explained in [4], the main utility of imposing strong convexity of the objective function $F(\theta)$ is a notable and provable speed-up in the convergence rate of gradient-based iterative methods. Specifically, GD can achieve linear convergence rate (i.e., exponentially fast convergence) when the objective function is strongly convex [2, 4]. In addition, the strong convexity assumption often leads to far simpler proofs of convergence and oracle complexity, and easier illustration of the benefits of improvements like momentum (or heavy-ball method) [10] and acceleration [11] on oracle complexity (cf. [2, 21]). For these reasons, this assumption has also been used extensively in the optimization literature to prove convergence guarantees in a variety of settings, e.g., [9, 8, 5, 13, 16, 6, 14, 21], etc. We also assume strong convexity for these reasons; indeed, strong convexity is needed for Lemma 2 in Section V, which is a restatement of [5, Lemma 2.1], to hold and the first term in the bound in Lemma 2 illustrates the aforementioned exponential convergence.

Smoothness in data is a standard assumption in the non-parametric estimation literature (see, e.g.,[27, 46, 47], and the references therein). In simple terms, we can think of a function in a Hölder class with positive integer parameter $\eta$ as an $(\eta-1)$-times differentiable function with Lipschitz continuous $(\eta-1)$th partial derivatives (or even an $\eta$-times differentiable function with bounded $\eta$th partial derivatives). Hence, the size of the parameter $\eta$ controls the degree of smoothness of functions in a Hölder class. Such Hölder class assumptions are used in various non-parametric regression techniques, e.g., kernel regression such as Nadaraya-Watson estimators [48, 49, 27], local polynomial regression [25, 26, 27], and nearest neighbor methods [50, 51, 46]. At a high-level, the basic idea of imposing smoothness assumptions for regression is the following: If an unknown function living in a “large” (or non-parametric) class is known to be smoother, then it can be better estimated from samples using Taylor polynomials. When defining the Hölder condition in (6), one can in principle use any $\ell^{q}$-norm (see, e.g., [42, 47, 43, 23]). We use the $\ell^{1}$-norm to define (6) because this yields the largest class of functions for fixed parameters $(\eta,L_{2})$; indeed, if (6) is satisfied for any $\ell^{q}$-norm, then it is also satisfied for the $\ell^{1}$-norm using the monotonicity of $\ell^{q}$-norms. In this sense, our definition includes the largest class of functions. In Appendix A, we provide an example of a commonly used loss function in machine learning that satisfies the aforementioned smoothness and strong convexity assumptions.

In a recent sequence of works in high-dimensional statistics, it has been shown that matrices defined by bivariate graphon functions or latent variable models are approximately low rank when the function satisfies Hölder class assumptions on one of the variables [41, 42, 43]. (We generalize these results in Theorem 1.) Partial derivatives of loss functions $g_{i}$ can be construed as bivariate functions of the data $x$ and parameter $\theta$, and we impose Hölder class smoothness in data on them in order to exploit the resulting approximate low rank structure for federated optimization (as we will elaborate in Section III). As noted earlier and in [23], smoothness in data exists in a variety of ERM problems, e.g., (regularized) linear regression [3, Chapter 3], (regularized) logistic regression [3, Section 4.4], deep neural network training with smooth loss and activation functions [24, Chapter 5], etc. The close connection between smoothness of bivariate functions and approximate low rank structure may seem peculiar at first glance. Perceiving a bivariate function as a kernel of an integral operator, the results in [41, 42, 43] suggest that smoothness of the kernel is closely related to approximate low rankness of the associated integral operator. This relation has its roots in more familiar ideas from harmonic analysis. Indeed, in the setting of difference (or convolution) kernels that are defined by a univariate function, increasing the smoothness of the function is essentially equivalent to faster decay of its Fourier transform [52, Theorem 60.2] (also see [53, 54]). This, in turn, can be interpreted as approximate low rankness of the associated convolution operator, because the Fourier transform is precisely the spectrum of the convolution operator.

Approximate low rank structure, broadly construed, has been utilized in a variety of areas, e.g., matrix estimation [55], reinforcement learning [56], function approximation via ridge functions [57, 58], semidefinite programming [59], uncertainty quantification [60], etc. In this work, we seek to exploit approximate low rank structure in the context of federated optimization. Just as smoothness is prevalent in many ERM problems, a litany of empirical investigations in the literature have directly demonstrated various forms of approximate low rank structure in large-scale machine learning tasks. For instance, the empirical risk and loss functions corresponding to deep neural network training problems display approximate low rank structure in their gradients (which live in low-dimensional subspaces) or Hessians (which have decaying eigenvalues, thereby inducing low rank gradients via Taylor approximation arguments) [61, 62, 63, 64, 65, 66]. In particular, [61] shows that when training deep neural networks for $K$-ary classification, the computed gradients belong to $K$-dimensional eigenspaces of the Hessian corresponding to its dominant eigenvalues. On a related front, low rankness has also been observed in weight matrices of deep neural networks [67, 68], neural collapse has been demonstrated when training deep neural networks [69, 70, 71], and training memory for neural networks has been saved using low rank parameterizations of gradient weight matrices [72]. Our work on exploiting approximate low rank structure for federated optimization, albeit different, is inspired by many of these observations.

The non-singularity, conditioning, and boundedness assumptions in Section II-B and Section III-C have been imposed for technical purposes for the proofs. Variants of such assumptions pertaining to matrix invertibility and spectra are often made for analytical tractability in theoretical studies in high-dimensional statistics (see, e.g., the invertibility and spectral boundedness assumptions made in [73, Section 2] to analyze linear regression). While it is difficult to provide simple and easily interpretable sufficient conditions that ensure the non-singularity assumptions in Section II-B and Section III-C, we provide some intuition for when we expect such assumptions to hold. To this end, consider the non-singularity assumption in Section II-B. As before, let us construe $g_{i}(x,\theta)$ as the kernel of an integral operator. Then, the spectral theory (specifically, singular value decomposition) of compact operators states that the kernel can be written as [56, Theorem 9] (also see [74]):

where $\{\sigma_{j}\geq 0:j\in{\mathbb{N}}\}$ are singular values that converge to $0$ as $j\rightarrow\infty$, $\{u_{j}:j\in{\mathbb{N}}\}$ and $\{v_{j}:j\in{\mathbb{N}}\}$ are orthonormal bases of singular vector functions, and although the equality above usually holds in an ${\mathcal{L}}^{2}$-sense, it can also hold in a pointwise sense under some assumptions. Hence, for data samples $y^{(1)},\dots,y^{(r)}$ and parameter values $\vartheta^{(1)},\dots,\vartheta^{(r)}$, we may write the matrix $G^{(i)}$ as

using (7) and (14), where we let $\mathsf{u}_{j}=\big{[}u_{j}(y^{(1)})\,\cdots\,u_{j}(y^{(r)})\big{]}^{\mathrm{T}}$ and $\mathsf{v}_{j}=\big{[}v_{j}(\vartheta^{(1)})\,\cdots\,v_{j}(\vartheta^{(r)})\big{]}^{\mathrm{T}}$. Since the integral operator corresponding to $g_{i}$ is a map between infinite-dimensional spaces, we typically have infinitely many non-zero $\sigma_{j}$’s. Moreover, since $\{u_{j}:j\in{\mathbb{N}}\}$ and $\{v_{j}:j\in{\mathbb{N}}\}$ are linearly independent sets of functions, we typically have $r$ non-zero unit-rank terms $\sigma_{j}\mathsf{u}_{j}\mathsf{v}_{j}^{\mathrm{T}}$ in (15), with linearly independent $\mathsf{u}_{j}$’s and $\mathsf{v}_{j}$’s, whose linear independence is not cancelled out by the other terms. Thus, $G^{(i)}$ has rank $r$ and is invertible. This provides some intuition behind the non-singularity assumption in Section II-B. The non-singularity of monomials and approximation assumptions in Section III-C have similar underlying intuition based on sampling of linearly independent functions (cf. (45) and (44), (46), respectively). In closing this subsection, we remark that one well-known setting where all sampled matrices of a symmetric bivariate function are invertible, and in fact, positive definite, is when the bivariate function itself defines a valid positive definite kernel (in the reproducing kernel Hilbert space sense), cf. [75, Section 4-2-4.6]. In the context of Section II-B, we cannot directly assume that every $g_{i}$ is a positive definite kernel since $g_{i}$ is asymmetric in general and we only require invertibility of sampled matrices, not positive definiteness, which is a much stronger condition. However, alternative versions of such kernel conditions could be developed in future work to provide cleaner assumptions that imply various non-singularity conditions.

Finally, we note that, strictly speaking, the boundedness and strong convexity assumptions make sense simultaneously when the parameter space is compact. Although our analysis is carried out with the parameter space ${\mathbb{R}}^{p}$, since iterates of inexact GD in Algorithm 1 only live in some compact space, our analysis carries over to the compact setting. For example, one way to rigorously carry over our analysis when it is assumed that the parameter space is compact is to run projected inexact GD instead of inexact GD in Algorithm 1. Then, the parameter iterates of Algorithm 1 remain in the compact space, and we can execute the analysis of our algorithm using the convergence analysis in, e.g., [8, Proposition 3]. For simplicity of exposition, however, we do not use projected inexact GD and simply take the parameter space as ${\mathbb{R}}^{p}$ in this paper.

In the traditional (non-distributed) optimization literature, the running time of first-order iterative methods, such as GD and SGD, is abstractly measured by the notion of (first-order) oracle complexity, i.e., the total number of queries to an oracle that returns the gradient of the loss function with respect to the parameter [76] (also see [2, Section 1.1]). This is because gradient computations are usually the most expensive operations in such algorithms. However, the problem of measuring running time is quite different for federated optimization algorithms, because the cost of clients communicating with servers is typically much higher than the running time of gradient computations. Currently, simulations based on several informal principles (such as the aforementioned one) are often used to empirically compare the efficiency of different federated optimization algorithms, rather than theoretically analyzing any formal notion of complexity; see, e.g., [30] and the references therein. To partially remedy this, inspired by [30], we propose a new definition of federated oracle complexity by formalizing some of the informal rules delineated in Section I-A using simple probabilistic approximations. We note that other notions of oracle complexity for distributed optimization have either implicitly or explicitly been posed in the literature, cf. graph oracle models [77] or [78] and the references therein. In fact, our definition of federated oracle complexity is reminiscent of the complexity measure introduced in [79], but our definition is adapted to the centralized, synchronous federated learning setting rather than distributed optimization frameworks.

Consider any algorithm belonging to the family outlined in Section I-A. In the sequel, we will approximate the running times of the two time intensive phases of each epoch: the gradient computation phase and the client-to-server communication phase.

To establish the federated oracle complexity of our algorithm, we will require a basic result about uniformly approximating a smooth function using piecewise polynomials. To illustrate this result, consider any function $g:[0,1]^{d}\times{\mathbb{R}}^{p}\rightarrow{\mathbb{R}}$, $g(x;\theta)$ such that for all $\vartheta\in{\mathbb{R}}^{p}$, the function $g(\cdot;\vartheta):[0,1]^{d}\rightarrow{\mathbb{R}}$ belongs to the $(\eta,L_{2})$-Hölder class as defined in (6) in Section II-B. Moreover, for any $q>0$, consider a minimal $\ell^{1}$-covering of the hypercube $[0,1]^{d}$ using closed $\ell^{1}$-balls (or cross-polytopes) with radius $q^{-1}$ and centers ${\mathcal{N}}=\{z^{(1)},\dots,z^{(|{\mathcal{N}}|)}\}\subseteq[0,1]^{d}$, which forms the associated $q^{-1}$-net, such that:

In particular, (27) shows that the union of the closed $\ell^{1}$-balls with radius $q^{-1}$ and centers ${\mathcal{N}}$ contains the set $[0,1]^{d}$, and the $q^{-1}$-net ${\mathcal{N}}$ is minimal in the sense that its cardinality $|{\mathcal{N}}|$ is as small as possible while ensuring (27), i.e., $|{\mathcal{N}}|$ is the $\ell^{1}$-covering number of $[0,1]^{d}$ (cf. [83, Definition 4.2.2]). Let $B_{1},\dots,B_{|{\mathcal{N}}|}\subseteq{\mathbb{R}}^{d}$ be an enumeration of the closed $\ell^{1}$-balls with radius $q^{-1}$ and centers $z^{(1)},\dots,z^{(|{\mathcal{N}}|)}\in{\mathcal{N}}$, respectively, and construct the sets $I_{1},\dots,I_{|{\mathcal{N}}|}\subseteq[0,1]^{d}$ such that $I_{1}=B_{1}\cap[0,1]^{d}$ and $I_{i}=(B_{i}\backslash(B_{1}\cup\cdots\cup B_{i-1}))\cap[0,1]^{d}$ for $i\in\{2,\dots,|{\mathcal{N}}|\}$. Then,

forms a partition of the hypercube $[0,1]^{d}$ such that every point in $I_{j}$ is at an $\ell^{1}$-distance of at most $q^{-1}$ from $z^{(j)}$. With this $\ell^{1}$-ball based partition in place, the next lemma conveys that $g$ can be uniformly approximated using a piecewise polynomial function akin to the proofs of [84, Proposition 1] and [43, Proposition 3.3].

Lemma 1 is proved in Section IV-A. While this lemma will suffice to establish federated oracle complexity in our problem, we also present a corollary of it in the context of latent variable models for completeness.

For any $n,k\in{\mathbb{N}}$, any $y^{(1)},\dots,y^{(n)}\in[0,1]^{d}$, and any $\vartheta^{(1)},\dots,\vartheta^{(k)}\in{\mathbb{R}}^{p}$, define the matrix $M\in{\mathbb{R}}^{n\times k}$ such that:

We refer to $M$ as a latent variable model with latent parameters $y^{(1)},\dots,y^{(n)}\in[0,1]^{d}$ and $\vartheta^{(1)},\dots,\vartheta^{(k)}\in{\mathbb{R}}^{p}$, and latent function $g$ (which satisfies the Hölder smoothness assumptions outlined earlier); see, e.g., [43] and the references therein. Furthermore, in the special case where $g$ is symmetric in its two inputs and $M$ is a symmetric matrix, $g$ is known as a graphon function (which naturally appears in Aldous-Hoover representations of exchangeable random graphs [85]), cf. [41, 86, 42]. Next, let

be the singular value decomposition of $M$ with singular values $\sigma_{1}\geq\sigma_{2}\geq\cdots\geq\sigma_{\min\{n,k\}}\geq 0$ and orthonormal bases of singular vectors $\{u_{1},\dots,u_{n}\}\subseteq{\mathbb{R}}^{n}$ and $\{v_{1},\dots,v_{k}\}\subseteq{\mathbb{R}}^{k}$. Then, by the Eckart-Young-Mirsky theorem, for any $r\in[\min\{n,k\}]$, the best rank $r$ approximation of $M$ in Frobenius norm is given by

i.e., $M_{r}=\operatorname*{arg\,min}_{A\in{\mathbb{R}}^{n\times k}:\,\operatorname{{rank}}(A)\leq r}{\|M-A\|_{\mathrm{F}}}$, cf. [87, Section 7.4.2]. Akin to [42, Proposition 1] and [43, Proposition 3.3], the ensuing theorem demonstrates that the latent variable model $M$ can be well-estimated by its low rank approximation due to the smoothness of $g$. We believe that this result may be of broader interest to the matrix and tensor estimation community.

Theorem 1 is proved in Section IV-B; we will use ideas from this proof in our analysis of federated oracle complexity. It is worth comparing this result with the related results in [41], [84, Proposition 1], [42, Proposition 1], and [43, Proposition 3.3]. The author of [41] proves an analog of Theorem 1 in the special case where $g$ is a graphon function and $g(\cdot;\vartheta)$ is Lipschitz continuous with respect to the $\ell^{\infty}$-norm for all $\vartheta\in{\mathbb{R}}^{p}$. In this special case, a piecewise constant function approximation of $g$ suffices for the proof. The results in [84, Proposition 1] and [42, Proposition 1] extend the result of [41] to the setting where $g$ is still a graphon function, but $g(\cdot;\vartheta)$ belongs to an $(\eta,L_{2})$-Hölder class with respect to the $\ell^{\infty}$-norm. The key insight in [84, 42] is to utilize piecewise polynomial function approximations instead of piecewise constant functions; we also use this key idea in Lemma 1. Finally, the authors of [43, Proposition 3.3] apply the argument of [84, 42] to extend their result further to the setting where $g$ is a general latent variable model, similar to our result, but $g(\cdot;\vartheta)$ only belongs to an $(\eta,L_{2})$-Hölder class with respect to the $\ell^{\infty}$-norm. In contrast, our result in Theorem 1 applies to general latent variable models $g$ such that $g(\cdot;\vartheta)$ belongs to an $(\eta,L_{2})$-Hölder class with respect to the $\ell^{1}$-norm, which subsumes the settings of [41, 84, 42, 43]. To prove this more general result, we resort to using a modified minimal $\ell^{1}$-covering to construct our piecewise polynomial function approximation in the intermediate result in Lemma 1; this differs from the $\ell^{\infty}$-coverings used in [41], [84, Lemma 3], [42], and [43], and it can be verified that $\ell^{\infty}$-coverings produce a worse bound in Lemma 1.⁷⁷7We note that our modified $\ell^{1}$-covering of $[0,1]^{d}$ is more sophisticated than the $\ell^{\infty}$-coverings used in [41, 84, 42, 43], because $\ell^{\infty}$-balls can be used to form a (geometric) honeycomb in the hypercube $[0,1]^{d}$, but $\ell^{1}$-balls do not form such a honeycomb in $[0,1]^{d}$ for general $d\in{\mathbb{N}}$. Therefore, in contrast to [41, 84, 42, 43], our proofs of Lemma 1 and Theorem 1 are adapted to deal with the $\ell^{1}$-norm Hölder class assumption on $g(\cdot;\vartheta)$ and the modified minimal $\ell^{1}$-covering of the hypercube $[0,1]^{d}$.

We now describe the new Federated Low Rank Gradient Descent (${\mathsf{FedLRGD}}$) algorithm for federated optimization under the setup of Section I-A and Section II-B. Recall that we are given a central server with $r$ private training data samples $x^{(0,1)},\dots,x^{(0,r)}\in[0,1]^{d}$, and $m$ clients with $s$ private training data samples each; client $c\in[m]$ has samples $x^{(c,1)},\dots,x^{(c,s)}\in[0,1]^{d}$. The number of samples at the server, $r$, depends on the approximate rank of the loss function, as we will see in the next subsection. The server’s objective is to compute an $\epsilon$-approximate solution in the sense of (11) to the ERM problem (3).