Convergence and Accuracy Trade-Offs in Federated Learning and Meta-Learning

Zachary Charles, Jakub Konečný

Introduction

Federated learning (McMahan et al., 2017) is a distributed framework for learning models without directly sharing data. In this framework, clients perform local updates (typically using first-order optimization) on their own data. In the popular FedAvg algorithm (McMahan et al., 2017), the client models are then averaged at a central server. Since the proposal of FedAvg, many new federated optimization algorithms have been developed (Li et al., 2020a; Reddi et al., 2020; Hsu et al., 2019; Xie et al., 2019; Basu et al., 2019; Li et al., 2020b; Karimireddy et al., 2019). These methods typically employ multiple local client epochs in order to improve communication-efficiency. We defer to Kairouz et al. (2019) and Li et al. (2019) for more detailed summaries of federated learning.

Local updates have also been used extensively in meta-learning. The celebrated MAML algorithm (Finn et al., 2017) employs multiple local model updates on a set of tasks in order to learn a model that quickly adapt to new tasks. MAML has inspired a number of model-agnostic meta-learning methods that also employ first-order local updates (Balcan et al., 2019; Fallah et al., 2020a; Nichol et al., 2018; Zhou et al., 2019). There are strong connections between federated learning and meta-learning, despite differences in practical concerns. Formal connections between the two were shown by Khodak et al. (2019) and have since been explored in many other works (Jiang et al., 2019; Fallah et al., 2020b).

We refer to methods that utilize multiple local updates across clients (or in the language of meta-learning, tasks) as local update methods (see Section 2.1 for a formal characterization). In practice, local update methods frequently outperform “centralized” methods such as SGD (McMahan et al., 2017; Finn et al., 2017; Hard et al., 2018; Yang et al., 2018; Hard et al., 2020). However, the empirical benefits of local update methods are not fully explained by existing theoretical analyses. For example, Woodworth et al. (2020) show that FedAvg often obtains convergence rates comparable to or worse than those of mini-batch SGD.

We focus on two difficulties that arise when analyzing local update methods. First, analyses must account for client drift (Karimireddy et al., 2019). As clients perform local updates on heterogeneous datasets, their local models drift apart. This hinders convergence to globally optimal models, and makes theoretical analyses more challenging. Similar phenomena were examined by Li et al. (2020a); Malinovsky et al. (2020); Pathak and Wainwright (2020) and Fallah et al. (2020b), who show that various local update methods do not converge to critical points of the empirical loss.

Second, local update methods are difficult to compare. Analyses of different methods may use different hyperparameters regimes, or make different assumptions. Even comparing seemingly similar methods can require significant theoretical insight (Karimireddy et al., 2019; Fallah et al., 2020b). Moreover, comparisons can be made in fundamentally different ways. One may wish to maximize the final accuracy, or minimize the number of communication rounds needed to attain a given accuracy. Thus, it is not even clear how local update methods should be compared.

In this work, we invert the conventional narrative that issues such as client drift harm convergence. Instead, we view such phenomena as improving convergence, but to sub-optimal points.

More generally, we show that local update methods face a fundamental trade-off between convergence and accuracy that is explicitly governed by algorithmic hyperparameters. Perceived failures of methods such as FedAvg actually correspond to operating points prioritizing convergence over accuracy. We use this trade-off to develop a novel framework for comparing local update methods. We compare methods based on their entire convergence-accuracy trade-off, not just their convergence to optimal points. In more detail:

We show that for quadratic models, local update methods are equivalent to optimizing a single surrogate loss function. The condition number of the surrogate is controlled by algorithmic choices. Popular local update methods, including FedAvg and MAML, reduce the surrogate’s condition number, but increase the discrepancy between the empirical and surrogate losses. Our results also encompass proximal local update methods (Li et al., 2020a; Zhou et al., 2019).

We derive novel convergence rates that showcasing this trade-off between convergence and accuracy. Our bounds demonstrate the benefit of local update methods over methods such as mini-batch SGD in communication-limited settings.

We use this theory to develop a framework for comparing local update methods through a novel Pareto frontier, which compares convergence-accuracy trade-offs of classes of algorithms. We use this to derive novel comparisons of many popular local update methods.

We use this technique to shed light on a broad range of phenomena, including the benefit of server momentum, the effect of proximal local updates, and differences between the dynamics of FedAvg and MAML.

While our theoretical results are restricted to quadratic models, we show that such convergence-accuracy trade-offs occur empirically in non-convex settings. We also validate our theoretical observations regarding server momentum and proximal updates on a non-convex task.

We view our work as a step towards holistic understandings of local update methods. Using the aforementioned Pareto frontiers, we highlight a number of new phenomena and open problems. One particularly intriguing observation is that the convergence-accuracy trade-off for FedAvg with heavy-ball server momentum appears to be completely symmetric. For more details, see Section 5. Our proof techniques may be of independent interest. We derive a novel analog of the Bhatia-Davis inequality (Bhatia and Davis, 2000) for mean absolute deviations, and use this to understand the accuracy of local update methods.

Notation

Accuracy and Meta-Learning

We study the accuracy of local update methods on the training population. However, meta-learning algorithms are designed to learn a model that adapts well to new tasks; The empirical loss is not necessarily indicative of the “post-adaptation” accuracy of such methods (Finn et al., 2017). Despite this our focus still yields novel insights into qualitative differences between the training dynamics of federated learning and meta-learning methods. Perhaps surprisingly, we show that in certain hyperparameter regions, these methods exhibit identical trade-offs between convergence and pre-adaptation accuracy (see Figures 4 and 5). While we believe our results can be adapted to post-adaptation accuracy via techniques developed by Fallah et al. (2020a), we leave the analysis to future work.

Problem Setup

For i∈Ii\in\mathcal{I}, we define the client loss function fif_{i} and the overall loss function ff as follows:

The joint distribution (I,Z)(\mathcal{I},\mathcal{Z}) defines a distribution over Z\mathcal{Z}, recovering standard risk minimization, as well as distributed risk minimization in which P\mathcal{P} and all Di\mathcal{D}_{i} are uniform over finite sets. For i∈Ii\in\mathcal{I}, define:

We assume these expectations exist and are finite. One can show that up to some additive constant,

We make the following assumptions throughout.

There are μ,L>0\mu,L>0 such that for all ii, μI⪯Ai⪯LI\mu I\preceq A_{i}\preceq LI.

There is some C>0C>0 such that for all ii, ∥ci∥≤C\|c_{i}\|\leq C.

After receiving all client updates, the server treats their average qtq_{t} as an estimate of the gradient of the loss function ff, and applies qtq_{t} to a first-order optimization algorithm ServerOpt. For example, the server could perform a gradient descent step using the “pseudo-gradient” qtq_{t}. We refer to this process (parameterized by α,γ,Θ\alpha,\gamma,\Theta and ServerOpt) as LocalUpdate and give pseudo-code in Algorithms 1 and 2.

LocalUpdate recovers many well-known algorithms for various choices Θ\Theta. For convenience, define

