Agnostic Federated Learning

Motivation

A key learning scenario in large-scale applications is that of federated learning. In that scenario, a centralized model is trained based on data originating from a large number of clients, which may be mobile phones, other mobile devices, or sensors (Konečnỳ, McMahan, Yu, Richtárik, Suresh, and Bacon, 2016b, a). The training data typically remains distributed over the clients, each with possibly unreliable or relatively slow network connections.

Federated learning raises several types of issues and has been the topic of multiple research efforts. These include systems, networking and communication bottleneck problems due to frequent exchanges between the central server and the clients . To deal with such problems, McMahan et al. (2017) suggested an averaging technique that consists of transmitting the central model to a subset of clients, training it with the data locally available, and averaging the local updates. Smith et al. (2017) proposed to further leverage the relationship between clients, assumed to be known, and cast the problem as an instance of multi-task learning to derive local client models benefiting from other similar ones.

The optimization task in federated learning, which is a principal problem in this scenario, has also been the topic of multiple research work. That includes the design of more efficient communication strategies (Konečnỳ, McMahan, Yu, Richtárik, Suresh, and Bacon, 2016b, a; Suresh, Yu, Kumar, and McMahan, 2017), devising efficient distributed optimization methods benefiting from differential privacy guarantees (Agarwal, Suresh, Yu, Kumar, and McMahan, 2018), as well as recent guarantees for parallel stochastic optimization with a dependency graph (Woodworth, Wang, Smith, McMahan, and Srebro, 2018).

Another key problem in federated learning which appears more generally in distributed machine learning and other learning setups is that of fairness. In many instances in practice, the resulting learning models may be biased or unfair: they may discriminate against some protected groups (Bickel, Hammel, and O’Connell, 1975; Hardt, Price, Srebro, et al., 2016). As a simple example, a regression algorithm predicting a person’s salary could be using that person’s gender. This is a key problem in modern machine learning that does not seem to have been specifically studied in the context of federated learning.

While many problems related to federated learning have been extensively studied, the key objective of learning in that context seems not to have been carefully examined. We are also not aware of statistical guarantees derived for learning in this scenario. A crucial reason for such questions to emerge in this context is that the target distribution for which the centralized model is learned is unspecified. Which expected loss is federated learning seeking to minimize? Most centralized models for standard federated learning are trained on the aggregate training sample obtained from the subsamples drawn from the clients. Thus, if we denote by ${\mathscr{D}}_{k}$ the distribution associated to client $k$, $m_{k}$ the size of the sample available from that client and $m$ the total sample size, intrinsically, the centralized model is trained to minimize the loss with respect to the uniform distribution

But why should $\overline{\mathscr{U}}$ be the target distribution of the learning model? Is $\overline{\mathscr{U}}$ the distribution that we expect to observe at test time? What guarantees can be derived for the deployed system?

Notice that, in practice, in federated learning, the probability that an individual data source participates in training depends on various factors such as whether the mobile device is connected to the internet or whether it is being charged. Thus, the training data may not truly reflect the usage of the learned model in inference. Additionally, these uncertainties may also affect the size of the sample $m_{k}$ acquired from each client, which directly affects the definition of $\overline{\mathscr{U}}$.

We argue that in many common instances, the uniform distribution is not the natural objective distribution and that seeking to minimize the expected loss with respect to the specific distribution $\overline{\mathscr{U}}$ is risky. This is because the target distribution may be in general quite different from $\overline{\mathscr{U}}$. In many cases, that can result in a suboptimal or even a detrimental performance. For example, imagine a plausible scenario of federated learning where the learner has access to a large population of expensive mobile phones, which are most commonly adopted by software engineers or other technical users (say $70\%$) than other users ($30\%$), and a small population of other mobile phones less used by non-technical users ($5\%$) and significantly more often by other users ($95\%$). The centralized model would then be essentially based on the uniform distribution based on the expensive clients. But, clearly, such a model would not be adapted to the wide general target domain formed by the majority of phones with a $5\%\mathord{-}95\%$ population of general versus technical users. Many other realistic examples of this type can help illustrate the learning problem resulting from a mismatch between the target distribution and $\overline{\mathscr{U}}$. In fact, it is not clear why minimizing the expected loss with respect to $\overline{\mathscr{U}}$ could be beneficial for the clients, whose distributions are ${\mathscr{D}}_{k}$s.

Thus, we put forward a new framework of agnostic federated learning (AFL), where the centralized model is optimized for any possible target distribution formed by a mixture of the client distributions. Instead of optimizing the centralized model for a specific distribution, with the high risk of a mismatch with the target, we define an agnostic and more risk-averse objective. We show that, for some target mixture distributions, the cross-entropy loss of the hypothesis obtained by minimization with respect to the uniform distribution $\overline{\mathscr{U}}$ can be worse, by a constant additive term, than that of the hypothesis obtained in AFL, even if the learner has access to an infinite sample size (Section 3.2).

We further show that our AFL framework naturally yields a notion of fairness, which we refer to as good-intent fairness (Section 3.3). Indeed, the predictor solution of the optimization problem for our AFL framework treats all protected categories similarly. Beyond federated learning, our framework and solution also cover related problems in cloud-based learning services, where customers may not have any training data at their disposal or may not be willing to share that data with the cloud. In that case too, the server needs to train a model without access to the training data. Our framework and algorithm can also be of interest to other learning scenarios such as domain adaptation, drifting, and other contexts where the training and test distributions do not coincide.

The rest of the paper is organized as follows. In Section 2, we give an extensive discussion of related work, including connections with the broad literature of domain adaptation. In Section 3, we give a formal description of the learning scenario of federated learning and the formulation of the problem as AFL. Next, we give a detailed theoretical analysis of learning in the AFL framework, including data-dependent Rademacher complexity generalization bounds (Section 4). These bounds lead to a natural learning algorithm with a regularization term based on a skewness term that we define (Section 5). We also present an efficient convex optimization algorithm for solving the optimization problem defining our algorithm (Section 5.2). Our algorithm is a stochastic gradient-descent solution for minimax problems, for which we give a detailed analysis, including the proof of convergence in terms of the variances of the stochastic gradients. In Section 6, we present a series of experiments comparing our AFL algorithm and solution with existing federated learning solutions. In Section 7, we discuss several extensions of AFL.

Related work

Here, we briefly discuss several learning scenarios and work related to our study of federated learning.