Special cases of LocalUpdate when ServerOpt is gradient descent are given in Table 1. For details on the relation between FedAvg and LocalUpdate, see Appendix A. By changing ServerOpt, we can recover methods such as FedAvgM (Hsu et al., 2019) (server gradient descent with momentum), and FedAdam (Reddi et al., 2020) (server Adam (Kingma and Ba, 2014)).

Local Update Methods as First-Order Methods

LocalUpdate can vary drastically from first-order optimization methods on the empirical loss. Despite this, we will show that Algorithm 2 is equivalent in expectation to ServerOpt applied to a single surrogate loss. This surrogate loss is determined by the inputs α,γ\alpha,\gamma and Θ\Theta to Algorithm 2. For each client i∈Ii\in\mathcal{I}, we define its distortion matrix Qi(α,γ,Θ)Q_{i}(\alpha,\gamma,\Theta) as

We define the surrogate loss function of client ii as

and the overall surrogate loss function as

When Θ=Θ1\Theta=\Theta_{1}, Qi(α,γ,Θ)=IQ_{i}(\alpha,\gamma,\Theta)=I, in which case there is no distortion. In general, Qi(α,γ,Θ)Q_{i}(\alpha,\gamma,\Theta) can amplify the heterogeneity of the AiA_{i}. We derive the following theorem linking the surrogate losses to Algorithm 2.

A version of Theorem 1 was shown for α=0,Θ=Θ2\alpha=0,\Theta=\Theta_{2} by Fallah et al. (2020b). We take this a step further and show that in certain settings, MAML is equivalent in expectation to ServerOpt on a surrogate loss.

MAML with KK local steps can be viewed as a modification of LocalUpdate. Algorithm 1 remains the same, and in Algorithm 2, each client executes KK mini-batch SGD steps. However, the client’s message to the server is different. Let XKi(x)X_{K}^{i}(x) be the function that runs KK steps of mini-batch SGD, starting from xx, for fixed mini-batches S1,…,SKS_{1},\dots,S_{K} of size BB drawn independently from Di\mathcal{D}_{i}. Define

Each client ii sends a stochastic estimate of ∇mKi(x)\nabla m_{K}^{i}(x) to the server. The rest is identical to LocalUpdate; The server averages the client outputs and uses this as a gradient estimate for ServerOpt. While MAML is not a special case of LocalUpdate, we show that if the clients use gradient descent, MAML is equivalent in expectation to LocalUpdate with Θ=Θ2K+1\Theta=\Theta_{2K+1}.

If XKi(x)X_{K}^{i}(x) is the function that runs KK steps of gradient descent on Di\mathcal{D}_{i} with learning rate γ\gamma starting at xx, then

Convergence and Accuracy of Local Update Methods

We wish to better understand (9) in cases of interest. We first consider Θ=Θ1:K\Theta=\Theta_{1:K}, as in FedAvg. Define

When γ=0,α=0\gamma=0,\alpha=0, we recover the condition number L/μL/\mu of the empirical loss ff. We next consider Θ=ΘK\Theta=\Theta_{K}, as in MAML-style algorithms. Define

As K→1K\to 1 or γ→0\gamma\to 0, κ(α,γ,ΘK)→L/μ\kappa(\alpha,\gamma,\Theta_{K})\to L/\mu, which bounds the condition number of the empirical loss ff. If γ\gamma is not close to 0, we get an exponential reduction (in terms of KK) of the condition number. While the analysis is not as clear for Θ1:K\Theta_{1:K}, one can show that κ(α,γ,Θ1:K)≤L/μ\kappa(\alpha,\gamma,\Theta_{1:K})\leq L/\mu, with equality if and only if α=0\alpha=0, and either γ=0\gamma=0 or K=1K=1. Moreover, κ(α,γ,Θ1:K)\kappa(\alpha,\gamma,\Theta_{1:K}) decreases as K→∞K\to\infty or γ→(L+α)−1\gamma\to(L+\alpha)^{-1}. For both Θ1:K\Theta_{1:K} and ΘK\Theta_{K}, increasing α\alpha decreases κ\kappa.

Here we see the impact of local update methods on convergence: Popular methods such as FedAvg, FedProx, MAML, and Reptile reduce the condition number of the surrogate loss function they are actually optimizing. In the next section, we translate this into concrete convergence rates for LocalUpdate.

2 Convergence Rates

Suppose γ<(L+α)−1\gamma<(L+\alpha)^{-1} and ServerOpt is gradient descent with Nesterov, heavy-ball, or no momentum. Then for some hyperparameter setting of ServerOpt, and ρ\rho as in Table 2, the iterates {xt}t≥1\{x_{t}\}_{t\geq 1} of LocalUpdate satisfy

Thus, (properly tuned) server momentum improves the convergence of LocalUpdate, giving theoretical groundingYuan and Ma (2020) first showed that momentum can accelerate FedAvg, though they use a different momentum scheme with extra per-round communication. to the improved convergence of FedAvgM shown by Hsu et al. (2019) and Reddi et al. (2020). Since ServerOpt does not change the surrogate loss, this improvement in convergence does not degrade the accuracy of the learned model.

3 Distance Between Global Minimizers

We are interested in ∥x∗(α,γ,Θ)−x∗∥\|x^{*}(\alpha,\gamma,\Theta)-x^{*}\|. While we focus on the setting where P\mathcal{P} is a discrete distribution over some finite I\mathcal{I}, our analysis can be generalized to arbitrary probability spaces (I,F,P)(\mathcal{I},\mathcal{F},\mathcal{P}). We derive the following bound.

Let b=max⁡i∈Iλmax⁡(Qi(α,γ,Θ))b=\max_{i\in\mathcal{I}}\lambda_{\max}(Q_{i}(\alpha,\gamma,\Theta)) and a=min⁡i∈Iλmin⁡(Qi(α,γ,Θ))a=\min_{i\in\mathcal{I}}\lambda_{\min}(Q_{i}(\alpha,\gamma,\Theta)). Then

When d=1d=1, we can reduce the constant factor to 2C2C, which we show is tight (see Appendix C.2). While we conjecture that this bound holds with a constant of 2C2C for all dd, we leave this to future work.

Our proof technique for Lemma 5 may be of independent interest. We derive this result by first proving an analog of the Bhatia-Davis inequality (Bhatia and Davis, 2000) for mean absolute deviations of bounded random variables (Theorem 5 in Appendix C).

Let κ0:=L/μ\kappa_{0}:=L/\mu. Specializing to Θ=Θ1:K\Theta=\Theta_{1:K} or ΘK\Theta_{K}, we derive a link between κ(α,γ,Θ)\kappa(\alpha,\gamma,\Theta) in (11) and (13) and the distance between optimizers.

Suppose that either (I) γ<(L+α)−1\gamma<(L+\alpha)^{-1} and Θ=Θ1:K\Theta=\Theta_{1:K} or (II) γ<(KL+α)−1\gamma<(KL+\alpha)^{-1} and Θ=ΘK\Theta=\Theta_{K}. Then for all i∈Pi\in\mathcal{P}, cond⁡(Qi(α,γ,Θ))≤κ0κ(α,γ,Θ)−1\operatorname*{cond}(Q_{i}(\alpha,\gamma,\Theta))\leq\kappa_{0}\kappa(\alpha,\gamma,\Theta)^{-1}.

Applying Theorem 3, we bound the convergence of LocalUpdate to the empirical minimizer x∗x^{*}.

Under the same settings as Theorem 3, for some hyperparameter setting of ServerOpt, the iterates {xt}t≥1\{x_{t}\}_{t\geq 1} of LocalUpdate satisfy

where κ=κ(α,γ,Θ)\kappa=\kappa(\alpha,\gamma,\Theta) is given in (11) and (13), and ρ\rho is given in Table 2.

Here we see the benefit of local update methods in communication-limited settings. When TT is small and ∥x0−x∗∥\left\lVert x_{0}-x^{*}\right\rVert is large, we can achieve better convergence by decreasing ρ\rho and leaving the second term fixed. In such settings, FedAvg can arrive at a neighborhood of a critical point in fewer communication rounds than mini-batch SGD, but may not ever actually reach the critical point. If ∥x0−x∗∥\left\lVert x_{0}-x^{*}\right\rVert is small, we may be better served by using mini-batch SGD instead.

Comparing Local Update Methods

Comparing optimization algorithms is a fundamental theoretical effort. Many past works compare local update methods based on their convergence to critical points of the empirical loss. By Theorem 1, LocalUpdate is only guaranteed to converge to critical points of ff if γ=0\gamma=0 or K(Θ)=1K(\Theta)=1. Thus, existing analyses ignore many useful cases of LocalUpdate.

To remedy this, we compare local update algorithms on the basis of both convergence and accuracy. Instead of fixing γ\gamma and Θ\Theta, we analyze LocalUpdate as γ\gamma and K(Θ)K(\Theta) vary. To do so, we use our theory from Section 4. Given α,γ\alpha,\gamma and Θ\Theta, we define the convergence rate ρ(α,γ,Θ)\rho(\alpha,\gamma,\Theta) as the infimum over all ρ\rho such that for all T≥1T\geq 1, (14) holds. Values of ρ\rho when ServerOpt is gradient descent are given in Table 2. For Θ=Θ1:K\Theta=\Theta_{1:K} or ΘK\Theta_{K}, we define the suboptimality Δ(α,γ,Θ)\Delta(\alpha,\gamma,\Theta) by

By Theorem 4, this captures the asymptotic worst-case suboptimality of LocalUpdate.

Note that ρ,Δ∈\rho,\Delta\in. Therefore, fixing μ,L\mu,L and ServerOpt, we obtain a Pareto frontier in 2^{2} by plotting (ρ,Δ)(\rho,\Delta) for various γ\gamma and K(Θ)K(\Theta). This curve represents the worst-case convergence/accuracy trade-off of a class of local update methods. We generally want the curve to be as close to (0,0)(0,0) as possible.

For example, in Figure 1 we let ServerOpt be gradient descent and set α=0,Θ=Θ1:K\alpha=0,\Theta=\Theta_{1:K}. We plot (ρ,Δ)(\rho,\Delta) as we vary KK and fix γ\gamma, and vice-versa. When L=5L=5, we obtain nearly identical curves. The curves for L=50L=50 are similar, except that when we fix KK and vary γ\gamma, we do not reach ρ≈0\rho\approx 0. While γ\gamma and KK have similar impacts on convergence-accuracy trade-offs, varying KK leads a larger set of attainable (ρ,Δ)(\rho,\Delta). Formally, this is because in (11), lim⁡γ→L−1(κ)≠0\lim_{\gamma\to L^{-1}}(\kappa)\neq 0. Intuitively, K→∞K\to\infty recovers one-shot averaging while γ→L−1\gamma\to L^{-1} does not. Notably, the convergence-accuracy trade-off becomes closer to a linear trade-off as L/μL/\mu decreases.

The Pareto frontiers contain more information than just the convergence rate to a critical point (the curve’s intersection with the xx-axis). This information is useful in communication-limited regimes, where we wish to minimize the number of rounds needed to attain a given accuracy. The curves also help visualize various hyperparameter settings of an algorithms simultaneously. To illustrate this, we use the Pareto frontiers to derive novel findings regarding server momentum, proximal client updates, and qualitative differences between FedAvg and MAML. The results are all given below. For more results, see Appendix D.

As shown empirically by Hsu et al. (2019) and as reflected in Table 2, server momentum can improve convergence. To understand this, in Figure 2 we compare Pareto frontiers where Θ=Θ1:K\Theta=\Theta_{1:K} and ServerOpt is gradient descent with various types of momentum (Nesterov, heavy-ball, or no momentum). We see a strict ordering of the server optimization methods. Heavy-ball momentum is better than Nesterov momentum, which is better than no momentum.

One important finding is that the benefit of momentum is more pronounced as L/μL/\mu increases. On the other hand, the benefit of server momentum diminishes for sufficiently large KK: In Figure 4, the various types of momentum lead to similar suboptimality when the convergence rate is close to 0. Intuitively, as K→∞K\to\infty, we recover one-shot averaging, which converges in a single communication round with or without momentum.

Another intriguing observation: The Pareto frontiers for heavy-ball momentum appear to be symmetric about the line ρ=Δ\rho=\Delta. We conjecture this is true for any μ,L\mu,L. While we believe that this may be provable by careful algebraic manipulation of our results above, ideally a proof would explain the root causes of this symmetry. Thus, we leave a proof to future work.

Proximal Client Updates

So far we have only considered α=0\alpha=0. One might posit that as α\alpha varies, the Pareto frontier moves closer to the origin. This appears to not be the case. In all settings we examined, changing α\alpha did not bring the Pareto frontier closer to . Instead, the frontier for α>0\alpha>0 was simply a subset of the frontier for α=0\alpha=0.

To illustrate this, we plot Pareto frontiers for varying α\alpha in Figure 3. As α\alpha increases, the frontier becomes a smaller subset of the frontier for α=0\alpha=0. Thus, proximal client updates may not enable faster convergence. Rather, their benefit may be in guarding against setting γ\gamma too small or KK too large. Figure 3 shows that FedAvg can always attain the same (ρ,Δ)(\rho,\Delta) as FedProx, but it may require different hyperparameters. The reverse is not true, as FedProx cannot recover one-shot averaging. Our findings are consistent with work by Wang et al. (2020), who show that FedProx can reduce the “objective inconsistency” of FedAvg, at the expense of increasing convergence time.

Comparing MAML to FedAvg

We now turn our attention to comparing FedAvg-style algorithms (Θ=Θ1:K\Theta=\Theta_{1:K}) to MAML-style algorithms (Θ=ΘK)\Theta=\Theta_{K}). We plot Pareto frontiers for the ρ,Δ\rho,\Delta guaranteed by Theorems 3 and 4. The results are in Figure 4.

For each ServerOpt, the MAML frontier is a subset of the FedAvg frontier. Recall that in Theorem 3, we require γ<(L+α)−1\gamma<(L+\alpha)^{-1} for FedAvg, but γ<(KL+α)−1\gamma<(KL+\alpha)^{-1} for MAML. In Figure 4 this causes the frontier for ΘK\Theta_{K} to be more restrictive than for Θ1:K\Theta_{1:K}. However, it is still notable that these two fundamentally different methods, attain the same frontier when ρ\rho is large.