The problem of federated learning is closely related to other learning scenarios where there is a mismatch between the source distribution and the target distribution. This includes the problem of transfer learning or domain adaptation from a single source to a known target domain (Ben-David, Blitzer, Crammer, and Pereira, 2006; Mansour, Mohri, and Rostamizadeh, 2009b; Cortes and Mohri, 2014; Cortes, Mohri, and Muñoz Medina, 2015), either through unsupervised adaptation techniques (Gong et al., 2012; Long et al., 2015; Ganin and Lempitsky, 2015; Tzeng et al., 2015), or via lightly supervised ones (some amount of labeled data from the target domain) (Saenko et al., 2010; Yang et al., 2007; Hoffman et al., 2013; Girshick et al., 2014). This also includes previous applications in natural language processing (Dredze et al., 2007; Blitzer et al., 2007; Jiang and Zhai, 2007; Raju et al., 2018), speech recognition (Legetter and Woodland, 1995; Gauvain and Chin-Hui, 1994; Pietra et al., 1992; Rosenfeld, 1996; Jelinek, 1998; Roark and Bacchiani, 2003), and computer vision (Martínez, 2002)

A problem more closely related to that of federated learning is that of multiple-source adaptation, first formalized and analyzed theoretically by Mansour, Mohri, and Rostamizadeh (2009c, a) and later studied for various applications such as object recognition (Hoffman et al., 2012; Gong et al., 2013a, b). Recently, Zhang et al. (2015) studied a causal formulation of this problem for a classification scenario, using the same combination rules as Mansour et al. (2009c, a). The problem of domain generalization (Pan and Yang, 2010; Muandet et al., 2013; Xu et al., 2014), where knowledge from an arbitrary number of related domains is combined to perform well on a previously unseen domain is very closely related to that of federated learning, though the assumptions about the information available to the learner and the availability of unlabeled data may differ.

In the multiple-source adaptation problem studied by Mansour, Mohri, and Rostamizadeh (2009c, a) and Hoffman, Mohri, and Zhang (2018), each domain $k$ is defined by the corresponding distribution ${\mathscr{D}}_{k}$ and the learner has only access to a predictor $h_{k}$ for each domain and no access to labeled training data drawn from these domains. The authors show that it is possible to define a predictor $h$ whose expected loss ${\mathscr{L}}_{\mathscr{D}}(h)$ with respect to any distribution ${\mathscr{D}}$ that is a mixture of the source domains ${\mathscr{D}}_{k}$ is at most the maximum expected loss of the source predictors: $\max_{k}L_{{\mathscr{D}}_{k}}(h_{{\mathscr{D}}_{k}})$. They also provide an algorithm for determining $h$.

Our learning scenario differs from the one adopted in that work since we assume access to labeled training data from each domain ${\mathscr{D}}_{k}$. Furthermore, the predictor determined by the algorithm of Hoffman, Mohri, and Zhang (2018) belongs to a specific hypothesis set ${\mathscr{H}}^{\prime}$, which is that of distribution weighted combinations of the domain predictors $h_{k}$, while, in our setup, the objective is to determine the best predictor in some global hypothesis set ${\mathscr{H}}$, which may include ${\mathscr{H}}^{\prime}$ as a subset, and which is not depending on some domain-specific predictors.

Our optimization solution also differs from the work of Farnia and Tse (2016) and Lee and Raginsky (2017) on local minimax results, where samples are drawn from a single source ${\mathscr{D}}$, and where the generalization error is minimized over a set of locally ambiguous distributions $\widehat{\mathscr{D}}$, where $\widehat{\mathscr{D}}$ is the empirical distribution. The authors propose this metric for statistical robustness. In our work, we obtain samples from $p$ unknown distributions, and the set of distributions $D_{\lambda}$ over which we optimize the expected loss is fixed and independent of samples. Furthermore, the source distributions can differ arbitrarily and need not be close to each other. In reverse, we note that our stochastic algorithm can be used to minimize the loss functions proposed in (Farnia and Tse, 2016; Lee and Raginsky, 2017).

Learning scenario

In this section, we introduce the learning scenario of agnostic federated learning we consider. Next, we first argue that the uniform solution commonly adopted in standard federated learning may not be an adequate solution, thereby further justifying our agnostic model. Second, we show the benefit of our model in fairness learning.

and by $h_{\mathscr{D}}$ its minimizer: $h_{\mathscr{D}}=\operatorname*{argmin}_{h\in{\mathscr{H}}}{\mathscr{L}}_{\mathscr{D}}(h)$.

We consider a learning scenario where the learner receives $p$ samples $S_{1},\ldots,S_{p}$, with each $S_{k}=((x_{k,1},y_{k,1}),\ldots,(x_{k,m_{k}},y_{k,m_{k}}))\in({\mathscr{X}}\times{\mathscr{Y}})^{m_{k}}$ of size $m_{k}$ drawn i.i.d. from a different domain or distribution ${\mathscr{D}}_{k}$. The learner’s objective is to determine a hypothesis $h\in{\mathscr{H}}$ that performs well on some target distribution. We will also denote by $\widehat{\mathscr{D}}_{k}$ the empirical distribution associated to sample $S_{k}$ of size $m$ drawn from ${\mathscr{D}}^{m}$.

This scenario coincides with that of federated learning where training is done with the uniform distribution over the union of all samples $S_{k}$, that is $\widehat{\mathscr{U}}=\sum_{k=1}^{p}\frac{m_{k}}{\sum_{k=1}^{p}m_{k}}\widehat{\mathscr{D}}_{k}$, and where the underlying assumption is that the target distribution is $\overline{\mathscr{U}}=\sum_{k=1}^{p}\frac{m_{k}}{\sum_{k=1}^{p}m_{k}}{\mathscr{D}}_{k}$. We will not adopt that assumption since it is rather restrictive and since, as discussed later, it can lead to solutions that are disadvantageous to domain users. Instead, we will consider an agnostic federated learning (AFL) scenario where the target distribution can be modeled as an unknown mixture of the distributions ${\mathscr{D}}_{k}$, $k=1,\ldots,p$, that is ${\mathscr{D}}_{\lambda}=\sum_{k=1}^{p}\lambda_{k}{\mathscr{D}}_{k}$ for some $\lambda\in\Delta_{p}$. Since the mixture weight $\lambda$ is unknown, here, the learner must come up with a solution that is favorable for any $\lambda$ in the simplex, or any $\lambda$ in a subset $\Lambda\subseteq\Delta_{p}$. Thus, we define the agnostic loss (or agnostic risk) ${\mathscr{L}}_{{\mathscr{D}}_{\Lambda}}(h)$ associated to a predictor $h\in{\mathscr{H}}$ as