The Pareto frontiers are identical for small KK (mirroring Figure 4), but diverge when γ≥(KL+α)−1\gamma\geq(KL+\alpha)^{-1}. The frontier for MAML then moves further from 0. Intuitively, FedAvg tries to learn a global model, while MAML tries to learn a model that adapts quickly to new tasks (Finn et al., 2017); MAML need not minimize (ρ,Δ)(\rho,\Delta). The MAML frontier is noisy for d=5d=5 (as ρ\rho, Δ\Delta depend on random eigenvalues of AA), but stabilizes for d=100d=100. While we posit that this reflects a semi-circle law for eigenvalues of random matrices (Alon et al., 2002), we leave an analysis to future work.

One final observation that highlights the similarities and differences of FedAvg- and MAML-style methods: In Figure 4, the curve for MAML when d=100d=100 has a clear cusp. This seems to occur at the same suboptimality (ie. yy-value) as the intersection of the FedAvg curve with the yy-axis. In other words, the behavior of MAML diverges substantially from FedAvg, but only after it reaches the same suboptimality as FedAvg for K→∞K\to\infty (which corresponds to one-shot averaging). We are unsure why the suboptimality of one-shot averaging corresponds to a cuspidal operating point of MAML, but this observation highlights significant nuance in the behavior of these methods.

Limitations and Discussion

Our convergence-accuracy framework and the resulting Pareto frontiers can be useful tools in understanding how algorithmic choices impact local update methods. The obvious limitation is that they only apply to quadratic models. While this is restrictive, we show empirically in Appendix F that even for non-convex functions, the client learning rate governs a convergence-accuracy trade-off for FedAvg.

Our framework may also be useful in identifying important phenomena underlying LocalUpdate, even in non-quadratic settings. To demonstrate this, we show that many of the observations in Section 5 hold in non-convex settings. We train a CNN on the FEMNIST dataset (Caldas et al., 2018) using LocalUpdate where Θ=Θ1:50\Theta=\Theta_{1:50}. We tune client and server learning rates. See Appendix E for full details. In Figure 6, we illustrate how server momentum and α\alpha change convergence. Our results match the Pareto frontiers in Figures 2 and 3: Server momentum improves convergence, while α\alpha has little to no effect, provided we tune learning rates.

This brief example illustrates that our framework can identify crucial facets of local update methods. While our framework may not capture all relevant details of such methods, we believe it greatly simplifies their analysis, comparison, and design. In the future, we hope to extend this framework to more general loss functions. Other important extensions include stochastic settings with partial client participation, as well as trade-offs between convergence and post-adaptation accuracy of local update methods.

Appendix A Relations between FedAvg, FedProx, and LocalUpdate

We focus on the following (simplified) version of FedAvg, otherwise known as Local SGD (Zinkevich et al., 2010; Stich, 2019): At each iteration tt, we sample some set of clients ItI_{t} of size MM from the client population I\mathcal{I}. Each client i∈Iti\in I_{t} receives the server’s model xtx_{t}, and applies KK steps of mini-batch SGD to its local model, resulting in an updated local model xtix_{t}^{i}. The server receives these models from the sampled clients, and updates its model via

Fix tt, and let gkig_{k}^{i} denote the kk-th mini-batch gradient of client ii. Suppose we use a learning rate of γ\gamma on each client when performing mini-batch SGD. Then we have

A similar analysis holds for FedProx, but with the usage of a proximal term with parameter α>0\alpha>0. In both cases, this is exactly LocalUpdate (see Algorithms 1 and 2) with Θ=Θ1:K\Theta=\Theta_{1:K}, η=γ\eta=\gamma, and where ServerOpt is gradient descent. However, by instead using a server learning rate of η\eta that is allowed to vary independently of γ\gamma in LocalUpdate, we can obtain markedly different convergence behavior. We note that a form of this decoupling has previously been explored by Karimireddy et al. (2019) and Reddi et al. (2020). However, these versions instead perform averaging on the so-called “model delta”, in which the server model is updated via

Thus, while this does decouple η\eta and γ\gamma to some degree, it does not fully do so. In particular, if we set γ=0\gamma=0, then xt+1=xtx_{t+1}=x_{t}, in which case we can make no progress overall. This is particularly important because, as implied by Theorem 1, for K(Θ)>1K(\Theta)>1 we can only guarantee that the surrogate loss has the same critical points as the true loss by setting γ=0\gamma=0. More generally, we see that the effective learning rate used in the model-delta approach is the product ηγ\eta\gamma. This can result in conflations between the effect of changing the server learning rate γ\gamma and changing the client learning rate η\eta. By disentangling these, we can better understand differences in the impact of these parameters on the underlying optimization dynamics.

Appendix B Omitted Proofs

B.1 Proofs of Theorems 1 and 2

Suppose we perform KK iterations of SGD on hh with learning rate γ\gamma. That is, starting at x1x_{1} we generate a sequence of independent random vectors {gk}k=1K\{g_{k}\}_{k=1}^{K} and corresponding SGD iterates {xk}k=1K+1\{x_{k}\}_{k=1}^{K+1} satisfying, for 1≤k≤K1\leq k\leq K,

We then have the following lemma regarding the gkg_{k}.

By (18), the law of total expectation, and the independence of the gkg_{k},

Recall that for i∈Ii\in\mathcal{I}, α≥0,γ≥0\alpha\geq 0,\gamma\geq 0 and Θ=(θ1,…,θK(Θ))\Theta=(\theta_{1},\dots,\theta_{K(\Theta)}) we define the matrix Qi(α,γ,Θ)Q_{i}(\alpha,\gamma,\Theta) by

We can now prove Theorem 1. For convenience, we restate the theorem here.

Fix i∈Ii\in\mathcal{I}, x1x_{1}, and α≥0\alpha\geq 0. For z∈Zz\in\mathcal{Z}, define

Note that Aα=Ai+αIA_{\alpha}=A_{i}+\alpha I, where AiA_{i} is as in (3). Straightforward manipulation of (24) implies

Note that the stochastic gradients g1,…,gKg_{1},\ldots,g_{K} computed in Algorithm 2 are therefore independent stochastic gradients of hαh_{\alpha}. By applying Lemma 7 and noting that the constant term τα\tau_{\alpha} does not impact these stochastic gradients, we have that for k≥1k\geq 1,

Expanding and using the fact that in Algorithm 2, x1=xx_{1}=x, we have

This last step follows from (3). Taking a sum and using the linearity of expectation,

A similar analysis using Lemma 7 can be used to derive Theorem 2, which we also restate.

Fix i∈Ii\in\mathcal{I}, and for convenience of notation, let Xk:=Xki(x),X1=xX_{k}:=X_{k}^{i}(x),X_{1}=x. Thus, for k≥1k\geq 1 we have

Since ∇2fi(y)=Ai\nabla^{2}f_{i}(y)=A_{i} for all yy, we have

The result follows from applying (23) for Θ=Θ2K+1\Theta=\Theta_{2K+1}. ∎

B.2 Proofs of Lemmas 1, 2, 3, 4, and 6

These lemmas will follow from a spectral analysis of Qi(α,γ,Θ)Q_{i}(\alpha,\gamma,\Theta). We defer the proof of Lemma 5 to Appendix C due to its more elaborate nature. We first state a general result about eigenvalues of expected values of matrices.

We also compute the spectrum of Qi(α,γ,Θ)Q_{i}(\alpha,\gamma,\Theta) and QI(α,γ,Θ)AiQ_{I}(\alpha,\gamma,\Theta)A_{i}.

For each eigenvalue λ\lambda of AiA_{i}, Qi(α,γ,Θ)Q_{i}(\alpha,\gamma,\Theta) has an eigenvalue

and Qi(α,γ,Θ)AiQ_{i}(\alpha,\gamma,\Theta)A_{i} has an eigenvalue

both with the same multiplicity as λ\lambda.

Let vv be an eigenvector of AiA_{i} with eigenvalue λ\lambda. Then vv is an eigenvector of (I−γ(Ai+αI))k−1(I-\gamma(A_{i}+\alpha I))^{k-1} with eigenvalue (1−γ(λ+α))k−1(1-\gamma(\lambda+\alpha))^{k-1}, implying (29) is an eigenvalue of Qi(α,γ,Θ)Q_{i}(\alpha,\gamma,\Theta) with eigenvector vv. Similarly, we note that vv is an eigenvector of (I−γ(Ai+αI))k−1Ai(I-\gamma(A_{i}+\alpha I))^{k-1}A_{i} with eigenvalue (1−γ(λ+α))k−1λ(1-\gamma(\lambda+\alpha))^{k-1}\lambda, implying (30) is an eigenvalue of Qi(α,γ,Θ)AiQ_{i}(\alpha,\gamma,\Theta)A_{i} with eigenvector vv. The statement about multiplicities follows directly. ∎

Lemma 9 can be used in a straightforward manner to prove Lemma 1, which we restate and prove below.

By Assumption 1, we have that for every eigenvalue λ\lambda of AiA_{i}, λ≤L\lambda\leq L. Therefore, for any such λ\lambda,

Since the θk\theta_{k} are nonnegative and not all zero by assumption, we see that (29) is a sum of nonnegative terms, at least one of which must be positive. Therefore, all eigenvalues of Qi(α,γ,Θ)Q_{i}(\alpha,\gamma,\Theta) are positive, and Qi(α,γ,Θ)Q_{i}(\alpha,\gamma,\Theta) is therefore symmetric and positive definite.

By Assumption 1, AiA_{i} is also symmetric and positive definite, hence the product Qi(α,γ,Θ)AiQ_{i}(\alpha,\gamma,\Theta)A_{i} is symmetric positive definite. However, since

Applying Lemma 8, we conclude the proof. ∎

By Lemma 10, it suffices to derive upper and lower bounds on the eigenvalues of Qi(α,γ,Θ1:K)AiQ_{i}(\alpha,\gamma,\Theta_{1:K})A_{i}. Fix i∈Ii\in\mathcal{I}. By Lemma 9, we see that the eigenvalues of Qi(α,γ,Θ1:K)AiQ_{i}(\alpha,\gamma,\Theta_{1:K})A_{i} are exactly of the form ϕ(λ,α,γ,K)\phi(\lambda,\alpha,\gamma,K), where λ\lambda is an eigenvalue of AiA_{i}. By Assumption 1, each such λ\lambda satisfies λ∈[μ,L]\lambda\in[\mu,L] where μ>0\mu>0.

Fix α,γ,K\alpha,\gamma,K, and define g(λ):=ϕ(λ,α,γ,K)g(\lambda):=\phi(\lambda,\alpha,\gamma,K). Since γ(L+α)<1\gamma(L+\alpha)<1, basic properties of geometric sums imply that for λ∈[μ,L]\lambda\in[\mu,L],

For a given λ\lambda, let ξ=1−γ(λ+α)\xi=1-\gamma(\lambda+\alpha). Simple but tedious computations show that if we take a derivative with respect to λ\lambda, we have

Note that since γ<(L+α)−1\gamma<(L+\alpha)^{-1} by assumption, 0≤ξ≤10\leq\xi\leq 1 for λ∈[μ,L]\lambda\in[\mu,L]. Therefore, g′(λ)≥0g^{\prime}(\lambda)\geq 0 for λ∈[μ,L]\lambda\in[\mu,L], so any eigenvalue χ\chi of Qi(α,γ,Θ1:K)AiQ_{i}(\alpha,\gamma,\Theta_{1:K})A_{i} must satisfy

We use a similar proof as that of Lemma 3 to prove Lemma 4, which we restate and prove below.

By Lemma 10, it suffices to bound the eigenvalues of Qi(α,γ,ΘK)AiQ_{i}(\alpha,\gamma,\Theta_{K})A_{i}. Fix i∈Ii\in\mathcal{I}. By Lemma 9, we see that the eigenvalues of Qi(α,γ,ΘK)Q_{i}(\alpha,\gamma,\Theta_{K}) are of the form ψ(λ,α,γ,K)\psi(\lambda,\alpha,\gamma,K) where λ\lambda is an eigenvalue of AiA_{i}. Note that by Assumption 1, any such λ\lambda satisfies λ∈[μ,L]\lambda\in[\mu,L].

Fix α,γ,K\alpha,\gamma,K, and define h(λ):=ψ(λ,α,γ,K)h(\lambda):=\psi(\lambda,\alpha,\gamma,K). Let ζ=1−γ(λ+α)\zeta=1-\gamma(\lambda+\alpha). Straightforward computations show

Since γ<(KL+α)−1\gamma<(KL+\alpha)^{-1}, we in particular have γ<(L+α)−1\gamma<(L+\alpha)^{-1} so 0≤ζ≤10\leq\zeta\leq 1 for λ∈[μ,L]\lambda\in[\mu,L]. Since γ<(KL+α)−1\gamma<(KL+\alpha)^{-1}, we also have the term 1−γ(Kλ+α)≥01-\gamma(K\lambda+\alpha)\geq 0 for λ∈[μ,L]\lambda\in[\mu,L]. Thus, h′(λ)≥0h^{\prime}(\lambda)\geq 0 for λ∈[μ,L]\lambda\in[\mu,L]. Thus, any eigenvalue χ\chi of Qi(α,γ,ΘK)AiQ_{i}(\alpha,\gamma,\Theta_{K})A_{i} must satisfy

Finally, we are now equipped to prove Lemma 6, which we restate here for posterity.

This will follow almost immediately from Lemmas 3, 4, and 9. First, consider the case Θ=Θ1:K\Theta=\Theta_{1:K}. Then by Lemma 9 and Assumption 1, we see that

Here we used the fact that κ0:=L/μ\kappa_{0}:=L/\mu and Lemma 3. An almost identical proof gives the analogous result for Θ=ΘK\Theta=\Theta_{K}. ∎

B.3 Tightness of Lemmas 3 and 4