We will extend our previous definitions and denote by $h_{{\mathscr{D}}_{\Lambda}}$ the minimizer of this loss:

In practice, the learner has access to the distributions ${\mathscr{D}}_{k}$ only via the finite samples $S_{k}$. Thus, for any $\lambda\in\Delta_{p}$, instead of the mixture ${\mathscr{D}}_{\lambda}$, only the $\lambda$-mixture of empirical distributions, $\overline{\mathscr{D}}_{\lambda}=\sum_{k=1}^{p}\lambda_{k}\widehat{\mathscr{D}}_{k}$, is accessible.Note, $\overline{\mathscr{D}}_{\lambda}$ is distinct from an empirical distribution $\widehat{\mathscr{D}}_{\lambda}$ which would be based on a sample drawn from ${\mathscr{D}}_{\lambda}$. $\overline{\mathscr{D}}_{\lambda}$ is based on samples drawn from ${\mathscr{D}}_{k}$s. This leads to the definition of ${\mathscr{L}}_{\overline{\mathscr{D}}_{\Lambda}}(h)$, the agnostic empirical loss of a hypothesis $h\in{\mathscr{H}}$ for a subset of the simplex $\Lambda$:

We will denote by $h_{\overline{\mathscr{D}}_{\Lambda}}$ the minimizer of this loss: $h_{\overline{\mathscr{D}}_{\Lambda}}=\operatorname*{argmin}_{h\in\mathcal{H}}{\mathscr{L}}_{\overline{\mathscr{D}}_{\Lambda}}(h)$. In the next section, we will present generalization bounds relating the expected and empirical agnostic losses ${\mathscr{L}}_{{\mathscr{D}}_{\Lambda}}(h)$ and ${\mathscr{L}}_{\overline{\mathscr{D}}_{\Lambda}}(h)$ for all $h\in{\mathscr{H}}$.

Notice that the domains ${\mathscr{D}}_{k}$ discussed thus far need not coincide with the clients. In fact, when the number of clients is very large and $\Lambda$ is the full simplex, $\Lambda=\Delta_{p}$, it is typically preferable to consider instead domains defined by clusters of clients, as discussed in Section 7. On the other hand, if $p$ is small or $\Lambda$ more restrictive, then the model may not perform well on certain domains of interest. We mitigate the effect of large $p$ values using a suitable regularization term derived from our theory.

2 Comparison with federated learning

Here, we further argue that the uniform solution $h_{\overline{\mathscr{U}}}$ commonly adopted in federated learning may not provide a satisfactory performance compared with a solution of the agnostic problem. This further motivates our AFL model.

As already discussed, since the target distribution is unknown, the natural method for the learner is to select a hypothesis minimizing the agnostic loss ${\mathscr{L}}_{{\mathscr{D}}_{\Lambda}}$. Is the predictor minimizing the agnostic loss coinciding with the solution $h_{\widehat{\mathscr{U}}}$ of standard federated learning? How poor can the performance of the standard federated learning be? We first show that the loss of $h_{\widehat{\mathscr{U}}}$ can be higher than that of the optimal loss achieved by $h_{{\mathscr{D}}_{\Lambda}}$ by a constant loss, even if the number of samples tends to infinity, that is even if the learner has access to the distributions ${\mathscr{D}}_{k}$ and uses the predictor $h_{\overline{\mathscr{U}}}$. Similar results are known for universal compression, where the goal is to compress a sequence of random variables without knowledge of the generating distribution (Grünwald, 2007).

Consider the following two distributions with support reduced to a single element $x\in{\mathscr{X}}$ and two classes ${\mathscr{Y}}=\{0,1\}$: ${\mathscr{D}}_{1}(x,0)=0$, ${\mathscr{D}}_{2}(x,1)=1$, ${\mathscr{D}}_{2}(x,0)=\frac{1}{2}$, and ${\mathscr{D}}_{2}(x,1)=\frac{1}{2}$. Let $\Lambda=\{\delta_{1},\delta_{2}\}$, where $\delta_{k}$, $k=1,2$, denotes the Dirac measure on index $k$. We will consider the case where the sample sizes $m_{k}$ are all equal, that is $h_{\overline{\mathscr{U}}}=\frac{1}{2}({\mathscr{D}}_{1}+{\mathscr{D}}_{2})$. Let $p_{0}$ denote the probability that $h$ assigns to class and $p_{1}$ the one it assigns to class $1$. Then, the cross-entropy loss of a predictor $h$ can be expressed as follows:

where the last inequality follows the non-negativity of the relative entropy. Furthermore, equality is achieved when $p_{0}=1-p_{1}=\frac{1}{4}$, which defines $h_{\overline{\mathscr{U}}}$, the minimizer of ${\mathscr{L}}_{\overline{\mathscr{U}}}(h)$. In view of that, $\mathcal{L}_{{\mathscr{D}}_{\Lambda}}(h_{\overline{\mathscr{U}}})$ is given by the following:

We now compute the loss of $h_{{\mathscr{D}}_{\Lambda}}$:

since $\frac{1}{2}$ is the solution of the convex optimization in $p_{1}$, in view of \max\Big{\{}\frac{1}{p_{1}},\frac{1}{\sqrt{p_{1}(1-p_{1})}}\Big{\}}=\frac{1}{\sqrt{p_{1}(1-p_{1})}}\leq\frac{1}{2} for $p_{1}>\frac{1}{2}$.

3 Good-intent fairness in learning

Here, we further discuss the relationship between our model of AFL and fairness in learning.

Fairness in machine learning has received much attention in recent past (Bickel et al., 1975; Hardt et al., 2016). There is now a broad literature on the topic with a variety of definitions of the notion of fairness. In a typical scenario, there is a protected class $c$ among $p$ classes $c_{1},c_{2},\ldots,c_{p}$. While there are many definitions of fairness, the main objective of a fairness algorithm is to reduce bias and ensure that the model is fair to all the $p$ protected categories, under some definition of fairness. The most common reasons for bias in machine learning algorithms are training data bias and overfitting bias. We first provide a brief explanation and illustration for both:

the training data is biased: consider the regression task, where the goal is to predict the salary of a person based on features such as education, location, age, gender. Let gender be the protected class. If in the training data, there is a consistent discrimination against women irrespective of their education, e.g., their salary is lower, then we can conclude that the training data is inherently biased.

the training procedure is biased: consider an image recognition task where the protected category is race. If the model is heavily trained on images based on certain races, then the resulting model will be biased because of over-fitting.

Our model of AFL can help define a notion of good-intent fairness, where we reduce the bias in the training procedure. Furthermore, if training procedure bias exists, it naturally highlights it.

Suppose we are interested in a classification problem and there is a protected feature class $c$, which can be one of $p$ values $c_{1},c_{2},\ldots,c_{p}$. Then, we define ${\mathscr{D}}_{k}$ as the conditional distribution with the protected class being $c_{k}$. If ${\mathscr{D}}$ is the true underlying distribution, then

Let $\Lambda=\{\delta_{k}\colon k\in[p]\}$ be the collection of Dirac measures over the indices $k$ in $[p]$. With this definition, we define a good-intent fairness algorithm as one seeking to minimize the agnostic loss ${\mathscr{L}}_{{\mathscr{D}}_{\Lambda}}$. Thus, the objective of the algorithm is to minimize the maximum loss incurred on any of the underlying protective classes and hence does not overfit the data to any particular model at the cost of others. Furthermore, it does not degrade the performance of the other classes so long as it does not affect the loss of the most-sensitive protected category. We further note that our approach does not reduce bias in the training data and is useful only for mitigating the training procedure bias.

Learning bounds

where $S_{k}=((x_{k,1},y_{k,1}),\ldots,(x_{k,m_{k}},y_{k,m_{k}}))$ is a sample of size $m_{k}$ and ${\boldsymbol{\sigma}}=(\sigma_{k,i})_{\begin{subarray}{c}k\in[p],i\in[m_{k}]\end{subarray}}$ a collection of Rademacher variables, that is uniformly distributed random variables taking values in $\{-1,+1\}$. We also defined the minimax weighted Rademacher complexity for a subset $\Lambda\subseteq\Delta_{p}$ by

Let \overline{\mathbf{m}}=\frac{{\mathbf{m}}}{m}=\big{(}\frac{m_{1}}{m},\ldots,\frac{m_{p}}{m}\big{)} denote the empirical distribution over $\Delta_{p}$ defined by the sample sizes $m_{k}$, where $m=\sum_{k=1}^{p}m_{k}$. We define the skewness of $\Lambda$ with respect to $\overline{\mathbf{m}}$ by

where $C(\Lambda,\epsilon)$ is a set of distributions $\Lambda^{\prime}$ such that for every $\lambda\in\Lambda$, there exists $\Lambda^{\prime}$ such that $\sum^{p}_{k=1}|\lambda_{k}-\lambda^{\prime}_{k}|\leq\epsilon$.

Our first learning guarantee is presented in terms of ${\mathfrak{R}}_{\mathbf{m}}({\mathscr{G}},\Lambda)$, the skewness parameter ${\mathfrak{s}}(\Lambda\,\|\,\overline{\mathbf{m}})$ and the $\epsilon$-cover $\Lambda_{\epsilon}$.

The proof is an extension of the standard proofs for Rademacher complexity generalization bounds (Koltchinskii and Panchenko, 2002; Mohri et al., 2018). Fix $\lambda\in\Lambda$. For any sample $S=S_{1},\ldots,S_{p}$, define $\Psi(S_{1},\ldots,S_{p})$ by

Let $S^{\prime}=(S^{\prime}_{1},\ldots,S^{\prime}_{p})$ be a sample differing from $S=(S_{1},\ldots,S_{p})$ only by point $x^{\prime}_{k,i}$ in $S^{\prime}_{k}$ and $x_{k,i}$ in $S_{k}$. Then, since the difference of suprema over the same set is bounded by the supremum of the differences, we can write

Thus, by McDiarmid’s inequality, for any $\delta>0$, the following inequality holds with probability at least $1-\delta$ for any $h\in{\mathscr{H}}$:

Therefore, by the union over $\Lambda_{\epsilon}$, with probability at least $1-\delta$, for any $h\in{\mathscr{H}}$ and $\lambda\in\Lambda_{\epsilon}$ the following holds:

By definition of $\Lambda_{\epsilon}$, for any $\lambda\in\Lambda$, there exists $\lambda^{\prime}\in\Lambda_{\epsilon}$ such that ${\mathscr{L}}_{{\mathscr{D}}_{\lambda}}(h)\leq{\mathscr{L}}_{{\mathscr{D}}_{\lambda}^{\prime}}(h)+M\epsilon$. In view of that, with probability at least $1-\delta$, for any $h\in{\mathscr{H}}$ and $\lambda\in\Lambda$ the following holds:

The expectation appearing on the right-hand side can be bounded following standard proofs for Rademacher complexity upper bounds (see for example (Mohri et al., 2018)), leading to

The sum $\sum^{p}_{k=1}\frac{\lambda^{2}_{k}}{m_{k}}$ can be expressed in terms of the skewness of $\Lambda$, using the following equalities:

It can be proven that the skewness parameter appears in a lower bound on the generalization bound. We will include that result in the final version of this paper. The theorem yields immediately upper bounds for agnostic losses by taking the maximum over $\lambda\in\Lambda$: for any $\delta>0$, with probability at least $1-\delta$, for any $h\in{\mathscr{H}}$,

The following result shows that, for a family of functions taking values in $\{-1,+1\}$, the Rademacher complexity ${\mathfrak{R}}_{\mathbf{m}}({\mathscr{G}},\Lambda)$ can be bounded in terms of the VC-dimension and the skewness of $\Lambda$.

For any ${\mathbf{a}}\in A_{\lambda}$, $\|{\mathbf{a}}\|_{2}=\sqrt{\sum_{k=1}^{p}m_{k}\frac{\lambda_{k}^{2}}{m_{k}^{2}}}=\sqrt{\sum_{k=1}^{p}\frac{\lambda_{k}^{2}}{m_{k}}}\leq\sqrt{\frac{{\mathfrak{s}}(\Lambda\,\|\,\overline{\mathbf{m}})}{m}}$. Then, by Massart’s lemma, for any $\lambda\in\Lambda$, the following inequalities hold:

By Sauer’s lemma, the following holds for $m\geq d$: $|A_{\lambda}|\leq\left(\frac{em}{d}\right)^{d}$. Plugging in the right-hand side in the inequality above completes the proof.

Both Lemma 4.3 and the generalization bound of Theorem 4.1 can thus be expressed in terms of the skewness parameter ${\mathfrak{s}}(\Lambda\,\|\,\overline{\mathbf{m}})$. Note that modulo the skewness parameter, the results look very similar to standard generalization bounds (Mohri et al., 2018). Furthermore, when $\Lambda$ contains only one distribution and is the average distribution, that is $\lambda_{k}=m_{k}/m$, then the skewness is equal to one and the results coincide with the standard guarantees in supervised learning.

Theorem 4.1 and Lemma 4.3 also provide guidelines for choosing the domains and $\Lambda$. When $p$ is large and $\Lambda=\Delta_{p}$, then, the number of samples per domain could be small, the skewness parameter ${\mathfrak{s}}(\Lambda\,\|\,\overline{\mathbf{m}})=\max_{1\leq k\leq p}\frac{1}{m_{k}}$ would then be large and the generalization guarantees for the model would become weaker. We suggest some guidelines for choosing domains in Section 7. We further note that for a given $p$, if $\Lambda$ contains distributions that are close to $\overline{\mathbf{m}}$, then the model generalizes well.

One straightforward choice of the parameter $\epsilon$ is $\epsilon=\frac{1}{\sqrt{m}}$, but, depending on $|\Lambda_{\epsilon}|$ and other tperms of the bound, more favorable choices may be possible. We conclude this section by adding that alternative learning bounds can be derived for this problem, as discussed in Appendix A.

Algorithm

In this section, we introduce a learning algorithm for agnostic federated learning using the guarantees proven in the previous section and discuss in detail an optimization solution.

The learning guarantees of the previous section suggest minimizing the asum of the empirical AFL term ${\mathscr{L}}_{\overline{\mathscr{D}}_{\Lambda}}(h)$, a term controlling the complexity of ${\mathscr{H}}$ and a term depending on the skewness parameter. Observe that, since ${\mathscr{L}}_{\overline{\mathscr{D}}_{\lambda}}(h)$ is linear in $\lambda$, the following equality holds:

where $\operatorname*{conv}(\Lambda)$ is the convex hull of $\Lambda$. Assume that ${\mathscr{H}}$ is a vector space that can be equipped with a norm $\|\cdot\|$, as with most hypothesis sets used in learning applications. Then, given $\Lambda$ and the regularization parameters $r\geq 0$ and $\gamma\geq 0$, our learning guarantees suggest minimizing the regularized loss ${\mathscr{L}}_{\overline{\mathscr{D}}_{\Lambda_{r}}}(h)+\gamma\|h\|$, where $\|\cdot\|$ is a suitable norm controlling the complexity of ${\mathscr{H}}$ and where $\Lambda_{r}$ is defined by $\Lambda_{r}=\{\lambda\in\operatorname*{conv}(\Lambda)\colon 1+\chi^{2}(\lambda\,\|\,\overline{\mathbf{m}})\leq r\}$. This can be equivalently formulated as the following minimization problem:

where $\mu\geq 0$ is a hyperparameter. This defines our algorithm for AFL.

2 Optimization algorithm

where $\mathsf{L}_{k}(w)$ stands for ${\mathscr{L}}_{\widehat{\mathscr{D}}_{k}}(h)$, the empirical loss of hypothesis $h\in{\mathscr{H}}$ (corresponding to $w$) on domain $k$:

Since the regularization terms do not make the optimization problem harder, to simplify the discussion, we will consider the unregularized version of problem (6). Thus, we will study the following problem given by the set of variables $w$:

Observe that problem (8) admits a natural game-theoretic interpretation as a two-player game, where nature selects $\lambda\in\Lambda$ to maximize the objective, while the learner seeks $w\in{\mathscr{W}}$ minimizing the loss. We are interested in finding the equilibrium of this game, which is attained for some $w^{*}$, the minimizer of Equation 8 and $\lambda^{*}\in\Lambda$, the hardest domain mixture weights. At the equilibrium, moving $w$ away from $w^{*}$ or $\lambda$ from $\lambda^{*}$, increases the objective function. Hence, $\lambda^{*}$ can be viewed as the center of $\Lambda$ in the manifold imposed by the loss function $\mathsf{L}$, whereas $\overline{\mathscr{U}}$, the empirical distribution of samples, may lie elsewhere, as illustrated by Figure 2.

By Equation (5), using the set $\operatorname*{conv}(\Lambda)$ instead of $\Lambda$ does not affect the solution of the optimization problem. In view of that, in what follows, we will assume, without loss of generality, that $\Lambda$ is a convex set. Observe that, since $\mathsf{L}_{k}(w)$ is not an average of functions, standard stochastic gradient descent algorithms cannot be used to minimize this objective. We will present instead a new stochastic gradient-type algorithm for this problem.

Let $\nabla_{w}\mathsf{L}(w,\lambda)$ denote the gradient of the loss function with respect to $w$ and $\nabla_{\lambda}\mathsf{L}(w,\lambda)$ the gradient with respect to $\lambda$. Let $\delta_{w}\mathsf{L}(w,\lambda)$, and $\delta_{\lambda}\mathsf{L}(w,\lambda)$ be unbiased estimates of the gradient, that is,

We first give an optimization algorithm Stochastic-AFL for the AFL problem, assuming access to such unbiased estimates. The pseudocode of the algorithm is given in Figure 3. At each step, the algorithm computes a stochastic gradient with respect to $\lambda$ and $w$ and updates the model accordingly. It then projects $\lambda$ to $\Lambda$ by computing a value in $\Lambda$ via convex minimization. If $\Lambda$ is the full simplex, then there is a near-linear time algorithm for this projection Wang and Carreira-Perpinán (2013). It then repeats the process for $T$ steps and return the average of the weights. We provide guarantees for this algorithm in terms of the variance of the stochastic gradients when the loss function $\mathsf{L}$ is convex and when the set of $w$s, ${\mathscr{W}}$, is a compact set.