In fact, Lemmas 3 and 4 are tight. Fix any α≥0\alpha\geq 0 and γ<(α+L)−1\gamma<(\alpha+L)^{-1}. Let P\mathcal{P} be supported on a single client ii, and let this client’s dataset be supported on a single example zz where

Note that by (3), we then have Ai=Bi,ci=0A_{i}=B_{i},c_{i}=0. In fact, we will show that in this case, the bounds on the condition numbers given in Lemma 3 and 4 are tight. By direct computation,

If γ<(L+α)−1\gamma<(L+\alpha)^{-1} then similar reasoning to the proof of Lemma 3 implies that the condition number satisfies

By analogous reasoning to the proof of Lemma 4, if γ<(KL+α)−1\gamma<(KL+\alpha)^{-1}, we have

Appendix C Proof of Lemma 5

In order to prove the results in this section, we will use the following straightforward lemma regarding the structure of x∗(α,γ,Θ)x^{*}(\alpha,\gamma,\Theta).

To prove Lemma 5, we will reduce it to a statement about mean absolute deviations of bounded random variables. We define the mean absolute deviation of a random variable below.

To derive our results, we bound the mean absolute deviation of bounded random variables. Our result is inspired by the bound by Bhatia and Davis (2000) on the variance of bounded random variables.

Moreover, this holds with equality iff XX is supported on {a,b}\{a,b\}.

Suppose XX takes on values x1,…,xnx_{1},\dots,x_{n} with probabilities p1,…pnp_{1},\dots p_{n}. We will first show that there is a random variable YY supported on {a,b}\{a,b\} such that D(X)≤D(Y)D(X)\leq D(Y).

Without loss of generality, suppose x1∈(a,b)x_{1}\in(a,b). Define

First note that s,t∈s,t\in. Simple analysis also shows

By iterating this procedure, (which is guaranteed to terminate after at most nn iterations), we obtain some random variable YY supported on {a,b}\{a,b\} such that D(X)≤D(Y)D(X)\leq D(Y), with equality if and only if XX is already supported on {a,b}\{a,b\}. Suppose YY takes on a,ba,b with probabilities (1−p),p(1-p),p. Straightforward calculation shows

Using the fact that for any real xx, (b−a)2≥4(b−x)(x−a)(b-a)^{2}\geq 4(b-x)(x-a) (with equality iff x=(b+a)/2x=(b+a)/2) we arrive at the following corollary, analogous to Popoviciu’s inequality on variances (Popoviciu, 1935).

If XX is a discrete random variable on [a,b][a,b], then

with equality iff XX takes on the values aa and bb, each with probability 1/21/2.

We can now prove a stronger version of Lemma 5 when d=1d=1. For simplicity, we assume P\mathcal{P} is a discrete distribution on some finite I\mathcal{I}, though the analysis can be generalized to arbitrary distributions.

Let b=max⁡i∈Iλmax⁡(Qi(α,γ,Θ))b=\max_{i\in\mathcal{I}}\lambda_{\max}(Q_{i}(\alpha,\gamma,\Theta)), a=min⁡i∈Iλmin⁡(Qi(α,γ,Θ))a=\min_{i\in\mathcal{I}}\lambda_{\min}(Q_{i}(\alpha,\gamma,\Theta)). Then for d=1d=1,

Since each Qi∈[a,b]Q_{i}\in[a,b], we also have v∈[a,b]v\in[a,b].By Lemma 11,

Maximizing the right-hand side for v∈[a,b]v\in[a,b], we get

C.2 Tightness of Lemma 12

Let A1=4,A2=1,c1=1,c2=−1A_{1}=4,A_{2}=1,c_{1}=1,c_{2}=-1. Let α=0,γ=1/8\alpha=0,\gamma=1/8. Then by (6), we have

Therefore, in this setting, b=(7/8)K,a=2−Kb=(7/8)^{K},a=2^{-K}. By applying L’Hopital’s rule, we find

Let P\mathcal{P} be the distribution that selects i=1i=1 with probability pp and i=2i=2 with probability 1−p1-p. Then

For any fixed p∈(0,1)p\in(0,1), straightforward but tedious applications of L’Hopital’s rule yields the fact that

By (36) and (37), we see that by selecting KK sufficiently large and pp sufficiently close to 1, we can ensure that

C.3 General Case

The matrix-weighted mean of {X1,…Xn}\{X_{1},\dots X_{n}\} with respect to {Y1,…Yn}\{Y_{1},\dots Y_{n}\} is given by

When d=1d=1, and the Yi>0Y_{i}>0, this gives the standard mean of a discrete random variable XX taking values X1,…XnX_{1},\dots X_{n} with probabilities Y1,…YnY_{1},\dots Y_{n}. When the context is clear, we will simply denote this by f(X∣Y)f(X|Y). We first prove a simple lemma regarding the Loewner ordering and matrix-weighted means.

Suppose that for all ii, aI⪯Xi⪯bIaI\preceq X_{i}\preceq bI. Then aI⪯f(X∣Y)⪯bIaI\preceq f(X|Y)\preceq bI.

Since the Xi,YiX_{i},Y_{i} are commuting positive definite matrices and YY is positive definite, f(X∣Y)f(X|Y) is similar to the matrix

By assumption, Xi⪯bIX_{i}\preceq bI. Therefore,

An analogous argument shows that P⪰aIP\succeq aI. By basic properties of matrix similarity, we therefore find aI⪯f(X∣Y)⪯bIaI\preceq f(X|Y)\preceq bI. ∎

We can use this matrix-weighted mean to define a normalized, matrix-weighted version of the mean absolute deviation.

The normalized matrix-weighted discrepancy of {X1,…Xn}\{X_{1},\dots X_{n}\} with respect to {Y1,…Yn}\{Y_{1},\dots Y_{n}\} is given by

where ∥⋅∥\left\lVert\cdot\right\rVert is the operator norm.

We will prove an analog of Theorem 5 for this normalized matrix-weighted discrepancy.

To prove this, we will require a straightforward lemma regarding eigenvalues of symmetric positive definite matrices.

Let P1,P2P_{1},P_{2} be symmetric positive definite matrices and let P=P1+P2P=P_{1}+P_{2}. Then

By basic properties of the Loewner ordering,

Since P1P−1P_{1}P^{-1} is similar to P−1/2P1P−1/2P^{-1/2}P_{1}P^{-1/2}, we have

We will proceed in a similar manner to the proof of Theorem 5. We will first show that we can always find a set of symmetric positive definite matrices X1′,…,Xm′,Y1′,…,Ym′X_{1}^{\prime},\dots,X_{m}^{\prime},Y_{1}^{\prime},\dots,Y_{m}^{\prime} such that:

For all ii, aI⪯Xi′⪯bIaI\preceq X_{i}^{\prime}\preceq bI.

For all ii, Xi′,Yi′X_{i}^{\prime},Y_{i}^{\prime} commute.

For t=∣{X1,…,Xn}\{aI,bI}∣t=|\{X_{1},\dots,X_{n}\}\backslash\{aI,bI\}|, we have ∣{X1′,…,Xm′}\{aI,bI}∣≤max⁡{0,t−1}|\{X_{1}^{\prime},\dots,X_{m}^{\prime}\}\backslash\{aI,bI\}|\leq\max\{0,t-1\}.