In the above analysis and in algorithm description in 3, we have ignored the regularization term. If the objective contains a regularization term such as Equation 6, then for $\lambda_{k}$, the regularization term yields a derivative of $-2\gamma\lambda_{k}/\overline{\mathbf{m}}_{k}$, which can be added to $\delta_{\lambda}\mathsf{L}(w,\lambda)$ in Step $3$ in Algorithm 3.

There are several natural candidates for the sampling method defining stochastic gradients. We highlight two techniques: PerDomain gradient and Weighted gradient. We analyze the time complexity and give bounds on the variance for both techniques in Lemmas 5.6 and 5.8 respectively.

Recently, Rakhlin and Sridharan (2013) and Daskalakis et al. (2017) gave an optimistic gradient descent algorithm for minimax optimizations. Our algorithm can also be modified to derive a stochastic optimistic algorithm, which we refer to as Optimistic-Stochastic-AFL. The pseudocode of this algorithm is also given in Figure 3. However, the convergence analysis we present in the next section does not cover this algorithm.

3 Analysis

Throughout this section, for simplicity, we adopt the notation introduced for Equation 7. Our convergence guarantees hold under the following assumptions, which are similar to those adopted for the convergence proof of gradient descent-type algorithms.

Assume that the following properties hold for the loss function $\mathsf{L}$ and sets ${\mathscr{W}}$ and $\Lambda\subseteq\Delta_{p}$:

Convexity: $w\mapsto\mathsf{L}(w,\lambda)$ is convex for any $\lambda\in\Lambda$.

Compactness: $\max_{\lambda\in\Lambda}\|\lambda\|_{2}\leq R_{\Lambda}$ and $\max_{w\in{\mathscr{W}}}\|w\|_{2}\leq R_{{\mathscr{W}}}$, for some $R_{\Lambda}>0$ and $R_{{\mathscr{W}}}>0$.

Bounded gradients: $\|\nabla_{w}\mathsf{L}(w,\lambda)\|_{2}\leq G_{w}$ and $\|\nabla_{\lambda}\mathsf{L}(w,\lambda)\|_{2}\leq G_{\lambda}$ for all $w\in{\mathscr{W}}$ and $\lambda\in\Lambda$.

Time complexity: $U_{w}$ denotes the time complexity of computing $\delta_{w}\mathsf{L}(w,\lambda)$, $U_{\lambda}$ that of computing $\delta_{\lambda}\mathsf{L}(w,\lambda)$, $U_{p}$ that of the projection, and $d$ denotes the dimensionality of ${\mathscr{W}}$.

Assume that the Properties 1 hold. Then, for the steps sizes $\gamma_{w}=\frac{2R_{{\mathscr{W}}}}{\sqrt{T(\sigma^{2}_{w}+G^{2}_{w})}}$ and $\gamma_{\lambda}=\frac{2R_{\Lambda}}{\sqrt{T(\sigma^{2}_{\lambda}+G^{2}_{\lambda})}}$, the following guarantee holds for Stochastic-AFL:

and the time complexity of the algorithm is in $\mathcal{O}((U_{\lambda}+U_{w}+U_{p}+d+k)T)$.

The time complexity of the algorithm follows the definitions of the complexity terms $U_{\lambda}$, $U_{w}$, and $U_{p}$ the dimension $d$ in Properties 1. To prove the convergence guarantee, we make a series of reductions. Let $w^{A}$ and $\lambda^{A}$ be a solution returned by the algorithm. First observe that since $\mathsf{L}$ is convex in $w$ and linear and thus concave in $\lambda$, by the generalized von Neumann’s theorem, the following holds:

Next, since the function is linear in $\lambda$ and convex in $w$,

In view of these inequalities, by the subadditivity of $\max$, the following inequality holds:

We now bound each of the terms above separately. For the first term, observe that for any $w\in{\mathscr{W}}$,

Since the right-hand side does not depend on $w$, taking the maximum of both sides over $w\in{\mathscr{W}}$ and the expectation yields

Similarly, using the projection property, the following inequality can be shown:

For the second term, by the Cauchy-Schwarz inequality, we can write

Taking the expectation of both sides and using Jensen’s inequality yields

For the third term, observe that the stochastic gradients at time $t$ are unbiased, conditioned on $\lambda_{t}$, and $w_{t}$, hence,

Combining the upper bounds just derived gives:

Setting $\gamma_{w}=\frac{2R_{{\mathscr{W}}}}{\sqrt{T((\sigma^{2}_{w}+G^{2}_{w}))}}$ and $\gamma_{\lambda}=\frac{2R_{\Lambda}}{\sqrt{T((\sigma^{2}_{\lambda}+G^{2}_{\lambda}))}}$ to minimize this upper bound completes the proof.

4 Stochastic gradients

For any $w\in{\mathscr{W}}$, $\lambda\in\Lambda$ and $k\in[p]$, the gradient with respect to $\lambda_{k}$ is given by

We first discuss the stochastic gradients for $\lambda$. Notice that the gradient for $\lambda$ is independent of $\lambda$. Thus, a natural choice for the stochastic gradient with respect to $\lambda$ is based on uniformly sampling a domain $K\in[p]$ and then sampling $x_{K,i}$ from domain $K$. This leads to the definition of the stochastic gradient $\delta_{\lambda}\mathsf{L}(w,\lambda)$ shown in Figure 4. The following lemma bounds the variance for that definition of $\delta_{\lambda}\mathsf{L}(w,\lambda)$.

The stochastic gradient $\delta_{\lambda}\mathsf{L}(w,\lambda)$ is unbiased. Further, if the loss function is bounded by $M$, then the following upper bound holds for the variance of $\delta_{\lambda}\mathsf{L}(w,\lambda)$:

The unbiasedness of $\delta_{\lambda}\mathsf{L}(w,\lambda)$ follows directly its definition. For the variance, observe that, for index $k\in[p]$, since the probability of not drawing domain $k$ is $(1-\frac{1}{p})$, the variance is given by the following

Summing over all indices from $k\in[p]$ completes the proof.

If the above variance is too high, then we can sample one $J_{k}$ for every domain $k$. This is same as computing the gradient of a batch and reduces the variance by a factor of $p$.

The gradient with respect to $w$ depends both on $\lambda$ and $w$. There are two natural stochastic gradients: the PerDomain-stochastic gradient and the Weighted-stochastic gradient. For a PerDomain-stochastic gradient, we sample an element uniformly from $[m_{k}]$ for each $k\in[p]$. For the Weighted-stochastic gradient, we sample a domain according to $\lambda$ and sample an element out of it. To bound the variance of these two stochastic gradients, we need a few definitions.