By iterating this procedure, we can replace X1,…,Xn,Y1,…,YnX_{1},\dots,X_{n},Y_{1},\dots,Y_{n} with matrices X1′′,…,Xl′′,Y1′′,…,Yl′′X^{\prime\prime}_{1},\dots,X^{\prime\prime}_{l},Y^{\prime\prime}_{1},\dots,Y^{\prime\prime}_{l} where the Xi′′X^{\prime\prime}_{i} are all in the set {aI,bI}\{aI,bI\}. It will then suffice to show that M(X′′∣Y′′)M(X^{\prime\prime}|Y^{\prime\prime}) satisfies the desired bound, which we do by a somewhat direct computation, though one that is made much easier due to the fact that aI,bIaI,bI are diagonal.

We now proceed in detail. Define the matrices

Note that since X1,Y1X_{1},Y_{1} commute, SS and TT are products of symmetric, positive definite, commuting matrices. They are therefore symmetric, positive definite, commuting matrices as well. One can easily verify that

We will use ZZ to denote the matrices in (41). Note that by Lemma 13, we know that aI⪯Z⪯bIaI\preceq Z\preceq bI. We will show that this replacement of (X1,Y1)(X_{1},Y_{1}) by (aI,S)(aI,S) and (bI,T)(bI,T) does not decrease the mean absolute deviation. By (40) and (41),

Note that since aI⪯Z⪯bI,aI⪯X1⪯bIaI\preceq Z\preceq bI,aI\preceq X_{1}\preceq bI and the YiY_{i} are positive definite, T1T_{1} and T2T_{2} are positive semi-definite matrices. Simple algebraic manipulation implies that T3=T1−T2T_{3}=T_{1}-T_{2}. Since T1,T2T_{1},T_{2} are positive definite matrices, we have

We therefore exhibit exactly the matrices satisfying properties (1)-(6) described above. By iterating this procedure, we obtain positive definite, symmetric matrices (X1′′,…Xm′′),(Y1′′,…,Ym′′)(X_{1}^{\prime\prime},\dots X_{m}^{\prime\prime}),(Y_{1}^{\prime\prime},\dots,Y_{m}^{\prime\prime}) such that Xi′′,Yi′′X_{i}^{\prime\prime},Y_{i}^{\prime\prime} commute, each Xi′′X_{i}^{\prime\prime} is equal to aIaI or bIbI, and such that

By consolidating Xi′′X_{i}^{\prime\prime} that are equal, we can assume without loss of generality that we have matrices (aI,bI)(aI,bI) with associated symmetric positive definite matrices (C1,C2)(C_{1},C_{2}). Let C=C1+C2C=C_{1}+C_{2}, and let R=aC1+bC2R=aC_{1}+bC_{2}. We then have

Let Z=RC−1Z=RC^{-1}. Then by direct computation,

After some straightforward but tedious algebraic manipulation, we find

Since C1,C2C_{1},C_{2} are symmetric positive definite matrices, we have

Letting p,qp,q denote λmin⁡(C2C1−1),λmin⁡(C1C2−1)\lambda_{\min}(C_{2}C_{1}^{-1}),\lambda_{\min}(C_{1}C_{2}^{-1}) respectively, and noting that we therefore have p,q>0p,q>0, we have

Suppose ∣I∣=n|\mathcal{I}|=n and P\mathcal{P} is the discrete distribution on I\mathcal{I} with associated probabilities pip_{i}. For brevity, we will let QiQ_{i} denote Qi(α,γ,Θ)Q_{i}(\alpha,\gamma,\Theta). We will let Yi=piAiY_{i}=p_{i}A_{i} and Y=∑i=1nYiY=\sum_{i=1}^{n}Y_{i}. By Lemma 11, we have

Here we used Lemma 9, which in particular shows that since γ<(L+α)−1\gamma<(L+\alpha)^{-1}, all the QiQ_{i} are positive definite. Moreover, Lemma 9 shows that Qi,AiQ_{i},A_{i} share the same eigenvectors, and therefore commute with one another. Hence, the QiQ_{i} commute with the YiY_{i}. Moreover, by Assumption 1, the YiY_{i} are positive definite symmetric matrices, and by Assumption 2, ∥ci∥≤C\left\lVert c_{i}\right\rVert\leq C. Applying Theorem 6, we have

The last inequality holds from simple algebraic manipulation.

Appendix D Additional Pareto Frontiers

To generate AA, we generate BB by sampling its entries independently from N(0,1)\mathcal{N}(0,1). We then set A=β1B⊺B+β2IA=\beta_{1}B^{\intercal}B+\beta_{2}I, where β1,β2\beta_{1},\beta_{2} are the unique scalars such that

We plot the resulting Pareto frontiers for varying LL and fixed μ\mu in Figure 7. We see that as LL increases with respect to μ\mu, the discrepancy between the MAML curves and the FedAvg curves grows. In particular, for small L/μL/\mu, we see that the MAML curve recovers most of the FedAvg curve before diverging, while for large L/μL/\mu, the two diverge almost immediately. Again, we see that when d=5d=5, there is some noise in (ρ,Δ)(\rho,\Delta), which seems to approach some limiting behavior for d=100d=100.

We perform a similar experiment, but where we fix L=10L=10 and vary ServerOpt over gradient descent with no momentum, with Nesterov momentum, and with heavy-ball momentum. The results are given in Figure 8. While the differences are not huge, we see that momentum helps convergence in all cases, FedAvg or MAML. Moreover, we see an interesting phenomenon where the type of momentum changes the concavity of the MAML Pareto frontier for d=100d=100. As we add momentum, the region to the right of the MAML curve becomes more convex, becoming more rounded for heavy-ball momentum than for Nesterov momentum.

D.2 Proximal MAML-style Pareto Frontiers

In Figure 9 we plot the analog of the Pareto frontiers in Figure 3, but for MAML-style algorithms where Θ=ΘK\Theta=\Theta_{K}. We see a similar, though more subdued, version of the behavior in Figure 3. That is, adding a proximal term simply alters how much of the Pareto frontier is traversed; it does not change the fundamental shape. Note that here we only used γ\gamma satisfying γ<(KL+α)−1\gamma<(KL+\alpha)^{-1}, as required by Theorem 3. In particular, the only restriction on the shape of the curve seems to be coming from the fact that larger α\alpha reduces the set of γ\gamma satisfying γ<(KL+α)−1\gamma<(KL+\alpha)^{-1}.

To see the effects of α\alpha when γ≥(KL+α)−1\gamma\geq(KL+\alpha)^{-1}, we use the same simulated approach as in Figure 5. We do this for varying α\alpha in Figure 10. We see that while increasing α\alpha shrinks the space of the Pareto curve of FedAvg, it does not seem to change the MAML curves by a meaningful amount.

Appendix E Experimental Setup

We use three datasets: the federated extended MNIST dataset (FEMNIST) (Caldas et al., 2018), CIFAR-100 (Krizhevsky and Hinton, 2009), and Shakespeare (Caldas et al., 2018). The first two are image datasets, and the third is a language dataset. All datasets are publicly available. We specifically use the versions available in TensorFlow Federated (Ingerman and Ostrowski, 2019), which gives a federated structure to all three. We keep the client partitioning when training, and create a test dataset by taking a union over all test client datasets. Statistics on the number of clients and examples in each dataset are given in Table 3.

The FEMNIST dataset consists of images hand-written alphanumeric characters. There are 62 total alphanumeric characters represented in the dataset. The images are partitioned among clients according to their author. The dataset has natural heterogeneity stemming from the writing style of each person. We train a convolutional network on the dataset (the same one used by Reddi et al. (2020)). The network has two convolutional layers. Each convolutional layer uses 3×33\times 3 kernels, max pooling, and then dropout with probability p=0.25p=0.25. The model has a final dense softmax output layer.

CIFAR-100

The CIFAR-100 dataset is a computer vision dataset consisting of 32×32×332\times 32\times 3 images with 100 possible labels. While this dataset does not have a natural partition among clients, a federated version was created by Reddi et al. (2020) using hierarchical latent Dirichlet allocation to enforce moderate amounts of heterogeneity among clients. We train a ResNet-18 on this dataset, where we replace all batch normalization layers with group normalization layers (Wu and He, 2018). The use of group norm over batch norm in federated learning was first advocated by Hsieh et al. (2019).

We perform small amounts of data augmentation and preprocessing, as is standard with CIFAR-100. We first perform a random crop to shape (24,24,3)(24,24,3), followed by a random horizontal flip. We then normalize the pixel values according to their mean and standard deviation. Thus, given an image xx, we compute (x−μ)/σ(x-\mu)/\sigma where μ\mu is the average of the pixel values in xx, and σ\sigma is the standard deviation.

Shakespeare

The Shakespeare dataset is derived from the benchmark designed by Caldas et al. (2018). The dataset corpus is the collected works of William Shakespeare, and the clients correspond to roles in Shakespeare’s plays with at least two lines of dialogue. To eliminate confusion, character here will refer to alphanumeric and other such symbols, while we will use client to denote the various roles in plays. We split each client’s lines into sequences of 80 characters, padding if necessary. We use a vocabulary size of 90: 86 characters contained in Shakespeare’s work, beginning and end of line tokens, padding tokens, and out-of-vocabulary tokens. We perform next-character prediction on the clients’ dialogue using an RNN. The RNN takes as input a sequence of 80 characters, embeds it into a learned 8-dimensional space, and passes the embedding through 2 LSTM layers, each with 256 units. Finally, we use a softmax output layer with 80 units, where we try to predict a sequence of 80 characters formed by shifting the input sequence over by one. Therefore, our output dimension is 80×9080\times 90. We compute loss using cross-entropy loss.

E.2 Implementation and Hyperparameters

We implement LocalUpdate in TensorFlow Federated (Ingerman and Ostrowski, 2019). We use LocalUpdate with Θ=Θ1:K\Theta=\Theta_{1:K} and client learning rate γ\gamma. In all experiments, ServerOpt is gradient descent with server learning rate η\eta, with either no momentum, Nesterov momentum, or heavy-ball momentum. We sample M=10M=10 clients per round. We sample without replacement within a given round, and with replacement across rounds. In order to derive fair comparisons between different hyperparameter settings, we use a random seed to fix which clients are sampled at each round. We use a batch size of B=20B=20 for FEMNIST and CIFAR-100, and B=4B=4 for Shakespeare.

E.3 Details of Figure 6

For posterity’s sake, we re-plot Figure 6 in Figure 11. To generate these plots, we perform two distinct experiments. In the first experiment (Figures 6 and 11, left), we fix α=0\alpha=0 and vary ServerOpt. Specifically, we let ServerOpt be gradient descent with no momentum (gradient), gradient descent with Nesterov momentum (nesterov), and gradient descent with heavy-ball momentum (momentum). When ServerOpt uses Nesterov or heavy-ball momentum, we use a momentum parameter of β=0.9\beta=0.9. In the second experiment (Figures 6 and 11, right), we fix ServerOpt to be gradient descent with no momentum, and vary the proximal strength α\alpha. In both cases, we fix Θ=Θ1:50\Theta=\Theta_{1:50}, and tune γ,η\gamma,\eta over the range

We select the values of γ,η\gamma,\eta attaining the best average test accuracy over the last 100 rounds.

Appendix F Additional Experiments

We wish to showcase the convergence-accuracy trade-off discussed in Section 4 in non-convex settings. We train LocalUpdate with α=0\alpha=0, Θ=Θ1:10\Theta=\Theta_{1:10}, and let ServerOpt be gradient descent with learning rate η\eta. First, we fix η=0.01\eta=0.01 and vary γ\gamma over

We plot the training loss over time on all three datasets in Figure 12, omitting results that diverge due to γ\gamma being too large.

We see that on all three tasks, especially CIFAR-100, the choice of client learning rate can impact not just the speed of convergence, but what point the algorithm converges to. In general, we see very similar behavior to that described in Sections 4 and 5, despite the non-convex loss functions involved in all three tasks. For both FEMNIST and CIFAR-100, smaller client learning rates eventually reach lower training losses than higher learning rates. This is particularly evident in the results for CIFAR-100. While γ=10−2\gamma=10^{-2} initially performs better than all other methods, it is eventually surpassed by γ=10−3\gamma=10^{-3}, and γ=0\gamma=0 ends up obtaining a comparable accuracy. This reflects the idea presented in Section 5 that hyperparameters should be chosen according to the desired convergence-accuracy trade-off. In communication-limited settings, we should use larger γ\gamma (or KK), while in cases where we can run many communication rounds, we should use smaller γ\gamma (or KK).

In short, we see clear evidence that the choice of client learning rate γ\gamma leads to a trade-off between convergence and accuracy. However, as shown in Lemmas 3 and 4, the condition number of the surrogate loss changes depending on parameters such as γ\gamma. To derive asymptotically optimal rates for strongly convex functions (such as the ones in Table 2), one must generally set the learning rate η\eta according to the condition number. Thus, we repeat the experiments in Figure 12, but where we tune the server learning rate η\eta instead of fixing it. This helps account for how the optimization dynamics can change as a function of the client learning rate γ\gamma. We vary the server learning rate η\eta over

and select η\eta that leads to the smallest average training loss over the last 100 rounds. The result is given in Figure 13. Again, we see similar behavior, but see that when the server learning rate is tuned, larger client learning rates may do much better initially. This reflects the fact that in Table 2, the best convergence rates can only be obtained by setting parameters of ServerOpt correctly.

This points to another benefit of the Pareto frontiers proposed in Section 5. Many comparisons of different algorithms, especially empirical ones, can miss good hyperparameter settings. This is heightened by the fact that many FL and ML methods have hyperparameters for both client and server optimizers, comprehensive tuning extremely difficult. This may lead to unfair comparisons between methods. By contrast, the Pareto frontiers showcase convergence-accuracy trade-offs when the hyperparameters of ServerOpt are selected in an “optimal” way, helping derive fair comparisons between methods.