the intra-domain variance with respect to $w$ is defined as follows:

the outer-domain variance with respect to $w$ is defined as follows:

the time complexity of computing the loss and gradient with respect to $w$ for a single sample is denoted by $U$.

With these definitions, we can bound the variance of both PerDomain and Weighted stochastic gradients.

PerDomain stochastic gradient is unbiased and runs in time $pU+\mathcal{O}(p\log m)$ and the variance satisfy,

The time complexity and the unbiasedness follow from the definitions. We now bound the variance. Since $\nabla_{w}\mathsf{L}_{k,J_{k}}$ is an unbiased estimate of $\nabla_{w}\mathsf{L}_{k}(w)$ and we have:

Weighted stochastic gradient is unbiased and runs in time $U+\mathcal{O}(k+\log n)$ and the variance satisfy the following inequality:

The time complexity and the unbiasedness follow from the definitions. We now bound the variance. By definition for any $w,\lambda$,

where the second equality follows from the unbiasedness of the stochastic gradients.

Since $R_{\Lambda}\leq 1$, at first glance, the above two lemmas may suggest that PerDomain stochastic is always better than Weighted stochastic gradient. Note, however, that the time complexities of the algorithms is dominated by $U$ and thus, the time complexity of PerDomain-stochastic gradient is roughly $k$ times larger than that of Weighted-stochastic gradient. Hence, if $k$ is small, it is preferable to choose the PerDomain-stochastic gradient.

For large values of $p$, to do a fair comparison, we need to average $p$ independent copies of the Weighted-stochastic gradient, which we refer to as $p$-Weighted, and compare it with the PerDomain-stochastic gradient. Since the variance of average of $p$ i.i.d. random variables is $1/p$ times the individual variance, by Lemma 5.8, the following holds:

Further, observe that $R_{\Lambda}=\max_{\lambda\in\Lambda}\sum^{p}_{k=1}\lambda^{2}_{k}\geq\frac{1}{p}$. Thus,

Hence, the right choice of the stochastic variance of $w$ depends on the application. If all domains are roughly equally weighted, then we have $R(\Lambda)\approx\frac{1}{p}$ and the PerDomain-variance is a more favorable choice. Otherwise, if $\sigma^{2}_{O}(w)$ is small, then the Weighted-stochastic gradient is more favorable.

Experiments

To study the benefits of our AFL algorithm, we carried out experiments with three datasets. Even though our optimization convergence guarantees hold only for convex functions and stochastic gradient, we show that our domain-agnostic learning performs well for non-convex functions and variants of stochastic gradient descent such as momentum and Adagrad too.

In all the three experiments, we compare the domain agnostic model with the model trained with $\widehat{\mathscr{U}}$, the uniform distribution over the union of samples, and the models trained on individual domains. In all of these experiments, we used PerDomain stochastic gradients and set $\Lambda=\Delta_{p}$. All algorithms were implemented in Tensorflow (Abadi et al., 2015).

The Adult dataset is a census dataset from the UCI Machine Learning Repository (Blake, 1998). It contains $32\mathord{,}561$ training samples with numerical and categorical features, each representing a person. The task consists of predicting if the person’s income exceeds \50\mathord{,}000$. We split this dataset into two domains depending on whether the person had a doctorate degree or not, resulting into domains: the doctorate domain containing$413$examples and the non-doctorate domain containing$32\mathord{,}148$examples. We trained a logistic regression model with just the categorical features and Adagrad optimizer. The performance of the models averaged over$50$runs is reported in Table 1. The performance on${\mathscr{D}}_{\Lambda}$of the model trained with$\widehat{\mathscr{U}}$, that is standard federated learning, is about$69.6\%$. In contrast, the performance of our AFL model is at least about$71.5\%$on any target distribution${\mathscr{D}}_{\lambda}$. The uniform average over the domains of the test accuracy of the AFL model is slightly less than that of the uniform model, but the agnostic model is less biased and performs better on${\mathscr{D}}_{\Lambda}$. Furthermore, of the two domains, the doctorate domain is the harder one for predictions. For this domain, the performance of the domain agnostic model is close to the model trained only on doctorate data and is better than that of the model trained with the uniform distribution$\widehat{\mathscr{U}}$.

2 Fashion MNIST

The Fashion MNIST dataset, originally announced by Xiao et al. (2017), is an MNIST-like dataset where images are classified into $10$ categories of clothing, instead of handwritten digits. The dataset includes $60\mathord{,}000$ training images and $10\mathord{,}000$ test images given as 28x28 arrays of grayscale pixel intensities, spread evenly among the ten categories. We first trained a simple logistic regression classifier and observed that the lowest performance was achieved for the following three categories: t-shirt/top, pullover, and shirt. Next, we extracted the subset of the data labeled with these three categories and split this subset into three domains, each consisting of one class of clothing. We then trained a classifier for the three classes using logistic regression and the Adam optimizer. The results are shown in Table 2. Since here the domain uniquely identifies the label, in this experiment, we did not compare against models trained on specific domains. Of the three domains or classes, the shirt class is the hardest one to distinguish from others. The domain-agnostic model improves the performance for shirt more than it degrades it on pullover and shirt, leading to both shirt-specific and overall accuracy improvement when compared to the model trained with the uniform distribution $\widehat{\mathscr{U}}$. Furthermore, in this experiment, note that our agnostic learning solution not only improves the loss of the worst domain, but also generalizes better and hence improves the average test accuracy. Our AFL model achieves a performance of about $\%74.5$ on any target distribution ${\mathscr{D}}_{\lambda}$, while the performance of standard federated learning can be as low as about $\%71.2$.

3 Language models

Motivated by the keyboard application (Hard et al., 2018), where a single client uses a trained language model in multiple environments such as chat apps, email, and web input, we created a dataset that combines two very different types of language datasets: conversation and document. For conversation, we used the Cornell movie dataset that contain movie dialogues Danescu-Niculescu-Mizil and Lee (2011). This dataset contains about $300\mathord{,}000$ sentences with an average sentence length of $8$. For documents, we used the Penn TreeBank (PTB) dataset that contains approximately $50\mathord{,}000$ sentences with an average sentence length of $20$ Marcus et al. (1993). We created a single dataset by combining both of the above corpuses, with conversation and document as domains. We preprocessed the data to remove punctuations, capitalized the data uniformly, and computed a vocabulary of $10\mathord{,}000$ most frequent words. We trained a two-layer LSTM model with LSTM and projection size of $512$ with momentum optimizer. The performance of the models are measured by their perplexity, that is the exponent of cross-entropy loss. The results are reported in Table 3. Of the two domains, the document domain is the one admitting the higher perplexity. For this domain, the test perplexity of the domain agnostic model is close to that of the model trained only on document data and is better than that of the model trained with the uniform distribution $\widehat{\mathscr{U}}$.

Extensions

In this section, we briefly discuss several extensions of the framework, theory and algorithms that we presented.

The choice of the domains can significantly impact learnability in federated learning. In view of our learning bounds, if the number of domains, $p$, is large and $\Lambda$ is the full simplex, $\Lambda=\Delta_{p}$, then the models may not generalize well. Thus, if the number of clients is very large, using each client as a domain may be a poor choice for better generalization. Ideally, each domain is represented with a sufficiently large number of samples and is relatively homogeneous or pure. This suggests using a clustering algorithm for defining the domains based on the similarity of the client distributions. Different Bregman divergences could be used to define the divergence or similarity between distributions. Thus, techniques such as those of Banerjee, Merugu, Dhillon, and Ghosh (2005) could be used to determine clusters of clients using a suitable Bregman divergence.

Client clusters can also be determined based on domain expertise. For example, in federated keyboard next word prediction (Hard et al., 2018), domains can be chosen to be the native language of the clients. If the model is used in variety of applications, domains can also be based on the application of interest. For example, the keyboard in (Hard et al., 2018) is used in chat apps, social apps, and web inputs. Here, domains can be the app that was used. Training models agnostically ensures that the user experience is favorable in all apps.

2 Incorporating a prior on ΛΛ\Lambda

In this formulation, larger weights are assigned to more likely domains. The generalization guarantees of Theorem 4.1 can be appropriately modified to include these changes. Furthermore, if the prior $p_{\Lambda}(\lambda)$ is a log-concave function of $\lambda$, then the new objective is convex in $h$ and concave in $\lambda$ and a slight modification of our proposed algorithm can be used to determine the global minima. We note that we could also adopt a multiplicative formulation with the prior multiplying the loss, instead of the additive one with the negative log of the probability in Equation 9.

3 Domain features and personalization

We studied agnostic federated learning, where we learn a model that performs well on all domains. First, notice that we do not make any assumption on the hypothesis set ${\mathscr{H}}$ and the hypotheses can use the domain $k$ as a feature. Such models could be useful for applications where the target domain is known at inference time. Second, while the paper deals with learning a centralized model, the resulting model $h_{{\mathscr{D}}_{\Lambda}}$ can be combined with a personalized model, on the client’s machine, to design better client-specific models. This can be done for example by learning an appropriate mixture weight $\alpha_{k}\in$ to use a mixture $\alpha_{k}h_{{\mathscr{D}}_{\Lambda}}+(1-\alpha_{k})h_{k}$ of the domain agnostic centralized model $h_{{\mathscr{D}}_{\Lambda}}$ and a client- or domain-specific model $h_{k}$.

Conclusion

We introduced a new framework of AFL for which we presented a detailed theoretical analysis. We also gave an algorithm for this problem benefiting from our theoretical analysis, as well as a new stochastic optimization solution needed for large-scale problems. Our experimental results suggest that our solution can lead to significant benefits in practice.

Acknowledgements

We thank Shankar Kumar, Rajiv Mathews, and Brendan McMahan for helpful comments and discussions.

References

Appendix A Alternative learning guarantees

An objective similar to that of AFL was considered in the context of multiple source domain adaptation by Liu et al. (2015). The authors presented generalization bounds for a scenario where the target is based on some specific mixture $\lambda$ of the source domains. Our theoretical results differ from those of this work in two ways. First, our generalization bounds do not hold for a single mixture weight $\lambda$ but for any subset $\Lambda$ of the simplex. Second, the complexity terms in the bounds presented by these authors are proportional to $\sqrt{m}\max_{k\in[p]}\frac{\lambda_{k}}{m_{k}}$, while our guarantees are in terms of $\sqrt{\sum_{k=1}^{p}\frac{\lambda^{2}_{k}}{m_{k}}}$, which is strictly tighter. In particular, in the special case where $k=2$, $\lambda_{1}=\frac{1}{\sqrt{m}}$, $\lambda_{2}=1-\lambda_{1}$ and $m_{1}=1$ and $m_{2}=m-1$, the bounds of Liu et al. (2015) are proportional to a constant and thus not informative, $\sqrt{m}\max_{k\in[p]}\frac{\lambda_{k}}{m_{k}}=1$, while our guarantees are in terms of $\frac{1}{\sqrt{m}}$.

Our generalization error in Theorem 4.1 is particularly useful when $\Lambda$ is a strict subset of the simple, $\Lambda\subset\Delta_{p}$. If $\Lambda=\Delta_{p}$, we can give the following alternative learning guarantee based.

For any $\delta>0$, with probability at least $1-\delta$ over the draw of samples $S_{k}\sim{\mathscr{D}}_{k}^{m_{k}}$, the following inequality holds for all $h\in{\mathscr{H}}$ and $\lambda\in\Lambda$:

where ${\mathfrak{R}}^{k}_{m_{k}}({\mathscr{G}})$ is the Rademacher complexity over domain ${\mathscr{D}}_{k}$ with $m_{k}$ samples.

The proof is a direct application of known Rademacher complexity bounds (Mohri et al., 2018) and a union bound and is omitted.

To relate the generalization bounds of Theorem 4.1 and Theorem A.1, observe that, by the sub-additivity of $\sup$ and the linearity of expectation, the following inequality holds:

Furthermore, by the sub-additivity of $\sqrt{\cdot}$, the following inequality holds:

Hence, up to the logarithmic factors in the second term, the guarantee of Theorem 4.1 is stronger than that of Theorem A.1. However, $\Lambda_{\epsilon}$ can be large and exponential in $p$, and it is not clear which of the bounds are stronger in general. This depends on $\overline{\mathbf{m}}$ and $\lambda$. Deriving learning bounds that improve upon both of the learning bounds above remains an interesting open question.