Adaptive Federated Optimization
Authors: Sashank J. Reddi
Zachary Charles
Manzil Zaheer
Zachary Garrett
Keith Rush
Jakub Konečný
Sanjiv Kumar
Hugh Brendan McMahan
Venue: ICLR 2021 Poster (2021)
Keywords: Federated learning
optimization
adaptive optimization
distributed optimization

Abstract: Federated learning is a distributed machine learning paradigm in which a large number of clients coordinate with a central server to learn a model without sharing their own training data. Standard federated optimization methods such as Federated Averaging (FedAvg) are often difficult to tune and exhibit unfavorable convergence behavior. In non-federated settings, adaptive optimization methods have had notable success in combating such issues. In this work, we propose federated versions of adaptive optimizers, including Adagrad, Adam, and  Yogi, and analyze their convergence in the presence of heterogeneous data for general non-convex settings. Our results highlight the interplay between client heterogeneity and communication efficiency. We also perform extensive experiments on these methods and show that the use of adaptive optimizers can significantly improve the performance of federated learning.

Published as a conference paper at ICLR 2021
ADAPTIVE FEDERATED OPTIMIZATION
Sashank J. Reddi; Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, H. Brendan McMahan Google Research
(sashank, zachcharles, manzilzaheer, zachgarrett, krush, konkey, sanjivk, momahan) @google. com
ABSTRACT
Federated learning is a distributed machine learning paradigm in which a large number of clients coordinate with central server to learn model without sharing their own training data. Standard federated optimization methods such as Fed- erated Averaging (FEDAVG) are often difficult to tune and exhibit unfavorable convergence behavior. non-federated settings, adaptive optimization methods have had notable success in combating such issues. In this work, we propose fed- erated versions of adaptive optimizers, including ADAGRAD, ADAM, and YOGI, and analyze their convergence in the presence of heterogeneous data for general nonconvex settings. Our results highlight the interplay between elient heterogeneity and communication efficiency. We also perform extensive experiments on these methods and show that the of adaptive optimizers can significantly improve the performance of federated learning.
INTRODUCTION
Federated learning (FL) is a machine learning paradigm in which multiple clients cooperate to learn model under the orchestration of central server (McMahan et al., 2017). In FL. raw client data is never shared with the server or other clients. This distinguishes FL from traditional distributed optimization, and requires contending with heterogeneous data. FL has two primary settings, cross- silo (eg. FL between large institutions) and cross-device (eg. FL across edge devices) (Kairouz et al., 2019, Table 1). In cross-silo FL, most clients participate in every round and can maintain state between rounds. In the more challenging cross-device FL, our primary focus, only a small fraction of clients participate in each round, and clients cannot maintain state across rounds. For a more in ie-depth discussion of FL and the challenges involved, we defer Kairouz et al. (2019) and Li al. (2019a).
Standard optimization methods, such as distributed SGD, are often unsuitable in FL and can incur high communication costs. To remedy this, many federated optimization methods use local client updates, in which clients update their models multiple times before communicating with the server. This can greatly reduce the amount of communication required to train model. One such method is FEDAvG (McMahan et al., 2017), in which clients perform multiple epochs of SGD on their local datasets. The clients communicate their models to the server, which averages them to form new global model. While FEDAVG has seen great success, recent works have highlighted its convergence issues in some settings (Karimireddy et al., 2019; Hsu et al., 2019). This is due variety of factors including (1) client drift (Karimireddy et al., 2019), where local client models move away from globally optimal models, and (2) lack of adaptivity. FEDAVG is similar in spirit to SGD, and may be unsuitable for settings with heavy-tail stochastic gradient noise distributions, which often arise when training language models (Zhang et al., 2019a). Such settings benefit from adaptive learning rates, which incorporate knowledge of past iterations perform more informed optimization.
In this paper, we focus on the second issue and present a simple framework for incorporating adaptivity in Fl.. In particular, we propose a general optimization framework in which (1) clients perform multiple epochs of training using a client optimizer to minimize loss on their local data and (2) server updates its global model by applying a gradient-based server optimizer to the average of the clients' model updates. We show that FEDAVG is the special case where SGD used as both client and server optimizer and server learning rate is 1. This framework can also seamlessly incorporate
Authors contrihuted equally to this work

---

Published as a conference paper at ICLR 2021
adaptivity by using adaptive optimizers client server optimizers. Building upon this, we develop novel adaptive optimization techniques for FL by using per-coordinate methods as server optimizers. By focusing on adaptive server optimization, we enable use of adaptive learning rates without increase in client storage or communication costs, and ensure compatibility with cross-device FI..
Main contributions In light of the above, we highlight the main contributions of the paper.
We study a general framework for federated optimization using server and client optimizers. This framework generalizes many existing federated optimization methods, including FEDAVG.
We use this framework to design novel, cross deoice compatible, adaptive federated optimization methods, and provide convergence analysis in general nonconvex settings. To the best of our knowledge, these are the first methods for FI. using adaptive server optimization. We show an important interplay between the number of local steps and the heterogeneity among clients.
We introduce comprchensive and reproducible empirical benchmarks for comparing federated optimization methods. These benchmarks consist of seven diverse and representative FL tasks involving both image and text data, with varying amounts heterogeneity and numbers of clients.
We demonstrate strong empirical performance of adaptive optimizers throughout, improving upon commonly used baselines. Our results show that our methods can be easier to tune, and highlight their utility in cross-device settings.
Related work FEDAVG was first introduced by McMahan et al. (2017), who showed it can dramati- cally reduce communication costs. Many variants have since been proposed to tackle issues such as convergence and client drift. Examples include adding a regularization term in the client objectives towards the broadcast model (Li et al., 2018), and server momentum (Hsu et al., 2019). When clients are homogeneous, FEDAVG reduces to local SGD (Zinkevich et al., 2010), which has been analyzed by many works (Stich, 2019; Yu et al., 2019; Wang & Joshi, 2018; Stich & Karimireddy, 2019; Basu et al., 2019). order to analyze FEDAVG in heterogeneous settings, many works derive convergence rates depending on the amount of heterogeneity (Li et al., 2018; Wang et al., 2019; Khaled et al., 2019; Li al., 2019b). Typically, the convergence rate of FEDAVG gets worse with client heterogeneity. By using control variates to reduce client drift, the SCAFFOLD method (Karimireddy et al., 2019) achieves convergence rates that are independent of the amount of heterogeneity. While effective in cross-silo FL, the method is incompatible with cross-device FL as it requires clients to maintain state across rounds. For more detailed comparisons, we defer Kairouz et al. (2019).
Adaptive methods have been the subject of significant theoretical and empirical study, in both convex (McMahan & Streeter, 2010b; Duchi al., 2011; 1; Kingma & Ba, 2015) and non-convex settings (Li & Orabona, 2018; Ward et al., 2018; Wu et al., 2019). Reddi et al. (2019); Zaheer et al. (2018) study convergence failures of ADAM in certain non-convex settings, and develop adaptive optimizer, YOGI, designed to improve convergence. While most work on adaptive methods focuses on non-FL settings, Xie et al. (2019) propose ADAALTER, method for FL using adaptive client optimization. Conceptually, our approach is also related to the LOOKAHEAD optimizer (Zhang et al., 2019b), which was designed for non-FL settings. Similar to ADAALTER, an adaptive FL variant of LOOKAHEAD entails adaptive client optimization (see Appendix B.3 for more details). note that both ADAALTES and LOOKAHEAD are, in fact, special cases of our framework (see Algorithm 1) and the primary novelty of our work comes in focusing on adaptive server optimization. This allows us to avoid aggregating optimizer states across clients, making our methods require at most half as much communication and client memory usage per round (see Appendix B.3 for details).
Notation For a, b E Rd, we let Va, a² and a/b denote the dlement-wise square root, square, and division of the vectors. For 0, E Rd,, we use both 0G, and [0.]; to denote its jth coordinate.
FEDERATED LEARNING AND FEDAVG
In federated learning, we solve optimization problem of the form:
1 m FF(x),
min f(x) aERd
F((x),
m
where Fx(x) EynD, (z, =)), 1, the loss function of the th client, 2 E Z, and D, is the data distribution for the 5th client. For * j. D, and D, may be very different. The functions F; (and
2

---

Published as a conference paper at ICLR 2021
therefore f) may be nonconvex. For each i and æ, assume access to unbiased stochastic gradient 9e(2)) the client's true gradient (x)).. In addition, we make the following assumptions. Assumption (Lipschitz Gradient). The function is L-smooth for all [m]] i.e. [VF(()) (V.())ll Ta| yll,for all 5. TR..
Assumption 2 (Bounded Variance). The function Fs have ot-bounded (local) variance i.e., [llV[f 2) - [VF((z)) oî for all x E RR, j E [d] and i E [mm] Furthermore, we assume the (global) variance is bounded, Rd and j E [d].
-
S og for all
Assumption 3 (Bounded Gradients). The function fo (x, 2) have G-bounded gradients i...foor any i [m]. x TRd and E 2 we have [V, z)],1 I G for all j [d]
With a slight abuse of notation, we use of and oî to denote Σj= o. and ΣØ= 02
9.3- Assumptions
and are fairly standard in nonconvex optimization literature (Reddi et al., 2016; Ward et al., 2018; Zaheer et al., 2018). We make no further assumptions regarding the similarity of clients datasets. Assumption 2 is a form of bounded variance, but between the client objective functions and the overall objective function. This assumption has been used throughout various works on federated optimization (Li et al., 2018; Wang et al., 2019). Intuitively, the parameter 0g quantifies similarity of client objective functions. Note 0g = O corresponds to the i.i.d. setting.
common approach to solving (1) federated settings is FEDAVG (McMahan et al., 2017). At each round of FEDAVG, a subset of clients are selected (typically randomly) and the server broadcasts its SGD their loss function, and send the resulting model to the server. The server then updates its global model as the average these local models. See Algorithm 3 in the appendix for more details.
global model each client. In parallel, the clients
Suppose that at round t, the server has model H, and samples set S of clients. Let at denote the model of each client i S after local training. We rewrite FEDAVG's update as
J++1
Σ LS
o =
1 S
Σ (x x;) iES
S
Let A :- and A, : (1/si) Σ Eies Ar Then the server update in FEDAVG is equivalent to applying SGD to the "pseudo-gradient" - with learning rate 1. This formulation makes it clear that other choices of 1) are possible. One could also utilize optimizers other than SGD on the clients, or use an alternative update rule on the server. This family of algorithms, which we refer to collectively FEDOPT, is formalized in Algorithm
Algorithm FEDOPT
1: Input: Zo. CLIENTOPT, SERVEROPT
2: for = 0, ... , 1 do
3: Sample a subset of clients
4: r².0
5: for each elient i E S in parallel do
6: for 1: ..... K 1 1 do
7:
Compute an unbiased estimate gi of (x(
8:
1,+1 = CLIENTOPT(Z), R. nu. t)
9:
A x 24
A.
S Eies A
11: T1+1 SERVEROPT(EL: -4, t)
In Algorithm 1, CLIENTOPT and SERVEROPT are gradient- optimizers with learning rates Ii and 17 respectively. Intuitively, CLIENTOPT aims to minimize (1) based on each client's local data while SERVEROPT optimizes from a global perspective. FEDOPT naturally allows the use of adaptive optimizers (eg. ADAM, YOGI, etc.), as well as techniques such server-side momentum (leading to FED,VGM, proposed by Hsu et al. (2019)). In its most general form, FEDOPT uses a CLIENTÖPT whose updates depend globally aggregated statistics (e.g. server updates in the
3

---

Published as a conference paper at ICLR 2021
previous iterations). We also allow v and ri to depend the round in order to encompass learning rate schedules. While we focus on specific adaptive optimizers in this work, we can in principle use any adaptive optimizer (e.g. AMSGRAD (Reddi et al., 2019), ADABOUND (Luo et al., 2019)).
While FEDOPT has intuitive benefits over FEDAVG, it also raises a fundamental question: Can the negative of the average model difference A be used as pseudo-g -radddnn in general server optimizer updates? In this paper, we provide an affirmative answer to this question by establishing theoretical basis for FEDOPT. We will show that the use of the term SERVEROPT is justified, as we can guarantee convergence across wide variety of server optimizers, including ADAGRAD, ADAM, and YOGI, thus developing principled adaptive optimizers for FL based on our framework.
3 ADAPTIVE FEDERATED OPTIMIZATION
In this section, we specialize FEDOPT to settings where SERVEROPT is an adaptive optimization method (one of ADAGRAD, YOGI or ADAM) and CLIENTOPT is SGD. By using adaptive methods (which generally require maintaining state) on the server and SGD on the clients, we ensure our methods have the same communication cost as FEDAVG and work in cross-device settings.
Algorithm provides pseudo-code for our methods. An alternate version using batched data and example-based weighting (as opposed to uniform weighting) of clients is given in Algorithm 5. The parameter controls the algorithmss' degree of adaptivity, with smaller values of representing higher degrees of adaptivity. Note that the server updates of our methods are invariant to fixed multiplicative changes to the client learning rate yi for appropriately chosen 7, though shall see shortly, will require 1/ to be sufficiently small in our analysis.
Algorithm² 2 FEDADAGRAD FEDYOGI, and FEDADAM
1: Initialization: zo, V- 2 7², decay parameters B1, 82 E [0,
2: for - ... T do
3: Sample subset S of clients
4: z².0
5:
for each client i E S in parallel do for l 0, K do
6:
7:
Compute an unbiased estimate Di of Fi(
8:
x 21,11 = 2I, mok
9:
A x 2t
10:
A 8 A 1 + (1 (1-B)) B1) (ST Eies 4!)
11:
Ut 1 A (FEDADAGRAD)
12:
U f-1 - (1 A3 sign(r:-1 A (FEDYOGI)
Ut 32 1 + (1 A2 (FEDADAM)
T++1 Ft + 1) +
We provide convergence analyses of these methods in general nonconvex settings, assuming full participation, i.e. = [m]] For expository purposes, we assume 31 = 0, though our analysis can be directly extended to 31 > 0. Our analysis can also be extended to partial participation (i.e. |S] < m, see Appendix A I for details). Furthermore, non-uniform weighted averaging typically used in FEDAVG (McMahan et al., 2017) can also be incorporated into our analysis fairly easily.
Theorem 1. Let Assumptions I to 3 hold, and let L, G, Gi og be as defined therein. Let o² = of + 6Ko² Consider the following conditions for mp:
1
1
1/3
(Condition I)
m < min
K
16L
T1/6
120L
1
1
33
1/5
1
72
1/4
(Condition II)
m <
min
3K
TT1/10
L2G3
T1/8
L*Gm
4

---

Published as a conference paper at ICLR 2021
Then the iterates of Algorithm for FEDADAGRAD sattisfy
under Condition I only,
min
E V
fxx)) 00
G
+
7
(vl + Dvar)
0StET-1
VT
KtKT
under both Condition & II,
min
E [z S 0
G
+
7
1p + L Fvar
0
VT
TTKT
Here, we define
f(xo) - f(x")
5m LTT
V
2r 0².
d(m +TnL)
7² 11 G²G
V
1 1
T
72
KT +TML
2n KT of + 10 KO m
var
72
All proofs are relegated to Appendix due to space constraints. When d satisfies the condition the second part the above result, we obtain a convergence rate depending on min V T var) J- To obtain an explicit dependence on T and K, simplify the above result for specific choice of 17 T and Corollary 1. Suppose qi is such that the conditions in Theorem I are satisfied and m ((LLLT)) Also suppose 1 O (K)) and G/t. Then, for sufficiently large T, the iterates of Al Algorithm 2 for FEDADAGR AD satisfy
E (x=)) o (f())) - -5(x))
207L. GVVKT
o²
o²T./m
min 0S1ST-1
VmKT
GKT
GVKT3/2
We defer detailed discussion about our analysis and its implication the end of the section.
Analysis of FEDADAM Next, we provide the convergence analysis of FEDADAM. The proof of FEDYOGI is very similar and hence, we omit the details of FEDYOGI'S analysis.
Theorem 2. Let Assumptions to 3 hold, and L, G, 01, g be as defined therein. Let 0² = of +KK 1/16LK and
Suppose the client learning rate satisfies <
1/2
1/4
1/3
S
min
6K
GL
GL
GL2
Then the iterates of Algorithm for FEDADAM satisfy
(xx)) = /BBTRKG +T
min 0StST-1
( + Uvar)
mKT
where
V f(xo) - -s(x*)
5ml L2T 02
2T
nI.) [n KT
2 m=²² of + 20nt
I.2T
var = G +
72
Similar to the FEDADAGRAD case, we restate the above result for specific choice of mhe 17 and T in order to highlight the dependence of K and T.
Corollary 2. Suppose tμ is chosen such that the conditions in Theorem 2 are satisfied and that 'Rt = O(1/KLLT Also, suppose 1 O(v Km) and T = G/t.. Then, for sufficiently large T, the iterates of Algorithm 2 for FEDADAM sutisfy
V[())] = o [(x)) ((w)) - -f(x')
207L
o²
aL
min
0StST-1
VmKT
GKT
GT
5

---

Published as a conference paper at ICLR 2021
The server learning rate 1) 1 rypically used in FEDAVG does not necessarity minimize the upper bound in Theorems & 2. The effect of o a measure of client heterogeneity, on convergence can be reduced by chwosing sufficiently yn and reasonably large o (e.g. see Corollary 1). Thus, the effect of client heterogeneity can be reduced by carefully choosing client and server learning rates, but not removed entirely. Our empirical analysis (eg. Figure 1) supports this conclusion.
Remark 1.
Discussion. We briefly discuss our theoretical analysis and its implications in the FL setting. The convergence rates for FEDADAGRAD and FEDADAM are similar, so our discussion applies to all the adaptive federated optimization algorithms (including FEDYOGI) proposed in the paper.
Comparison of convergence rates. When T is sufficiently large compared to K, O(/VmXT)) is the dominant term in Corollary 1 & 2. Thus, we effectively obtain a convergence rate of (1//mKT),, which matches the best known rate for the general non-convex setting of our interest (e.g. see (Karimireddy et al., 2019)). We also note that in the i.i.d setting considered in (Wang & Joshi, 2018), which corresponds to dg O, we match their convergence rates. Similar to the centralized setting, it is possible to obtain convergence rates with better dependence on constants for federated adaptive methods, compared to FEDAVG, by incorporating non-um non-uniform bounds on gradients across coordinates (Zaheer et al., 2018).
(ii) Learning rates & their decay. The client learning rate of 1/VT in our analysis requires knowl- edge of the number of rounds T' priori: however, it easy to generalize our analysis to the case where iy is decayed at a rate of 1/Vt. Observe that one must decay TJ- not the server learning rate 17, to obtain convergence. This is because the client drift introduced by the local updates does not vanish as T + oo when is constant. As we show in Appendix E.6, learning rate decay can improve empirical performance. Also, note the inverse relationship between yp and n in Corollary & 2, which we observe in our empirical analysis (see Appendix E.4).
(iii)
Communication efficiency & local steps. The total communication cost of the algorithms depends on the number of communication rounds T. From Corollary & 2, it is clear that larger K leads to fewer rounds of communication as long as K = of /a Thus, the number of local iterations can be large when either the ratio of /og or T is large. In the i.i.d setting where 0g 0. unsurprisingly, K can be very large.
(iv) Client heterogeneity. While careful selection of client and server learning rates can reduce the effect of client heterogeneity (see Remark 1), it does not completely remove it. In highly heterogeneous settings, it may be necessary to use mechanisms such as control variates (Karim- ireddy et al., 2019). However, our empirical analysis suggest that for moderate, naturally arising heterogeneity, adaptive optimizers are quite effective, especially in cross-device settings (see Figure 1). Furthermore, our algorithms can be directly combined with such mechanisms.
As mentioned earlier, for the sake of simplicity, our analysis assumes full-participation (S [m]) Our analysis can be directly generalized to limited participation at the cost of an additional variance term in our rates that depends on [S|/m, the fraction of clients sampled (see Section A.2.1 for details).
4 EXPERIMENTAL EVALUATION: DATASETS, TASKS, AND METHODS
We evaluate our algorithms on what we believe is the most extensive and representative suite of federated datasets and modeling tasks to date. We wish to understand how server adaptivity can help improve convergence, especially in cross-device settings. To accomplish this, conduct simulations on seven diverse and representative learning tasks across five datasets. Notably, three of the five have naturally-arising client partitioning, highly representative of real-world FL problems.
Datasets, models, and tasks We use five datasets: CIFAR-10, 10, CIFAR-100 (Krizhevsky & Hinton, 2009), EMNIST (Cohen et al., 2017), Shakespeare (McMahan al., 2017), and Stack Overflow (Au- thors. 2019). The first three are image datasets, the last two are text datasets. For CIFAR- 10 and CIFAR- 100, we train ResNet-18 (replacing batch norm with group norm (Hsieh et al., 2019)). For EMNIST, we train CNN for character recognition (EMNIST CR) and a bottleneck autoencoder (EMNIST AE). For Shakespeare, we train an RNN for next wwwracter-prediettiie
For Stack Overflow, we perform tag prediction using logistic regression on bag-of loug-of-wodde words vectors (SO LR) and train an RNN to do next-word prrdd-preiittion (SO NWP). For full details of the datasets, see Appendix C.
Implementation We implement all algorithms in TensorFlow Federated (Ingerman & Ostrowski, 2019). Clients are sampled uniformly at random, without replacement in a given round, but with replacement across rounds. Our implementation has two important characteristics. First, instead of

---

Published as a conference paper at ICLR 2021
doing K training steps per client, we do E epochs of training over each client's dataset. Second, to account for varying numbers of gradient steps per client, we weight the average of the client outputs A by each client's number of training samples. This follows the approach of (McMahan et al., 2017); and can often outperform uniform weighting (Zaheer et al., 2018). For full descriptions of the algorithms used, see Appendix B.
Optimizers and byperparameters We compare FEDADAGRAD, FEDADAM, and FEDYOGI (with adaptivity 7) to FEDOPT where CLIENTOPT and SERVEROPT are SGD with learning rates and n. For the server, we use momentum parameter of O (FEDAVG), and 0.9 (FEDAVGM). We fix the client batch size on a per-task level (see Appendix D.3). For FEDADAM and FEDYOGI, we fix momentum parameter 31 = 0.9 and second moment parameter 32 0.99. We also compare to SCAFFOLD (see Appendix B.2 for implementation details). For SO NWP, we sample 50 clients per round, while for all other tasks sample 10. We use E = 1 local epochs throughout.
We select TJt 17, and T by grid-search tuning. While this is often done using validation data in centralized settings, such data is often inaccessible in FL, especially cross-device FL. Therefore, tune by selecting the parameters that minimize the average training loss over the last 100 rounds of training. We run 1500 rounds of training on the EMNIST CR, Shakespeare, and Stack Overfoow tasks, 3000 rounds for EMNIST AE, and 4000 rounds for the CIFAR tasks. For more details and a record of the best hyperparameters, see Appendix D.
Validation metries For all tasks, we measure the performance on a validation set throughout training. For Stack Overflow tasks, the validation set contains 10,000 randomly sampled test examples (due to the size of the test dataset, see Table 2). For all other tasks, we use the entire test set. Since all algorithms exchange equal-sized objects between server and clients, we use the number of communication rounds a proxy for wall-clock training time.
5
EXPERIMENTAL EVALUATION: RESULTS
5.1 COMPARISONS BETWEEN METHODS
We compare the convergence of our adaptive methods to non- non-adaptive methods: FEDAVG, FEDAVGM
and SCAFFOLD. Plots of validation performances for each task/optimizer are in Figure 1, and Table 1 summarizes the last- 100-round validation performance. Due to space constraints, results for EMNIST CR are in Appendix E. and full test set results for Stack Overflow are in Appendix E.2.
Sparse-gradient tasks Text data often produces long-tailed feature distributions,
FED...
ADAGRAD
ADAM
YOGI
AvGM
AVG
leading to approximately-sparss gradients
CIFAR-10
72.1
77.4
78.0
77.4
72.8
which adaptive optimizers can capitalize
CIFAR-100
47.9
52.5
52.4
52.4
44.7
on (Zhang et al., 2019a). Both Stack Over- flow tasks exhibit such behavior, though they are otherwise dramatically different- in feature representation (bag-of-words vs. variable-length token sequence), model ar- chitecture (GLM vs deep network), and optimization landscape (convex vs non- convex). In both tasks, words that do not appear in a client's dataset produce near- zero client updates. Thus, the accumulators Vt, in Algorithm/ remain small for pa- rameters tied to rare words, allowing large updates to be made when they do occur. This intuition is born out in Figure 1, where adaptive optimizers dramatically outperform non-adaptive ones. For the non-convex NWP task, momentum is also critical, whereas it slightly hinders performance for the convex LR task.
EMNIST CR
85.1
85.6
85.5
85.2
84.9
SHAKESPEARE
57.5
57.0
57.2
57.3
56.9
SO NWP
23.8
25.2
25.2
23.8
19.5
SO LR
67.1
65.8
65.9
36.9
30.0
EMNIST AE
4.20
1.01
0.98
1.65
6.47
Table 1: Average validation performance over the last 100 rounds: % accuracy for rows 1-5; Recall (x 100) for Stack Overflow IR; and MSE ( (x 1000) for EMNIST AE. Performance within 0.5% of the best result for each task are shown in bold.
Dense-gradient tasks CIFAR-10/100, EMNIST AE/CR, and Shakespeare lack a sparse- -raddeeg structure. Shakespeare is relatively easy-moss optimizers perform well after enough rounds once suit- ably tuned, though FEDADAGRAD converges faster. For CIFAR-10/100 and EMNIST AE, adaptivity and momentum offer substantial improvements over FEDAVG. Moreover, FEDYOGI and FEDADAM have faster initial convergence than FEDAVGM on these tasks. Notably, FEDADAM and FEDYOGI perform comparably or better than non-adaptive optimizers throughout, and close
better than
7

---

Published as a conference paper at ICLR 2021
0.0
CI
0.8
QFAR 100
0.03
EM AE
Mean Squared Error
0.6
Acouacy
0.4
0.02
0.4
0.22
0.01
0.0
1000
2000
3000
4000
0.0
1000
2000
3000
4000
0.00
1000
2000
3000
Number of Rounds
Number of Rounds
Number of Rounds
0.5
Shalesspare
0.8
Stack Overfoow LR
0.3
Stack Overflow NWP
Acaracy
0.5
Felldiggnd
Rccalles
Accuracy
0.2
reskdam
Fotfigg
Fozhygg Follerg SCAFFOLD
0.3g
200 400
600 800 1000 1200 Number of Rounds
0.0
500
1000
1500
0.0
500
1000
1500
Numbber of Rounds
Num ber of Rounds
Figure 1: Validation accuracy of adaptive and non-adaptive methods, well as SCAFFOLD, using constant learning rates 1) and yn tuned to achieve the best training performance over the last 100 communication rounds; see Appendix D.2 for grids.
FEDADAGRAD throughout. As discuss below, FEDADAM and FEDYOGI actually enable easier learning rate tuning than FEDAvGM in many tasks.
Comparison to SCAFFOLD On all tasks, SCAFFOLD performs comparably to or worse than FEDAVG and our adaptive methods. On Stack Overflow, SCAFFOLD and FEDAVG are nearly identical. This is because the number of clients (342,477) makes it unlikely we sample any client more than once. Intuitively, SCAFFOLD does not have chance to use its client control variates. In other tasks, SCAFFOLD performs worse than other methods. We present two possible explanations: First, we only sample a small fraction of clients at each round, so most users are sampled infrequently, Intuitively, the client control variates can become stale, and may consequently degrade the perfor- mance. Second, SCAFFOLD is similar to variance reduction methods such SVRG (Johnson & Zhang, 2013). While theoretically performant, such methods often perform worse than SGD in practice (Defazio & Bottou, 2018). As shown by Defazio et al. (2014), variance reduction often only accelcrates convergence when close to a critical point. In cross-device settings (where the number of communication rounds are limited), SCAFFOLD may actually reduce empirical performance.
5.2
EASE OF TUNING
Stark Duerfioe NWP, TedAdam
Stack Overlow KMP, Tedtoogi
Stack Dverfow Fedhogg
1.0
1.0
1.0
0.5
25
0.5
25
0.5
25
0.0
0.0
0.0
20
20
20
-0.5
-0.5
Laamen
$2.0
$2.0
$2.0
-2.5
10
2.0
2.0
2.0
2.5
2.5
3.0
-3.0
-3.0
3.0
25
2.0
10.00055 0.5 0.0 05
25 2.0
20015100
0.0 05
3.0 25 2.0
15
1.0 055 0.0 0.5
Clend Law Kng Race (2g1a)
and Learring Rate Oegual
all Learting Rate Opgie)
Figure 2: Validation accuracy (averaged over the last 100 rounds) of FEDADAM, FEDYOGI, and FEDAVGM for various client/server learning rates combination on the SO NWP task. For FEDADAM and FEDY we set 7 10-3
Obtaining optimal performance involves tuning m. 7). and for the adaptive methods, т. To quantify how easy it is to tune various methods, we plot their validation performance as function of ? and Figure gives results for FEDADAM, FEDYOGI, and FEDAVGM on Stack Overflow NWP. Plots for all other optimizers and tasks are in Appendix E.3. For FEDAVGM, there are only a few good values of ny for each 1), while for FEDADAM and FRDYOGI, there are many good values of rp for range of
8

---

Published as a conference paper at ICLR 2021
Maan Squa Error
Accurayy
0.75
Accuracy
0.5
0.015
0.70
0.010
0.4
0.65
0.005
0.3
0.000
Adaptivity
Adaptiiity
Adaptivily
0.60
Shakesseare
0.70
Stack Overllow UR
0.30
Stack Overllow NWP
0.58
0.65
0.25
Recalies
Acorrcy
0.58
0.54
0.52
Pesesum
Reshdagrad
0.55
0.15
Pestog
0.50
6
16-1
10-2
16-4
10
0.50
10-4
10-2
16-4"
Yo-
0.10
16-5
16
10--
16-4"
10-
Acaptivity (T]
Adaptivity (T)
Adaptivity (T)
Figure 3: Validation performance of FEDADAGRAD, FEDADAM, and FEDYOGI for varying on various tasks. The learning rates and are tuned for each to achieve the best training performance on the last 100 communication rounds.
Thus, FEDADAM and FEDYOGI are arguably easier to tune in this setting. Similar results hold for other tasks and optimizers (Figures 5 to 11).
This leads to a natural question: Is the reduction in the need to tune yp and v offset by the need to tune the adaptivity 7? In fact, while tune in Figure 1, results relatively robust To demonstrate, we plot the best validation performance for various T in Figure For nearly all tasks 10-3 works almost as well all other values. This aligns with work by Zaheer
and optimizers, et al. (2018), who show that moderately large T yield better performance for centralized adaptive optimizers. FEDADAM and TEDY see only small differences in performance among all tasks except Stack Overflow LR (for which FEDADAGRAD the best optimizer, and robust to 7).
5.3 OTHER FINDINGS
We present additional empirical analyses in Appendix E. These include EMNIST CR results (Ap- pendix 1) Stack Overflow results on the full test dataset (Appendix E.2), client/server learning rate heat maps for all optimizers and tasks (Appendix E.3), an analysis of the relationship between 1) and n (Appendix E.4), and experiments with learning rate decay (Appendix E.6).
CONCLUSION
In this paper, we demonstrated that adaptive optimizers can be powerful tools in improving the conver- gence of FL. By using a simple client/server optimizer framework. we can incorporate adaptivity into FL in a principled, intuitive, and theoreticallyjjsstiiee manner. We also developed comprehensive benchmarks for comparing federated optimization algorithms. To encourage reproducibility and breadth of comparison, we have attempted to describe our experiments as rigorously as possible, and have created an open-source framework with all models, datasets, and code. We believe our work raises many important questions about how best perform federated optimization. Example directions for future research include understanding how the use of adaptivity affects differential privacy and fairness.
REFERENCES
The TensorFlow Federated Authors.
TensorFlow Federated Stack Overflow dataset, wwt.://///feeeeeee////a...ee///
2019. URL htt.ttt//// http:p::s//wwtttssofolllll./ s httpa/////attteeeetttttteeeeeeee

---

Published as a conference paper at ICLR 2021
Debraj Basu, Deepesh Data, Can Karakus, and Suhas Diggavi. Osparse-local-SOG Distributed SGD with quantization, sparsification and local computations. In Advances Neural Information Processing Systems, pp. 14668-14679, 2019.
Keith Bonawitz, Hubert Eichner, Wolfgang Grieskamp, Dzmitry Huba, Alex Ingerman, Vladimir Ivanov, Chloé Kiddon, Jakub Konečný, Stefano Mazzocchi, Brendan McMahan, Timon Van Overveldt, David Petrou, Daniel Ramage, and Jason Roselander. ated learning at scale: System design. In A. Talwalkar, V. Smith, and M. Zaharia (eds.), Proceedings of Machine Learning and Systems, volume 1, pp. 374-388. Proceed- ings of MI.Sss, 2019. URL https::/p/// http://///ccoeooeeeeeeooee ml sys... wwrg//.r///eooo www.o/ppggrppaeeeeoeee (aaper/2011 9/file/ 98e faeeeee0eea0000 fa301 P-Paper. pdf..
Towards feder-
Sebastian Caldas, Peter Wu, Tian Li, Jakub Konečný, H Brendan McMahan, Virginia Smith, and Ameet Talwalkar. LEAF; A benchmark for federated settings. arXiv preprint 1812.007 2018.
Gregory Cohen, Saeed Afshar, Jonathan Tapson, and Andre Van Schaik. EMNIST: Extending MNIST to handwritten letters. 2017 International Joint Conference on Neural Networks (LJCNN), pp- 2921-2926. IEEE, 2017.
Aaron Defazio and Léon Bottou. On the ineffectiveness of variance reduced optimization for deep learning. arXiv preprint 1 2018.
Aaron Defazio, Francis Bach, and Simon Lacoste-Julien SAGA: A fast incremental gradient method with support for non-strongly convex composite objectives. NIPS, pp. 1646-1634, 2014.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(uu))2121-2159, 2011.
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jía, and Kaiming He. Accurate, large minibatch SGD: Training ImageNet in hour. arXiv preprint arXiv:/006.02677, 2017.
Kevin Hsieh, Amar Phanishayee, Onur Mutlu, and Phillip B Gibbons. The non-IID data quagmire of decentralized machine learning. arXiv preprint urXiv:1910.00150 2019.
Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the effects of non- con-identical data distribution for federated visual classification. arXiv preprint vrXiv::909.063: 06335, 2019.
Alex erated, 2019.
Ingerman and Krzys
Ostrowski. URL
Introducing TensorFlow Fed-
fetps:://mddiie com/ http://ensorfw/////ooo
Rie Johnson and Tong Zhang. Accelcrating stochastic gradient descent using predictive variance reduction. In Advances in Neural Information Processing Systems, pp. 315-323, 2013.
Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Keith Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning. arXiv preprint arXiv: 2019.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank J Reddi, Sebastian U Stich, and Ananda Theertha Suresh. SCAFFOLD: Stochastic controlled averaging for on-device federated learning. arXiv preprint //X.0:19/00006378, 2019.
Ahmed Khaled, Konstantin Mishchenko, and Peter Richtárik. First analysis of local GD on heteroge- 2019.
neous data. arXiv preprint
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA. USA, May 7-9, 2015, Conference Track Proceedings, 2015.
Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images. Technical report, Citeseer, 2009.
10

---

Published as a conference paper at ICLR 2021
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization heterogeneous networks. arXiv preprint arXiv: 18/2.06177 2018.
Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Challenges, methods, and future directions. arXiv preprint arXiv:/908.0 978807773, 2019a.
Wei Li and Andrew McCallum. Pachinko allocation: DAG-structured mixture models of topic correlations. In Proceedings of the 23rd international conference Machine learning, pp. 577- 584, 2006.
Xiang Li, Kaixuan Huang, Wenhao Yang. Shusen Wang, and Zhihua Zhang. On the convergence of FedAvg on non-IID data. arXiv preprint 1907:02/0002/88 2019b.
Xiaoyu Li and Francesco Orabona. On the convergence of stochastic gradient descent with adaptive stepsizes. arXiv preprint arXiv:1805.081 14, 2018.
Liangchen Luo, Yuanhao Xiong, Yan Liu, and Xu Sun. Adaptive gradient methods with dynamic bound of learning rate. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL https::/ net/Eorum? i BKg
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Agüera y Arcas. Communication-efficiee learning of deep networks from decentralized data. In Proceedings of the 20th International Conference on Artificial Intelligence and Statisties, AISTATS 2017, 20-22 April 2017, Fort Lauderdale, FL, USA, pp. 1273-1282, 2017. URL http::t/poccceeiiiiii.// press 454 memahan171 7a html.
mlr.
H. Brendan McMahan and Matthew Streeter. Adaptive bound optimization for online convex optimization. In COLT, 2010a.
H. Brendan McMahan and Matthew Streeter. Adaptive bound optimization for online convex optimization. In COLT The 23rd Conference on Learning Theory, 2010b.
Sashank Reddi, Ahmed Hefny, Suvrit Sra, Barnabás Póczós, and Alex Smola. Stochastic variance reduction for nonconvex optimization. arXiv:1603.0616) 2016.
Sashank Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of ADAM and beyond. arXiv preprint arXiv:1904.092. 2019.
Sebastian U. Stich. Local SGD converges fast and communicates little. In International Confer- ence on Learning Representations, 2019. URL h/tttp:://penrreviee. weteeeiwwwwweeeeeeeeeeee...eeeee / S1g2JnRcFX.
http://wwwwweet..nmteeeom forum?id-
Sebastian U Stich and Sai Pranceth Karimireddy. The error-feedback framework: Better rates for SGD with delayed gradients and compressed communication. arXiv preprint arXiv:1909.0533: 2019.
Jianyu Wang and Gauri Joshi. Cooperative SGD: A unified framework for the design and analysis of caamuuncation-efffccen SGD algorithms. arXiv preprint 1.X.::8000.0770 2018
Shiqiang Wang. Tiffany Tuor, Theodoros Salonidis, Kin K Leung, Christian Makaya, Ting He, and Kevin Chan. Adaptive federated learning in resource constrained edge computing systems. IEEE Journal on Selected Areas in Communications, 37(6)) 1221, 2019.
Rachel Ward, Xiaoxia Wu, and Leon Bottou. Adagrad stepsizes: Sharp convergence over nonconvex landscapes, from any initialization. arXiv preprint 10X1v:1806.0/8//, 2018.
Xiaoxia Wu, Simon S Du, and Rachel Ward. Global convergence of adaptive gradient methods for an pwer-parameterized neural network. arXiv preprint
2019.
Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 3-19, 2018.
Cong Xie, Oluwasanmi Koyejo, Indranil Gupta, and Haibin Lin. Local AdaAlter: Communication- efficient stochastic gradient descent with adaptive learning rates. arXiv preprint 2019.
11

---

Published as a conference paper at ICLR 2021
Hao Yu, Sen Yang, and Shenghuo Zhu. Parallel restarted SGD with faster convergence and less communication: Demystifying why model averaging works for deep learning. In Proceedings of the AAAI Conference Artificial Intelligence, volume 33, pp. 5693-5700, 2019.
Manzil Zaheer, Sashank Reddi, Devendra Sachan, Satyen Kale, and Sanjiv Kumar. Adaptive methods for nonconvex optimization. In Advances in Neural Information Processing Systems, pp. 9815- 9825, 2018.
Jingzhao Zhang, Sai Pranecth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank J. Reddi, Sanjiv Kumar, and Suvrit Sra. Why ADAM beats SGD for attention models. arXiv preprint 1912::1912.03194 2019a.
Michael R. Zhang, James Lucas, Jimmy Ba, and Geoffrey E. Hinton. Lookahead optimizer: k steps forward, step back. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-BBuc Emily B. Fox, and Roman Gamett (eds.). Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, 8- 14 December 2019, Vancouver, BC, Canada, pp. 9593-9604, 2019b.
Martin Zinkevich, Markus Weimer, Lihong Li, and Alex J Smola. Parallelized stochastic gradient descent. Advances in neural information processing systems, pp. 2595-2603, 2010.
12

---

Published as a conference paper at ICLR 2021
A PROOF OF RESULTS
A.1 MAIN CHALLENGES
We first recap some of the central challenges to analysis. Theoretical analyses of optimization methods for federated learning are much different than analyses for centralized settings. The key factors complicating the analysis are:
1. Clients performing multiple local updates.
2. Data heterogeneity.
3. Understanding the communication complexity.
As result of (I), the updates from the clients the server not gradients, even unbiased estimates of gradients, they are pseudo-gradients gradients (see Section 2). These sseudo-gradient are challenging to analyze they can have both high bias (their expectation is not the gradient of the empirical loss function) and high variance (due to compounding variance across client updates) and are therefore challenging to bound. This is exacerbated by (2), which we quantify by the parameter dg in Section 2. Things are further complicated by (3), as we must obtain a good trade-off between the number of client updates taken per round (K in Algorithms I and 2) and the number of communication rounds T. Such trade-offs do not exist in centralized optimization.
A.2 PROOF OF THEOREM 1
Proof of Theorem I. Recall that the server update of FEDADAGRAD the following
A
Te+1.i Ne.i +
Vue. +
for all i E [d]. Since the function is L-smooth, have the following:
f(1141) 2 + (Vf(z)), I4+1 - -xe) +
|| 1oll2
D.
7
A
=f(x)) + vf(xx).
+
(2)
/ur +
2
Ut,
7)²
t=1
The second step follows simply from FEDADAGRAD" AD's update. We take the expectation of f(1++1) (over randomness at time step t) in the above inequality:
L
A
[.([(++1)] & f(x)) + 17 (f(x),, E.
B
4,
V +
(Vut,
T)2
t=1
A
A
A
= f(ax) + +1 (z(z)), E.
Lve +
/
Vt-1
+
Vt-I
+
172T.
A2
E,
( + +)2
j=1
A
A
4
=((x)) + +1 Vf(x)), E,
+1
V ff(x)),
/
VTe +
V +
T
Tb
d
A t.j
E,
(3)
2
(VTr.) + +
13

---

Published as a conference paper at ICLR 2021
We will first bound T2 in the following manner.
A.
T2 = V (1a)), E.
Vu-1
D.s
A 1
E, > [
Vor. +T
= E, j=1
(1))) x Ans
(vUt.) + +T)
and recalling o 1--1 + A so -3j (1-01-1 Vt. (V + vor,)) we have,
- [/a)+t)( + TT) (T((V0-11)) + vUss)
= E, Σ
f(st)) x Des
< E, [f(x))] x | x
(viis + [T( + (V/011 + vur))
j=1
A7
< E,
[f(z))] x |
since V2-1, 2 72.
(vr, + +2) (1-1. +T)
Here vt-1j since v-1 2 (see the initialization of Algorithm and Veyy increasing The above bound can be further upper bounded in the following manner.
A
T2 < E, Σ KKG
since [Vf(x)) S G and D., S KG
(UL
+
(2)
1
+T)
j=l
MKG2
2 L
S B.
since Vt-1j 2
(4)
T
1=
+
Bounding Ti We now turn our attention bounding the term Ti, which we need be sufficiently negative. We observe the following:
T
V (xxe), F.
(a()))
E, AA.
Tk (n
Vf(z)) Vut-1 +
K
E, [A. + {KV((x)))
(5)
Pt-15 +5
Ts
In order to bound Ti we use the following upper bound on T3 (which captures the difference between the actual update A, and an appropriate scaling of (xx))
(((t))
T3
E, A
m
KV [V
11-1
m
K
ngg + K (VV((a)))
E,
/1--1 +
m
Vf(xz)
K
+
E,
VU-T +T"
K=
14

---

Published as a conference paper at ICLR 2021
Here we used the fact that fx(=)) = #h E (2((2t) and gi,x is an unbiased estimator of the gradient at zi, we further bound T3 as follows using a simple application of the fact that ab s (a² + b2)
m
[
m E. 2K
m
VR(CTN))
m K-
VF((2)
T3 5
2
+T
V
m
V+-++
j=1
k=0
K
[V
M EE,, 2m
V.(( - (P)(
2
Vi- + 7
V
j=1
7
[V
m I.2
m
K-1
E,
I - x4 1??
using Assumption and 8t-1 2 0.
2
V
+T
2mr
j=1
(6)
The second inequality follows from Lemma 6. The last inequality follows from L-Lipschitz nature of the gradient (Assumption 1). We now prove a lemma that bounds the "drift" of the ri, from Tz:
Lemma 3. For any step-size satisfying mp S as
1 SLK
we can bound the drift for any l E (0,...,A K-13
1
m EE 2t S 5K E Σ (ots +
30
m
Proof. The result trivially holds for h = 1 since =F,00 2, for all E [m]. We now turn our attention
to the case where & 2 1. To prove the above result, we observe that for any client i E [m]] and k E [K].
E E| - xell2 = ==El== E1I - x2 - m99f,--ll
BB|=,,,-- - It. - n - F. + VR(CE, - V+((+)+ + VB)))) V/(+))))/((++++(
1
1 + 2K -1
E E - 2Fl1 + - VR(
+ (V
VF(( (ww.l//m//mm........ [f=)))) + [[KEIl/V<<=>)))
The first inequality uses the fact that -İ-1, is an unbiased estimator of (F and Lemma 7. The above quantity can be further bounded by the following:
1
E 21
-2e112 -
[+
El Bllai xell² + + Σ oi + (6 x2)1
2K
+ GXT[IIn - V + V<
1+ 2K-
1
+ +6K
Bll
zz )|1 ト Σ ot.
+ +6KE|lm - (( (I
Here, the first inequality follows from Assumption 1 and 2. Averaging over the clients i, we obtain the following:
El x=1225
(
1
1
6K
1
- xz
nH E Σ
oi
2
2K-
j=1
6K m
+
E (elll
m
m Σ
+ 6K 122
- xal1 + n E Σ
(o + 6Ko
2k-
1=1
15

---

Published as a conference paper at ICLR 2021
From the above, we get the following inequality:
m
1 m m
E - - 1rll² =
B B 2+112 + 1B Σ (oï + 6K 6Ko 93))
m
K
+ +6K
E
(
Unrolling the recursion, we obtain the following:
P
E|
Σ
1
+
T E
(o + +6Kog.) + 6K Bl
K
p=0
S S(K - 1)x (1=
1
-1
x
1/1
(i 6K)
K
d > (oî + 6Ko + 30K +003)
<
Sf E ;=1
concluding the proof of Lemma 3. The last inequality uses the fact that (1 + R-T) )k < for
K x>ı. 1.
Using the above lemma in Equation 6 and Condition I, we can bound T3 in the following manner:
K
I. 2mT
K 1 d
Ts S
+
E,
[1 (xxel))
2
+
j=1
1=1
km0 j=1
K d
[xx())))
?KL
d
5k E
(o + 6Ko 30K3 +[
2
+
2T
j=1
3mp
K
d
13
5m 22
<
E
(o + 6K0
4
1-11
27
j=1
Here used the fact that bound Equation 5, we get
1 S m KGVT and Condition I in Theorem 1. Using the above
K TI
5m
K2
EE (ot + +,
4
+T
2T
Putting the pieces together in Equation (4), we obtain
Substituting in Equation (3), bounds Ti in Equation (8) and bound T2
K
5m K2
B.ff(r+11 S f(xe) +
E
(oi
6Ko
1t-15
+T
27
j=1
KG2 m
A
11
A tj
x E
LIE
(9)
Σ
E
+
2
j=1
Σi=
A + 72
Rearranging the above inequality and summing it from = t T 1, we get
Σ
K
E
[
((x)) - [[[([r)]]
12
(oī + +6K02)
+
7
2T
A=
j=1
j=1
T-1
G2
nT.
A2
+
x
(10)
2
Σi=0
+
1=
j=1
16

---

Published as a conference paper at ICLR 2021
The first inequality uses simple telescoping sum. For completing the proof, we need the following result.
Lemma 4. The following upper bound holds for Algorithm 2 (FEDADAGRAD):
T-1 d
A
d
1+ o R
1
min
d
Σ
log
+
72
==1
4nF KT
Σ
d
20m
k I TE
(o
+ 6Kogg
+
40 10 T-1
Σ
[1V
H
mT² j=1
j=1 72
72
Proof. We bound the desired quantity in the following manner:
T
D <d+Ey j=1 Li Ai +
1+ - A
d <d+
1+ K2T
Σ log
1+
S
Σ
log
72
72
j=1
An alternate way of the bounding this quantity is follows:
T-1 d
A
T-1 d
A i
E
E
x2
t=
1=0 j=1
T
A m
KV f(x))
m
KV F (3+z)
2
[x+())) 27
2E
+ nf x2
(11)
The first quantity in the above bound can be further bounded as follows:
21E
Σ moc + kKV((z))
m
2F
Σ
- -mKV((m))
m
2m
Σ K-
E
Σ
(ol, -V-(()) + VV -
4v7
Σ
ΣΣ
(oin - -VR(Cttt)
+ΣΣ h=
K
(VFF(at -VR(()))
E
+
m²
0=0
4np KT
of
E
(VE((E)))- -
m
m
i=1
K=
t=
4np KT
4 K
K- 1
L.
(EiA- -20) l|
E
(by Assumptions 1 and 2)
72
m
j=1
k=
t=
4nf KT mr?
d
(of, + 6Ko2
22g)
40n K2 T-1
Σ
+0m
A
3
L
2T
E
[[V](())))
(by Lemma 3).
22
72
Here, the first inequality follows from simple application of the fact that ab S (a² + 12/ The result follows.
17

---

Published as a conference paper at ICLR 2021
Substituting the above bound in Equation (10), we obtain:
mK
T-1
Vo- +T
(3o) [IT]
T-1
E
( + +,)
7
2T
1=0
5=1
d
mKG²
n/
1h
miKKTT)
+
min
d dd
1+
2
72
j=1
4n KT
d
o +
+0
k3 1
d (of, =KKo2s)
+
40mf
T-1
Σ E
[
|
}
mT²
72
72
j=1
1=1
0
(12)
We observe the following:
T-1 ΣΣ j=1 W-
[Ef(a
T-1 d
[V( KKGGT++ n
T + +
ΣE t= j=1
E
Z
min E f
0stsT
The second part of Theorem I follows from using the above inequality in Equation (12). Note that the first part of Theorem is obtain from the part of Lemma 4.
A. LIMITED PARTICIPATION
For limited participation, the main changed in the proof is in Equation (11). The rest of the proof is similar so we mainly focus on Equation (11) here. Let S be the sampled set at the 7h iteration such that |S| = In partial participation, we assume that the set S is sampled uniformly from all subsets of [m]] with size 8. In this case, for the first term in Equation (11), we have
K-1 mol. + m VVKV((())
2E
Σ
S
iES
K=0
T-1
K-1
2E
(noin - TI /FFF(( + m - - - f(a(
mI
ES
km 0
6 T-1
K-
K
Σ
Σ
Σ
(oi. - vFL(ott))
+ ΣΣ (VF(CCN)-VA(() VF(()))
$
2² 1=0
iES 0=
iES R=0
27
1
+ KK²
7
iES
6n KT
Σ
of
6n
K
E
:. (VE((E) VF((o.))
67 KnR
(1-m)
8
8
72
$
72
m
j=1
iES
0=0 t=
Gn KT
Σ
o.
+
6m
E
K-
1
(z - xx1)
+
6n, onTK2To2
3
(1- m
8)
72
72
j=l
SS
0=
t=
6n KT
(oī + 6Ko2 2id
K2 T-1
Σ
ot
+
30m
+
E |}{({)))
72s
s
72
72
j=l
t=
6nf T2T
g
(1-
$
72
m
Note that the expectation here also includes S. The first inequality is obtained from the fact that (a ba)b+c)2= & 3(a2 + 12 +2). The second inequality is obtained from the fact that set S is uniformly
18

---

Published as a conference paper at ICLR 2021
sampled from all subsets of [m]] with size The third and fourth inequalities similar to the one used in proof of Theorem Substituting the above bound in Equation (10) gives the desired convergence rate.
A.3 PROOF OF THEOREM 2
Proof Theorem 2. The proof strategy is similar to that of FEDADAGRAD except that we need to handle the exponential moving average in FEDADAM. We note that the update of FEDADAM is the following
A
Jt+1 It +1
V +
for all E [d]. Using the L-smooth nature of function and the above update rule, have the following:
A
7 L
A (Vut,
$ f(x)) + +n Vf(xz).
+
(13)
Voe
2
+)²
The second step follows simply from FEDADAM's update. We take the expectation of f(x++1) (over randomness at time step t) and rewrite the above inequality as:
A
A V +T
A
12L
[.[[(+++)]] S (are) + +
Vs(x)),
R2
V0 +
v B2t0-11
+
2
Uv + T)
A.
A.
= f((x) +
Vf(xx), E.
+1
Vf(xz), Es
3201-1 +
Voi
+T
v B2ur-1+ +T
R1
R2
1) T
A bj
Σ j=1
E₁
(V + T)2
(14)
Bounding Rg.
We observe the following about R2:
A
A esj
R₂ R,
Σ
Vol. +T
VB2ur 1o +T
/=1
VB20e-135 1i Un.
E. Σ
[x:)) Aus x
(vor. + T)( 12 +
/=1
-(1 - -12)
R, Σ
[f()))) x A.j x
(Vors + +)(( ( 121r + (v
12 + Uv.5)
/=1
A7,
(1 - E
[ x |A
(Woes T)( B2Vr + 7) 32 / Vt.i)
==1
A7,
5 S E
[)(x))>1 x
/oe. + T) 32 +)
1=1
G
A2
V E3E Σ
Un.j +
19

---

Published as a conference paper at ICLR 2021
Bounding R1.
The term R1 can be bounded as follows:
D
R1
V [f(t), E
1/20--17
((()))
Br SA. (VV((() + -VV((())))
1 +T
[V(+))
Vf(z))
En [. m K ((r())))
(15)
12 +
B2u +T
Rs
Bounding R3. Theorem 1:
The term R3 can be bounded in exactly the same way as term T3 in proof of
A d [f(2))
mL 2m
K
R3
S
x211? 11
2
B20r-15 +
R=0
Substituting the above inequality in Equation (15), we get
A.K
[ + B2 +
m LL 2mn
m K
R1 S
E,
| 2z
2
k=0
Here used the fact that /ur-1 mKG and conditions in Theorem 2. Using Lemma 3, obtain the following bound on Ry:
d
[
KnpKKL2
R1
S
E,
> Cof (Ko
(16)
4
1
+7
2r
J=1
Putting pieces together.
Substituting bounds R₁ and R2 in Equation (14), have
mp K
[
5mm R 122
s f(x))
Σ
>>
(ats + (Kop.)
4
/Bu 1,
+T
2T
5=1
j=1
n
B2G
A
112
A ij
E
E,
U1,5
+
2
Ut.j + +2)
j=1
Summing over f = to and using telescoping sum, we have
mm K
T-1
d
5mm
1.2T
B[((xr) < fxx)
E
E
Oo + 6Ko 2.1
V B200-11
+T
27
t=0
j=1
j=1
I B&G
T-1
A
1121.
T-1
A
+
+
T)
2
V1, + 72)
t=
(17)
To bound this term further, we need the following result.
Lemma 5. The following upper bound holds for Algorithm (FEDADAM):
T-1
A7.
KT
d Σ oo
20m
K3
d
40m 40m/ T
E
(oî + 6KKo
+
|V
(UL. + +2)
mt2
72
72
1=1
j=1
Proof.
T-1
A E
T-1 d
E
Σ
< E
(1
B₂)
B +
72
1=0
j=1
1=0
1=1
20

---

Published as a conference paper at ICLR 2021
The rest of the proof follows along the lines of proof of Lemma 4. Using the same argument, get
4n KT
20m/
K
E
E
(ots + 6Ko
(Vt.j + 72)
mT²
72
j=1
40n K 1
T
[
2
which is the desired result.
Substituting the bound obtained from above lemma in Equation (17) and using a similar argument for bounding
B2G
T-
T)
1=
obtain
mi
T-1
[
5mm LTT
[.[((TT]] S f(ro)
+
>
(aï + +KKo
8
/B211-12 +
2T
1= j=l
j=1
IV I I- 3G+ +
1/2b
x
4n² KT
Σ
01
+
20n
K
L
> j=l
(o + 6Ko
2
mr²
72
j=1
The above inequality is obtained due to the fact that:
(VI- -1-3,c+ 2
nT. 40nf 12
11 K
( +
++)
7
16
VFmKG
The above condition follows from the condition on m in Theorem 2. We also observe the following:
T-1
[++)))]
ΣE
T-1 d
((())
T min Vfx)11
B2
V2mKKG+
VB2nKG + OS1ST
t=0 j=1
1=0 j=1
Substituting this bound in the above equation yields the desired result.
A.4 AUXILIARY LEMMATTA
Lemma 6. For random variables 21, Ze, we have
E [[|=1 + + zr [12] S ri [| +... || 113]
Lemma 7, For independent, mean O random variables 21,.... Zy. we have
E [[|21 + ... 2, [1-] E [||
||=
B FEDERATED ALGORITHMS: IMPLEMENTATIONS AND PRACTICAL. CONSIDERATIONS
B.1 FEDAVG AND FEDOPT
In Algorithm 3, we give simplified version of the FEDAVG algorithm by McMahan et al. (2017), that applies to the setup given Section We write SGD (æt, m. fi) to denote K steps of SGD using gradients V fx. 2) for 2 ~ D; with (local) learning rate starting from As noted in Section 2, Algorithm is the special case of Algorithm I where CLIENTOPT SGD, and SERVEROPT SGD with learning rate 1.
While Algorithms 2, and 3 useful for understanding relations between federated optimization methods, we are also interested in practical versions of these algorithms. In particular, Algorithms 1, 2, and 3 stated in terms of kind of 'gradient oracle', where compute unbiased estimates
21

---

Published as a conference paper at ICLR 2021
Algorithm 3 Simplified FEDAVG
Input: 20 for t = 0.... ... T -111 do
Sample subset S of clients
r = 2t
for each client E S in parallel do
x SGDx (I,, fi) for E S (in parallel)
F++1 = 1 [8] Di
of the client's gradient. practical scenarios, often only have access finite data samples, the number of which may vary between clients.
Instead, assume that in (1), each client distribution D; is the uniform distribution some finite set D, of size n4- The n; may vary significantly between clients, requiring extra care when implementing federated optimization methods. We assume the set D; is partitioned into collection of batches Br, each of size B. For b E B, we let fr;;; b) denote the average loss on this batch at a with corresponding gradient Vft(;;;) b). Thus, if b is sampled uniformly at random from B,, V b) an unbiased estimate V F((z).
When training, instead of uniformly using K gradient steps, in Algorithm 1, we will instead perform E' epochs training over each client's dataset. Additionally, we will take a weighted average of the client updates, where we weight according to the number of examples dataset. This leads a batched data version of FEDOPT in Algorithm 4, and batched data version of FEDADAGRAD, FEDADAM, and FEDYOGI given in Algorithm 5.
in each client's
Algorithm 4 FEDOPT Batched data
Input: 2g. CLIENTÖPT, SERVEROPT for t = ......................... ... T 1 do
Sample subset S of clients
x It
for each client E S in parallel do for e = 1,..., E do
for b B, do
g = Vi(lz(; b) x = CLIENTOPT oi m)
A xf Zt
n Σ ies the A. Eies nu A
n It+1 = SERVEROPT(EL. - D.. 17,
In Section 5,
Algorithm when implementing FEDAVG and FEDAVGM. In particular, FEDAVG and FEDAVGM correspond to Algorithm 4 where CLIENTOPT and SERVEROPT are SGD. FEDAVG uses vanilla SGD the server, while FEDAVGM uses SGD with momentum parameter of 0.9. In both methods, we tune both client learning rate !n and server learning rate 7). This means that the version of FEDAVG used in all experiments is strictly more general than that in (McMahan al., 2017), which corresponds FEDOPT where CLIENTÖPT and SERVEROPT are SGD, and SERVEROPT has learning rate of 1.
We use Algorithm 5 for all implementations FEDADAGRAD. FEDADAM, and FEDYOGI in Section For FEDADAGRAD, we set 31 = 32 0 (as typical versions of ADAGRAD do not use momentum). For FEDADAM and FEDYOGI
set 31 0.9, B2 0.99. While these parameters are generally good choices (Zaheer et al., 2018), we emphasize that better results may be obtainable by tuning these parameters.
B.2 SCAFFOLD
As discussed in Section 5, we compare all five optimizers above to SCAFFOLD (Karimireddy et al., 2019) our various tasks. There few important notes about the validity of this comparison.
22

---

Published as a conference paper at ICLR 2021
Algorithm 5 FEDADAGRAD
FEDYOGI, and FEDADAM Batched data
Input: 20, 1-1 2 72, optional $1,32 [0, 1) for FEDYOGI and FEDADAM for t = 0, ... T do
Sample subset S of clients
x = 2z
for each client i E S in parallel do for 1,.... E do
for b E B,
do
x = =
Th (af;;
A = of
Et
Σ ies Mii A, Σ Lies n
A
A = 31 A + (1 - 31)
0t = 1- + A (FEDADAGRAD)
1/ Pt1 (1 A sign(e,-1 A7) (FEDYOGI)
32 + (1 33 (FEDADAM)
I++1 = +
A +
1. In cross-device settings, this is not a fair comparison. In particular, SCAFFOLD does not work settings where clients cannot maintain state across rounds, as may be the case for federated learning systems on edge devices, such as cell phones.
2. SCAFFOLD has two variants described by Karimireddy et al. (2019). In Option I, the control variate of client is updated using full gradient computation. This effectively requires performing an extra pass over each client's dataset, as compared Algorithm 1. In order to normalize the amount of client work, we instead use Option II, in which the clients' control variates are updated using the difference between the server model and the client's learned model. This requires the same amount of client work as FEDAVG and Algorithm 2.
For practical reasons, we implement a version of SCAFFOLD mirroring Algorithm 4, in which we perform E' epochs over the client's dataset, and perform weighted averaging of client models. For posterity, we give the full pseudo-code the version SCAFFOLD used our experiments in Algorithm 6. This is simple adaptiation of Option II of the SCAFFOLD algorithm in (Karimireddy al., 2019) to the same setting as Algorithm particular, we let v., denote the number examples in client i's local dataset.
Algorithm 6 SCAFFOLD, Option II Batched data
Input: 2o. C, 1/,
for t= 0, ... . do
Sample subset S clients
o! = 3t.
for each client E S in parallel do
for E = 1,.... E do
for b B, do
g = V f((((;b) b)
r; x mg e₁ + c)
ct = G c (E|B, (1)) (a) 24)
Ar. x/ Tτ, Dc; = =c Ci
G c
ai Ac n = Eies Di, Ar = Eies ni Ari, Ae= Eies
n
It+1 = I, + Drz C c Ac
In this algorithm, &; is the control variate of client and e the running average of these control variates. In practice, we must initialize the control variates in some way when sampling a client for the first time. In our implementation,
G the first time sample client This has
23

---

Published as a conference paper at ICLR 2021
the advantage of exactly recovering FEDAVG when each client is sampled most once. To initialize e, we use the zeros vector.
We compare this version of SCAFFOLD FEDADAGRAD, FEDADAM, FEDYOGI, FEDAVGM, and FEDAVG on our tasks, tuning the learning rates the same way (using the same grids as in Appendix D.2). In particular, Tj. ? are tuned to obtain the best training performance over the last 100 communication rounds. We use the same federated hyperparameters for SCAFFOLD as discussed in Section 4, Namely, we set E' 1, and sample 10 clients per round for all tasks except Stack Overflow NWP, where sample 50, The results are given Figure Section 5.
B.3 LOOKAHEAD, ADAALTER, AND CLIENT ADAPTIVITY
The LOOKAHEAD optimizer (Zhang et al., 2019b) primarily designed for non-FL settings. LOOKA- HEAD uses a generic optimizer in the inner loop and updates its parameters using a "outer" learning rate. Thus, unlike FEDOPT, LOOKAHEAD ferent. In fact, LOOKAHEAD can be seen as a special case of FEDOPT in non-Fl. settings which uses a generic optimizer CLIENTÖPT client optimizer, and SGD the server optimizer. While there multiple ways LOOKAHEAD could be generalized to a federated setting, one straightforward version would simply
single generic optimizer and is thus conceptually dif-
adaptive method the CLIENTOPT. On the other hand, ADAALTER (Xie et al., 2019) is designed specifically for distributed settings. In ADAALTER, clients use a local optimizer similar ADAGRAD (McMahan & Streeter, 2010a; Duchi et al., 2011 1) to perform multiple epochs of training on their local datasets. Both LOOKAHIAAD and ADTELTER use client adaptivity, which is fundamentally different from the adaptive server optimizers proposed in Algorithm 2.
To illustrate the differences, consider the client-to-server communication in ADAALTER. This requires communicating both the model weights and the client accumulators (used to perform the adaptive optimization, analogous U, Algorithm 2). the case of ADAALTER, the client accumulator is the same size as the model's s trainable weights. Thus, the client-to-sere communication doubles for this method, relative FEDAVG. ADAALTER, the server averages the client accumulators and broadcasts the average to the next set of clients, who use this to initialize their adaptive optimizers. This means that the server-to-client communication also doubles relative FEDAVG. The same would occur for any adaptive client optimizer in the distributed version of LOOKAHEAD described above. For similar reasons, needed on the client (as they must store the current accumulator). By contrast, our adaptive server methods (Algorithm 2) do not require extra communication client memory relative to FEDAVG. Thus, we see that server-side adaptive optimization benefits from lower per-round communication and client memory requirements, which are of paramount importance for FL applications (Bonawitz al., 2019).
that client adaptive methods also increase the amount of memory
C DATASET & MODELS
Here we provide detailed description of the datasets and models used in the paper. We use federated versions of vision datasets EMNIST (Cohen al., 2017), CIFAR-10 (Krizhevsky & Hinton, 2009), and CIFAR- 100 (Krizhevsky & Hinton, 2009), and language modeling datasets Shakespeare (McMa- han et al., 2017) and StackOverflos (Authors, 2019). Statisties for the training datasets can be found in Table 2. We give descriptions of the datasets, models, and tasks below. Statistics on the number of clients and examples in both the training and test splits of the datasets are given in Table 2.
Table 2: Dataset statistics.
DATASET
TRAIN CLIENTS
TRAIN EXAMPLES
TEST CLIENTS
TEST EXAMPLES
CIFAR-10
S00
50.000
100
10,000
CIFAR-100
500
50,000
100
10,000
EMNIST-62
3,400
671,585
3,400
77,483
SHAKESPEARE
715
16,068
715
2,356
STACKOVERFLOW
342,477
135,818,730
204,088
16,586,035
24

---

Published as a conference paper at ICLR 2021
C.I CIFAR-IO/CIFAR
We create a federated version of CIFAR-10 by randomly partitioning the training data among 500 clients, with each client receiving 100 examples. We use the same approach Hsu et al. (2019), where we apply latent Dirichlet allocation (LDA) over the labels of CIFAR-10 to create a federated dataset. Each client has an associated multinomial distribution over labels from which its examples are drawn. The multinomial is drawn from a symmetric Dirichlet distribution with parameter 0.1.
For CIFAR-100, we perform a similar partitioning of 100 examples to 500 clients, but using more sophisticated approach. We use two step LDA process over the coarse and fine labels. We randomly partition the data to reflect the "coarse" and "fine" label structure of CIFAR- 100 by using the Pachinko Allocation Method (PAM) (Li McCallum, 2006). This creates more realistic client datasets, whose label distributions better resemble practical heterogeneous client datasets. We have made publicly available the specific federated version of CIFAR-I06 we used for all experiments, though we avoid giving a link in this work in order to avoid dc-anonymization. For complete details on how the dataset was created, see Appendix
We train a modified ResNet- 18 on both datasets, where the batch normalization layers are replaced by group normalization layers (Wu & He, 2018). We use two groups in each group normalization layer. As shown by Hsieh et al. (2019), group normalization can lead to significant gains in accuracy over batch normalization in federated settings.
Preprocessing CIFAR- 10 and CIFAR-100 consist of images with 3 channels of 32 32 pixels each. Each pixel is represented by an unsigned int8. As is standard with CIFAR datasets, we perform preprocessing on both training and test images. For training images, we perform a random crop to shape (24, 24, 3). followed by random horizontal flip. For testing images, we centrally crop the image (24, 24, 3). For both training and testing images, then normalize the pixel values according to their mean and standard deviation. Namely, given an image æ, we compute (æ p)/a where j the average of the pixel values in x, and the standard deviation.
C.2 EMNIST
EMNIST consists of images of digits and upper and lower case English characters, with 62 total classes. The federated version of EMNIST (Caldas al., 2018) () partitions the digits by their author. The dataset has natural heterogeneity stemming from the writing style of each person. We perform two distinct tasks EMNIST, autoencoder training (EMNIST AE) and character recognition (EMNIST CR). For EMNIST AE, we train the "MNIST" autoencoder (Zaheer et al., 2018). This is a densely connected autoencoder with layers of size (28 x 28) 1000 500- 250 30 and a symmetric decoder. full description network. The network has two convolutional layers (with 3 x 3 kernels), max pooling, and dropout, followed by 128 unit dense layer. full description of the model is in Table 4.
the model is in Table 3. For EMNIST CR, we use a convolutional
Table 3: EMNIST autoencoder model architecture. We use a sigmoid activation at all dense layers.
Layer
Output Shape
of Trainable Parameters
Input
784
0
Dense
1000
785000
Dense
500
500500
Dense
250
125250
Dense
30
7530
Dense
250
7750
Dense
500
125500
Dense
1000
501000
Dense
784
784784
25

---

Published as a conference paper at ICLR 2021
Table 4: EMNIST character recognition model architecture.
Layer
Ouiput Shape
Trainable Parameters
Activation
Hyperparameters
Input
(28, 28, 1)
0
Conv2d
(26, 26, 32)
320
kernel size strides=(1,
Conv2d
(24, 64)
18496
ReLU
kernel size 3; strides=(1,
MaxPool2d
(12, 12, 64)
0
pool size= (2, 2)
Dropout
(12, 12, 64)
0
p = 0.25
Flatten
9216
0
Dense
128
1179776
Dropout
128
0 7998
p 0.5
Dense
62
softmax
C.3 SHAKESPEARE
Shakespeare is language modeling dataset built from the collective works of William Shakespeare. In this dataset, each client corresponds to a speaking role with at least two lines. The dataset consists of 715 clients. Each client's 's lines are partitioned into training and test sets. Here, the task is to do next character prediction. We use an RNN that first takes a series of characters as input and embeds each of them into learned 8-dimensional space. The embedded characters are then passed through L.STM layers, each with 256 nodes, followed by densely connected softmax output layer. We split the lines of each speaking role into into sequences of 80 characters, padding if necessary. We use vocabulary size of 90; 86 for the characters contained the Shakespeare dataset, and extra characters for padding, out-of- off.ooobblllar. beginning of line and end of line tokens. We train our model take a sequence of 80 characters, and predict a sequence of 80 characters formed by shifting the input sequence by one (so that its last character is the new character we are actually trying to predict). Therefore, our output dimension is 80 90. A full description of the model is in Table
Table 5: Shakespeare model architecture.
Layer
Output Shape
# of Trainable Parameters
Input
80
Embedding
(80, 8)
720
LSTM
(80, 256)
271360
LSTM
(80, 256)
525312
Dense
(80, 90)
23130
C.4 STACK OVERFLOW
Stack Overflow is a language modeling dataset consisting of question and answers from the question and answer site, Stack Overflow. The questions and answers also have associated metadata, including tags. The dataset contains 342,477 unique users which we use as clients. We perform two tasks on this dataset: tag prediction via logistic regression (Stack Overflow LR, SO LR for short), and next-we word prediction (Stack Overflow NWP, SO NWP for short). For both tasks, we restrict to the 10,000 most frequently used words. For Stack Overflow LR, restrict to the 500 most frequent tags and adopt one-versus-res, classification strategy, where each question/answer is represented as a bag-of-words vector (normalized to have sum 1).
For Stack Overflow NWP, we restrict each client to the first 1000 sentences in their dataset (if they contain this many, otherwise we use the full dataset). We also perform padding and truncation ensure that sentences have 20 words. We then represent the sentence a sequence of indices corresponding to the 10,000 frequently used words, as well as indices representing padding, out-of- vocabulary words, beginning of sentence, and end of sentence. We perform next-word-prediette these sequences using an RNN that embeds each word in a sentence into a learned 96-dimensional space. It then feeds the embedded words into single LSTT layer of hidden dimension 670, followed
26

---

Published as a conference paper at ICLR 2021
by densely connected softmax output layer. A full description of the model is in Table 6. The metric used main body is the top-1 accuracy over the proper 10,000-word vocabulary; it does not include padding, out-of-vocab, or beginning or end of sentence tokens when computing the aceuracy.klkjlk
Table 6: Stack Overflow next word prediction model architecture.
Layer
Output Shape
# of Trainable Parameters
Input
20
Embedding
(20,90)
960384
LSTM
(20, 670)
2055560
Dense
(20,
64416
Dense
(20, 10004)
970388
D EXPERIMENT HYPERPARAMETERS
D.1 HYPERPARAMETER TUNING
Throughout our experiments, we compare the performance of different instantiations of Algorithm that use different server optimizers. We use SGD, SGD with momentum (denoted SGDM), ADAGRAD, ADAM, and YOGL. For the client optimizer, we use mini- batch SGD throughout. For all tasks, we tune the client learning rate 77 and server learning rate 7) by using large grid search. Full descriptions of the per-task server and client learning rate grids are given in Appendix D.2.
We use the version of FEDADAGRAD, FEDADAM, and FEDYOGT in Algorithm 5. We let 31 0 for FEDADAGRAD, and we let 31 = 0.9, 82 0.99 for FEDADAM, and FEDYOGI. For FEDAVG and FEDAVGM, we use Algorithm 4, where CLIENTÖPT, SERVERÔPT are SGD. For FEDAVGM, the server SGD optimizer uses a momentum parameter of 0.9. For FEDADAGRAD, FEDADAM, and FEDYOGI, we tune the parameter in Algorithm 5.
When tuning parameters, select the best hyperparameters (M 7), r) based the average training loss over the last 100 communication rounds. Note that at each round, we only see fraction of the total users (10 for each task except Stack Overflow NWP, which 50). Thus, the training loss at given round a noisy estimate of the population- level training loss, which why we averave over a window of communication rounds.
D.2 HYPERPARAMETEE GRIDS
Below, we give the client learning rate (1) in Algorithm 1) and server learning rate (n in Algorithm 1) grids used for each task. These grids were chosen based on an initial evaluation over the grids
M/| E (10 3,00 10.00-0 2.5
100.3
{10-3 3,
10 2.5 2
10'}
These grids were then refined for Stack Overflow LR and EMNIST AE in attempt to ensure that the best client/server learning rate combinations for each optimizer was contained in the interior of the learning rate grids. We generally found that these two tasks required searching larger learning rates than the other two tasks. The final grids were follows:
CIFAR-10:
W E {10
10
100.5)
{10-3
10 2.5
10
CIFAR-100:
E {10-3 10
10
(10-3 10-20
10-}
27

---

Published as a conference paper at ICLR 2021
EMNIST AE:
" E {10-1.5 1001
102)
E {10-?, 2 10 1.5
10
EMNIST CR:
E {10-3, 3 10 10020-.. 2.5
100.5
7 {10-3 10 2.5
10²}
Shakespeare:
" (10 10 2.5
10
7 {10
10
10
StackOverflow LR:
"h E {10-1, 10
103}
7 E {10-2, 10 1.5
10
StackOverflow NWP:
E {10 10 2.5
10
E {10- 10 2.5
10²}
For all tasks, we tune over the grid:
7 E (10-5 ..... 10 10-1)
D.3 PER-TASK BATCH SIZES
Given the large number of hyperparameters to tune, and to avoid conflating variables, we the batch size at per-task level. When comparing centralized training to federated training in Section 5, we use the same batch size in both federated and centralized training. full summary of the batch sizes is given in Table 7.
Table 7: Client batch sizes used for each task.
TASK
BATCH SIZE
CIFAR-10
20
CIFAR-100
20
EMNIST AE
20
EMNIST CR
20
SHAKESPEARE
4
STACKOVERFLOW LR
100
STACKOVERFLOW NWP
16
D.4 BEST PERFORMING HYPERPARAMETERS
In this section, we present, for each optimizer. the best client and server learning rates and values of r found for the tasks discussed in Section 5. Specifically, these are the hyperparameters used in Figure Figure and table Table 1. The validation metrics in Table I obtained using the learning rates 'aable 8 and the values of r in Table 9. As discussed in Section 4, we choose 7), !h and to be the parameters that minimizer the average training loss over the last 100 communication rounds.
28

---

-5 Published as a conference paper at ICLR 2021
Table 8: The base-10 10 logarithm the client (m) and server (7) learning rate combinations that achieve the accuracies from Table 1. See Appendix D.2 for full description of the grids.
FEDADAGRAD
FEDADAM
FEDYOGI
FEDAVGM
FEDAVG
T
T)
YR
7
Y
1
YR
1
CIFAR-10
%/2
-/2
-2
-
-2
-//
1/2
1/2
CIFAR-100
-1
-1
-3/2
0
-
0
-//2
0
-1
1/2
EMNIST AE
1/2
-3/2
-
-
1/2
0
EMNIST CR
-//
-1
-3/2
-
-
-
-
-//
-1
SHAKESPEARE
0
-1/2
0
-2
-2
0
-1/
0
0
STACKOVERFLOOW LR
2
-1/
-1/
0
STACKOVERFLOW NWP
1/2
-//
-1/2
-2
-//
-2
-//
0
-1/
Table 9: The base- 10 logarithm of the parameter (as defined in Algorithm 2) that achieve the validation metrics in Table 1.
FEDADAGRAD
FeDADAM FEDYOGI
CIFAR-10
-2
CIFAR-100
-2
EMNIST AE
-3
EMNIST CR
-2
-4
.4
SHAKESPEARE
-1
STACKOVERFLOW LR
-2
-5
STACKOVERFLOW NWP
-4
29

---

Published as a conference paper at ICLR 2021
EMNIST CR
0.9
Accuracy
0.8-
0.7
500
1000
1500
Number of Rounds
Figure 4: Validation accuracy on EMNIST CR using constant learning rates 7,, Yh and T tuned to achieve the best training performance on the last 100 communication rounds; see Appendix D for hyperparameter grids.
Table 10: Test set performance for Stack Overflow tasks after training: Accuracy for NWP and Recall @5 (x x000 for LR. Performance within within 0.5% of the best result are shown in bold.
FED...
ADAGRAD
ADAM
YOGI
AvGM
AvG
STACK OVERFLOW NWP
24.4
25.7
25.7
24.5
20.5
STACK OVERELOW LR
66.8
65.2
66.5
46.5
40.6
E ADDITIONAL EXPERIMENTAL RESULTS
E. RESULTS ON EMNIST CR
We plot the validation accuracy of FEDADAGRAD, FEDADAM, FEDYOGI, FEDAvGM, FEDAVG, and SCAFFOLD on EMNIST CR. As in Figure 1, we tune the learning rates m, 1) and adaptivity by selecting the parameters obtaining the smallest training loss, averaged over the last 100 training rounds. The results are given in Figure 4.
We see that all methods are roughly comparably throughout the duration training. This reflects the fact that the dataset is quite simple, and most clients have all classes in their local datasets, reducing any heterogeneity among classes. Note that SCAFFOLD performs slightly worse than FEDAVG and all other methods here. As discussed in Section 5, this may be due to the presence of stale client control variates, and the communication -imiiiiiti----eeeedddeeeeeeee regime of our experiments.
E.2 STACK OVERFLOW TEST SET PERFORMANCE
As discussed in Section 5, in order to compute measure of performance for the Stack Overflow tasks as training progresses, we use subsampled version of the test dataset, due to its prohibitively large number of clients and examples. In particular, at each round of training, we sample 10,000 random test samples, and use this completed, we also evaluate on the full test dataset. For the Stack Overflow experiments described in Section 5, the final test accuracy is given in Table 10.
measure of performance over time. However, training is
E.3 LEARNING RATE ROBUSTNESS
In this section, we showcase what combinations of client learning rate ? and server learning rate performed well for each optimizer and task. As in Figure plot, for each optimizer, task, and pair (m)7)) the validation set performance (averaged over the last 100 rounds). As in Section 5, we fix 10-3 throughout. The results, for the CIFAR-10, CIFAR-100 100, EMNIST AE, EMNIST CR,
30

---

Published as a conference paper at ICLR 2021
Shakespeare, Stack Overflow LR, and Stack Overflow NWP tasks given in Figures 5, 6, 7, 7, 9, 10, and 11, respectively.
While the general trends depend on the optimizer and task, we see that in many cases, the adaptive methods have rectangular regions that perform well. This implies a kind of robustness to fixing one of 1/, 7tt and varying the other. On the other hand, FEDAVGM and FEDAVG often have triangular regions that perform well, suggesting that 7) and NN should be tuned simultancously.
Fessiagaad
CIFARLO, FodAdam
OIFRRDD Fedrog
1.4
363
361
TME SE4 30E 361 351
12.
1.0
203
350
1.0
263
504
204
3.5
153 163 361 361 35X
70
0.5
300
363 353 194
70
0.5
12A
12.
153
70
2.4
T34
314
60
2.0
364 351 155
0.0
364
354 151
180
Learr maiiing Date
10.5
50
+0.5
302
31A
532
52.
50
1-0.5
11a
302
212
3.0
40
-1.4
114
112
13.
40
-1.4
114
40
-1.
113
12.0
144
414
.2.
122
122
100
71
114
514
100
712
71
10
324
3x4
-0.4
618
364
112
614
332
-1.4
413
524
K13
612
812
Glontaa Laa.mmmmm Rata (00g))
Clent Loo
mg
Rate 90gus)
Clenn Lo0 Loaminn Rate 110gu)
CIFARLD, FesAugh
FedAvg
1.0
132
1.0
70
70
302
132
91
621
9.
912
234
75
60
0.0
512
635
714
60
112
224
$2.0-21
214
07
41
40
33.5 144
111
30
433
41 40
22.0 313
343 213 284 251 811
20
203
10
22.5
545
318 118 150 164
10
3.0.:51 131 315
362 317 181 509
3300
KX 154 181
Clend Lewning Race Oeggal
Clen Learning Rate Orggal
Figure 5: Validation accuracy (averaged over the last 100 rounds) of FEDADAGRAD, FEDADAM, FEDYOGE FEDAVGM, and FEDAVG for various client/server learning rates combination on the CIFAR- 10 task. For FEDADAGRAD, FEDADAM, and FEDYOGI, we set 10
31

---

Published as a conference paper at ICLR 2021
CIMAR-005, Teskdegred
G0
CIPAR-TOE, TedAdem
60
GTAR-100 TedTog
60
Cenvor Leamiing Rate Wogee
Genver Leuming Rate (kgre)
Senvel Lauming Rate: Wogue
40
40
40
20
20
20
152
114
114
154
192
154
178
214
138 399
3.0 2.5
2.0
15
1.0
0.5
0.0 0.5
3.0
2.5
2.8
15
1.0 0.00000 0.5
2.5.2220 3.0
2.0
1.5 0.5.0. 0.00000:000
(gua)
Regual
CIFAR-100, FedAugM
60
CIFAR- Fedtvg
60
1.0
1.0
50
40
40
3.0
30
20
20
3.0
200 2. 2.015-111 10. 4.5 0.0 0.5
3.0
25
1.0 2.5
0.9 0.5
ring 00gual
00gual
Figure 6: Validation accuracy (averaged over the last 100 rounds) of FEDADAGRAD, FEDADAM, FEDYOGI, FEDAVGM, and FEDAVG for various client/server learning rates combination on the CIFAR- 100 task. For FEDADAGRAD, FEDADAM, and FEDYOGI, we set T = 10-3
Servel Lariing hata (bouu)
SUNIST AE Fedddograd
Server Laursing Rote lbgue
PNNST AE Reskdo
Server Lauring Ruta (bgue)
ENN AT Fedrogi
R00 20 2 (41000
1.3
1.3
1.
Client Lear-ing singiing Nate
Client ninging Nate
Clent Learning Hala
BTNSTT AE: Fediuggt
EMNIST AE Feddvg
Server Server Looming Rate Wigul
2.5
RMte
2.0
Soweer Losmingg
1.5
1.5
1.0
1.0
17
1.0
57
07
Clent Learring Rate
Clantt Learring Rate
Figure 7: Validation MSE (averaged over the last 100 rounds) of FEDADAGRAD, FEDADAM, FEDYOGI, FEDAVGM, and FEDAVG for various client/server learning rates combination on the EMNIST AE task. For FEDADAGRAD, FEDADAM, and FEDYOGI, we set T = 10
32

---

Published as a conference paper at ICLR 2021
EMNNSTOR CR. FosAdagrad
90
ENNSST CR. FodAsom
90
EMNST CR. Fodtogl
90
1.0
2.0-44
1.0
80
0.5-41
80
Laumingg hot Dooue
80
70
70
70
0.0.1
0.0:42
0.0
Senver Learning Rate
60
Learnin Rat
60
50
50
50
40
40
40
FLC
AL4 WD 1LT
1.5
MLN M.N
30
30
30
20
20
20
10
10
10
650
764
N. 0.3
exx 200
10-2312.--15-11000 (9gua)
0.0
-3.0-2.522
-3.0-2.92 10-2.--2--1..1
(logua)
HINETT ClI, TedAugM
50
ENNIST CV, hestig
50
1.0
1.0
80
80
70
70
60
60
50
50
3.0
40
40
30
30
20
20
10
10
3.0-22
-30-25--2--1.--00--
05
-3.02 300-2--2.------
0.0 0.
Oogual
Figure 8: Validation accuracy (averaged over the last 100 rounds) of FEDADAGRAD, FEDADAM, YEDYOGI, FEDAVGM, and FEDAVG for various client/server learning rates combination the EMNIST CR task. For FEDADAGRAR AD, FEDADAM, and FEDYOGI, we set r = 10-3
Sha
Fedicdegrad
60
Shacoeparr PedAdam
60
Shacespperee PedYoggi
1.0
1.0
1.0
Cercua Letrring Rate Loone
21A
50
Sencer Letring Rate (o0mel
50
Seree Leariing Rate loggel
50
0.0
0.0
0.5
3.0
550
515
30
1.0
21
04
30
1.0
30
563
312
345
535
512
342
34
564
334
514
314
35
312
12.0
34
312
222
10
39
10
36.7
382
10
3.0
314
314
314
314
304
3.0
02
61
11a
514
3.0
114
02
113
P7
3..2.5.20
Clans LearLemmnng
Rate
Clert Leeming flate Ougua)
Leeriing Rati (logs))
Shakesseare, Fedhoggh
Shakespeare, Fedhug
1.0
1.0
21-
2v2
123
50
50
0.0
13
0.0
Suever Some S
14
213
1..0
$2.0
17
Ar
2.0
2.0
90
31
2.55
50
10
-3.0
60
12
54
-3.0
83
80
2.0 1.0115:110 0.000000000
05
2.0
1.0
0.0
0.
Figure 9: Validation accuracy (averaged over the last 100 rounds) of FEDADAGRAD, FEDADAM, FEDYOGI, FEDAVGM, and FEDAVG for various client/server learning rates combination on the Shakespeare task. For FEDADAGRAD, FEDADAM, and FEDYOGI, we set T = 10-3
33

---

Published as a conference paper at ICLR 2021
Overlaw Foyldayyad
Fededom
Fedroog
Learsin Baae Cogyel
1052-4 SEGPPN
50
100 0 SANPE
6.0
6.0
M4
57 944
6.0
(00) 0
314
312 313 26:
312
312
-2.0
212
312
360 357 211
-2.0
302
301
M1
-2.0 15 365
302
233
3.00000000
2.5
-300000000
3.5 2.0
-1.000500 3.0
3.0
CRem Leaming
Rato (19g))
Cleenteam mg Race (igg.)
Client Leamiin Roce (ogos)
Stack Overfoow LX, Foolaggh
Stack Overflow Lh. FolAvg 141
1.5 SAt we aen 142 141
50
Ratee Pihee
3.5
3,0
60
50
100 N0
50
184 142
224
15 (31 1K1
083 184 547 180
182 122 180
Cler Leaming Pate 30ges)
CAAe Leaming Rate S0gts)
Figure 10: Validation recall @5 (averaged over the last 100 rounds) of FEDADAGRAD, FEDADAM, FEDYOGI, FEDAVGM, and FEDAVG for various client/server learning rates combination on the Stack Overflow LR task. For FEDADAGRAD, FEDADAM, and FEDYOGI, we set 10-3
Stack Overfow MNP. FedAdngrad
30
Steck Overlow KWP, TedMidm
30
Stace Overfow MNP, Tedfrogi
30
Server Lesming Sate (esgrs
Gewe LosmmPa Bato (Wgus!
20
Server Lesming Sate
20
20
10
10
10
30 23 2.0.15
1.0 0'5 0.0 0.5
30.25-22..15
10.00055000 mg Ongps)
3.0.2.5200 11225-20:1133 1001111000000 10.05:000
(egga)
(Inga))
Stack Overfow NVP. FedRigM
10
Stock Overfoow NAP, Foddup
10
Balte
20
20
130-25-20--15- 0.0 05
DD 05
Clent (oggul
Clent Leariin Oogual
Figure 11: Validation accuracy (averaged over the last 100 rounds) of FEDADAGRAD, FEDADAM, FEDYOGL, FEDAVGM, and FEDAVG for various client/server learning rates combination on the Stack Overflow NWP task. For FEDADAGRAD, FEDADAM, and FEDYOGI, we set T = 10-3
34

---

Published as a conference paper at ICLR 2021
E.4 ON THE RELATION BETWEEN CLIENT AND SERVER LEARNING RATES
In order to better understand the results in Appendix E.3, we plot the relation between optimal choices of client and server learning rates. For each optimizer, task, and client learning rate T], we find the best corresponding server learning rate 1) among the grids listed in Appendix D.2. Throughout, we fix r 10-3 3 for the adaptive methods. We omit any points for which the final validation loss is within 10% of the worst-recorded validation loss over all hyperparameters Essentially, we omit client learning rates that did not lead any training of the model. The results are given in Figure 12.
CIFAR-10
CIFAR-100
EMNIST AC
Server Leaming Rato
Sarver Leaming Rate
10'
10"
10"
10-1
10-1
10-1
Server
20.
10
10-2
10
10-2
10
10
16'
Client Learning
Learning Rate
Learning
EMNIST
Shakespeare
Stack
10
10
Servert Leaning Rake
Servart Lesmiing Rata
100
Leaning
103
10a
10.
10-2
10
10
10
10
10
10
Client Learning
Client Learning Rate
Client Learning Rate
Stack Overfiow NWP
Server Learning Rate
10a
100
Podl-jom
Fodldoogoo
Fedriggi
10-3
PedAugM
10-2
10a3
10
10
Client Learning Rate
Figure 12: The best server learning rate in our hyperparameter tuning grids for each client learning rate, optimizers, and task. We select the server learning rates based on the average validation performance over the last 100 communication rounds. For FEDADAGRAD, FEDADAM, and FEDYOGI, we fix 10-3, We omit all client learning rates for which all server learning rates did not change the initial validation loss by more than 10%.
In virtually all tasks, learning rate y for FEDAVG and FEDAVGM. As discussed in Section 5, this supports the observation that for the non-adaptive methods, 7) and 17 must in some sense be tuned simultancously. On the other hand, for adaptive optimizers on most tasks we see much more stability in the best server learning rate 77 the client learning rate T varies. This supports observation in Section 5 that for adaptive methods, tuning the client learning rate is more important.
clear inverse relationship between client learning rate 7) and server
Notably, we see clear exception to this on the Stack Overflow LR task, where there is definitive inverse relationship between learning rates among all optimizers. The EMNIST AE task also displays somewhat different behavior. While there are still noticeable inverse relationships between learning rates for FEDAVG and FEDAVGM, the range of good client learning rates is relatively small. We emphasize that this task is fundamentally different than the remaining tasks. As noted by Zaheer et al. (2018), the primary obstacle in training bottleneck autoencoders is escaping saddle points, not in converging to critical points. Thus, we expect there to be qualitative differences between EMNIST AE and other tasks, even EMNIST CR (which uses the same dataset).
35

---

Published as a conference paper at ICLR 2021
E.5 ROBUSTNESS OF THE ADAPTIVITY PARAMETER
As discussed in Section 5, we plot, for each adaptive optimizer and task, the validation accuracy as a function of the adaptivity parameter т. In particular, for each value of 7 (which we vary over (10-5, ..... 10-1), see Appendix D), we plot the best possible last- 100--ouosss validation set performance. Specifically, we plot the average validation performance over the last 100 rounds using the best a posteriori values of client and server learning rates. The results are given in Figure 13.
0.80
CIF AR-10
CIFAR-100
0.020
EMNIST AE
Mean Squared Eror
Acourdcy
0.75
Accuraly
0.5
0.015
0.70
0.010
0.4
0.65
0.005
0.60
10-3 Adaptiviiy (t)
0.3
10"
0.000
10:3
10-
Adaptivity (t)
Andaptivity It!
0.30
FMNIST
0.60
Shakesseare
0.70
Stack Overflow
0.85
0.58
Recaligs
0.65
Acoracy
0.56
0.80
0.60
0.54
0.75
0.52
reshangred
0.55
0.70
10
10-1
10-3
10-h
10
0.50 10
10-
10-3
10-h
10-
0.50
10
10
10-3
10-h
10-
Adaptivity
Adaptiityy (t)
0.30
Stack Overflow NWP
0.25
0.20
0.15
0.10
10
10th
10-3
10-4
10
Adaptivity (rl
Figure 13: Validation performance of FEDADAGRAD, FEDADAM, and FEDYOGT for varying on various tasks. The learning rates v) and 1): are tuned for each to achieve the best training performance on the last 100 communication rounds.
E.6 IMPROVING PERFORMANCE WITH LEARNING RATE DECAY
Despite the success of adaptive methods, it is natural to ask if there is still more to be gained. To test this, we trained the EMNIST CR model in a centralized fashion on a shuffled version of the dataset. We trained for 100 epochs and used tuned learning rates for each (centralized) optimizer, achieving an accuracy of 88% (see Table 11, CENTRALIZED row), significantly above the best EMNIST CR results from Table 1. The theoretical results in Section 3 point to partial explanation, as they only hold when the client learning rate is small or is decayed over time.
To validate this, we ran the same hyperparameter grids on the federated EMNIST CR task, but using a client learning rate schedule. We use a "staircase" exponential decay schedule (ExpDecAy) where the client learning rate π) is decreased by a factor of 0.1 every 500 rounds. This is analogous in some sense to standard staircase learning rate schedules in centralized optimization (Goyal et al., 2017). Table 11 gives the results. EXPDECAY improves the accuracy of all optimizers, and allows most to get close the best centralized accuracy. While we do not close the gap with centralized optimization entirely, we suspect that further tuning of the amount and frequency of decay may lead to even better accuracies. However, this may also require performing significantly more communication rounds, as the theoretical results in Section 3 are primarily asymptotic. In communication-lininn settings, the added benefit of learning rate decay seems to be modest.
36

---

Published as a conference paper at ICLR 2021
Table 11: (Top) Test accuracy (%) of a model trained centrally with various optimizers. (Bottom) Average test accuracy (%) over the last 100 rounds of various federated optimizers on the EMNIST CR task, using constant learning rates or the EXPDECAY schedule for mn- Accuracies (for the federated tasks) within 0.5% of the best result are shown in bold.
ADAGRAD
ADAM 87.9
YOGI 88.0
SGDM
SGD
CENTRALIZED 88.0
87.7
87.7
FED... ADAGRAD
ADAM
YOGI
AvGM
AvG
CONSTANT Tj EXPDECAY
85.1
85.6
85.5
85.2
84.9
85.3
86.2
86.2
85.8
85.2
F CREATING A FEDERATED CIFAR-100
Overview We use the Pachinko Allocation Method (PAM) (Li & McCallum, 2006) to create federated CIFAR-100 100. PAM is topic modeling method in which the correlations between individual words in a vocabulary are represented by rooted directed acyclic graph (DAG) whose leaves are the vocabulary words. The interior nodes are topics with Dirichlet distributions over their child nodes. To generate document, we sample multinomial distributions from each interior node's Dirichlet distribution. To sample a word from the document, we begin at the root, and draw a child node its multinomial distribution, and continue doing this until we reach a leaf node.
To partition CIFAR-100 across clients, we use the label structure of CIFAR-106 Each image in the dataset has a fine label (often referred to as its label) which is a member of a coarse label. For example, the fine label "seal" is member of the coarse label "aquatic mammals". There are 20 coarse labels in CIFAR-100, each with 5 fine labels. We represent this structure as a DAG G. with a root whose children are the coarse labels. Each coarse label is an interior node whose child nodes its fine labels. The root node has a symmetric Dirichlet distribution with parameter a over the coarse labels, and each coarse label has symmetric Dirichlet distribution with parameter B3
We associate each client to a document. That is, we draw a multinomial distribution from the Dirichlet prior at the root (Dir(a:) and multinomial distribution from the Dirichlet prior at each coarse label (Dir(3)). To create the client dataset, we draw leaf nodes from this DAG using Pachinko allocation, randomly sample an example with the given fine label, and assign it to the client's 's dataset. We do this 100 times for each of 500 distinct training clients.
While more complex than LDA, this approach creates more realistic heterogeneity among client datasets by creating correlations between label frequencies for fine labels within the same coarse label set. Intuitively, if a client's 's dataset has many images of dolphins, they are likely to also have pictures of whales. By using a small few coarse labels. By using a larger /3 for the coarse-to-fine label distributions, clients are more likely to have multiple fine labels from the same coarse label.
the root, client datasets become more likely to focus
One important note: Once we sample a fine label, we randomly select an element with that label without replacement. This ensures no two clients have overlapping examples. In more detail, suppose we have sample a fine label with coarse label for client m,, and there is only one remaining such example (z, y). We assign (I. e, y)) to client m''s dataset, and remove the leaf node y from the DAG G. We also remove y from the multinomial distribution O. that client m has associated to coarse label which we refer to renormalization with respect to y (Algorithm 8). If i has no remaining children after pruning node j. we also remove node from G and re-normalize the root multinomial 0, with respect to c. For all subsequent clients, we draw multinomials from this pruned G according to symmetric Dirichlet distributions with the same parameters as before, but with one fewer category.
Let C denote the set of coarse labels and y the set of fine labels, and let S denote the CIFAR- dataset. This consists of examples (x, c, y) where x an image vector, c is a coarse label set, and y E y is a fine label with E c. For e E C, E ), we let S. and S, denote the set of examples in S with coarse label e and fine label y. For E G we let [G[u] denote the set of children of v in G. For 7 E R, let Dir(7; k) denote the symmetric Dirichlet distribution with categories.
Notation and method
37

---

Published as a conference paper at ICLR 2021
Algorithm 7 Creating federated CIFAR-100
Input: N, .M -o, ax, 5 Rg0
for m = =, .M do
Sample 0, Dir(a, [G[r])
for c C G[r] do
Sample 0g ~ Dir(B, [G[c])
Dn =
for 1, ................................ N do
Sample (ultinomial(0.))
Sample Multinomial(E3)
Select (I, c, y) E uniformly at random
Dm Dm U cxr,
S S (3r(I,,,))
if S then
G Gi(E33
0, RENORMALIZE((0..)
if S. @ then
G = G'Vcc
0, RENORMALIZZE(O,,C)
Algorithm 8 RENORMALIZE
Initialization: = (Pi...., PK), [K]
En Pk: for E [K].k lt + do P/. = Pk/a
Return 0' (p1, .... 1/ .PK)
Let M denote the number of clients, N the number of examples per client, and D, the dataset for client m E {1, ... .M. A full description of our method is given in Algorithm 7. For our experiments, use N 100, MM 500, 0.1, 10. In Figure 14, we plot the distribution of unique labels among the 500 training clients. Each client has only a fraction of the overall labels in the distribution. Moreover, there variance in the number of unique labels, with most clients having between 20 and 30, and some having over 40. Some client datasets have very few unique labels. While this is primarily heterogeneity of the dataset in a way that can reflect practical concerns, as in many settings, clients may only have few types of labels in their dataset.
artifact of performing without replacement sampling, this helps increase the
CIFAR-100 Client Label Distribution
Frequency
50
40
30
20
10
10
20
30
40
Number of unique labels
Figure 14: The distribution of the number of unique labels among training client datasets in our federated CIFAR-100 dataset.
38


Reviews:
[{"review_id": "pOkj6QSTCjJ", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: This paper presents an adaptive federated optimization framework that induces three different adaptive federated learning algorithms, which are proposed to address the issues of client drift due to data heterogeneity and lack of adaptivity. The authors presented thorough literature survey on the federated learning and formulated the meta-algorithm, FedOpt, introducing both server and client optimizers. Different from the popular FedAvg, the server optimizer in this work is an adaptive protocol which is originated from the client drift. The authors analyzed mathematically the convergence rates of the proposed framework and showed extensive experimental results on different benchmark datasets to validate the efficacy of the developed algorithms.\n\nOverall, this paper is well organized and technically sound. It is also easy to follow. The overall proof ideas are also correct, though I didn’t derive step by step. However, the authors need to pay attention to the following points to improve the current draft.\n\n1. Server optimizer can be confusing. The server optimizer developed in this study generalizes to a more adaptive one. But it is not really an optimization in the server, instead just another form of local model averaging in order to obtain the global model. I completely understand the motivation for this is to improve the performance beyond FedAvg, which may not necessarily perform well with non-IID data. However, without using zeroth order or first order information of objectives, to me, the so-called server optimizer is just a nonlinear local model averaging.\n2. The server learning rate is bit counterintuitive. In Corollary 1 and 2, when defining the learning rate for the server, it looks like $\\eta$ can be quite large. Of course, given the fact that the server optimizer is not really an optimization, that could be understandable. Additionally, the authors empirically investigated the relationship between the server and client learning rates in the appendix. But to me, it is naturally a parameter that plays a similar role as learning rate, not exactly the learning rate. \n3. The communication efficiency needs more discussion. In Section 3, the authors presented one discussion on the communication efficiency, particularly quantifying $K$. I wonder how the authors obtained that. Also, it would be great to see more quantitative results on the communication efficiency. \n4. Client heterogeneity is only qualitative in the paper. Based on the empirical implementation, I couldn’t find out how the authors simulated the non-IID data distributions for different clients for each benchmark task. This needs more detail. Also, the authors mentioned that the proposed framework can work effectively for moderate and naturally arising heterogeneity. But under how much heterogeneity does the algorithms work well? For example, for CIFAR 10, if having 10 clients, each client has only one class without any overlapping, is FedOpt still effective?\n5. The experimental results look promising, but not showing significant outperforming capability. From Table 1, we can still observe that for some tasks, the AvgM is still favorably comparable. Also, when to select which method among FedAdam, FedAdaGrad, and FedYOGI is unclear. The authors need to give more discussion.\n6. Based on the Corollary 1 and 2, it looks like when $T$ is sufficiently large, the FedOpt can achieve linear speed up, correct? However, the authors failed to give the specific lower bound for $T$, even if they only mentioned the dominating term that induced a sublinear rate.\n\n***********************************\nAfter carefully considering the rebuttal from the authors, I think I am more positive about the paper so I raised my score. The rebuttal clarifies most of my confusion about the paper, though more improvement can be done. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "dXlZgiWlRFs", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes several federated variants of adaptive stochastic gradient methods.   Moreover,  the convergence rates of the proposed algorithms are also provided.  Based on the current submission,  The reviewer has several concerns: \n\n1. The adaptive learning rates in FedAdagrad/FedYogi/FedAdam are all related to parameter $\\tau$.  According to Theorem 1-2, Corollary 1-2,   $\\tau$ needs to set as a large constant $G/L$.   When  $\\tau$ is sufficiently large, FedAdagrad/FedYogi/FedAdam is actually reduced to FedSGD. Can the authors give several comments on this parameter? \n\n2. The second question is about the generalization ability of federated adaptive stochastic gradient methods.   It has been demonstrated in [1] that adaptive SGD generalizes poorly compared with SGD.  A natural question is that \"  does this dilemma still exists in federated stochastic gradient methods? \" \n\n[1] Wilson AC, Roelofs R, Stern M, Srebro N, Recht B. The marginal value of adaptive gradient methods in machine learning. In Advances in neural information processing systems 2017 (pp. 4148-4158).", "extra_scores": {}, "extra_text": {}}, {"review_id": "Mloyim_mx1Z", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper studies the convergence of well-known adaptive methods, ADAM, ADAGRAD, and YOGI, for the federated learning problem. In particular, while the nodes (clients) still use SGD for their local computations (same as Fed-Avg), the server uses one of the three adaptive methods mentioned above to update the model. The authors have provided convergence rates (based on gradient norms) for all three methods for the nonconvex settings. Moreover, various experiments have been conducted to compare the strengths of these methods against classic FedAvg. Overall, this paper is well-written, and the problem statement and the goal of the paper are clear. Moreover, the authors have shown the success of using adaptive methods in numerical settings.\n\nMy main questions are regarding the theoretical analysis of the paper. First, the analysis of adaptive methods in the nonconvex setting has been studied in the literature (as mentioned by the authors as well in the related work). Given the bounded gradient assumption imposed in this paper, I am not sure what would be the main challenge in extending those analyses to the federated setting. I encourage the authors to highlight possible novelties in their analysis compared to prior works on adaptive methods for nonconvex objective functions.   \n\n\nSecond, it is not clear to me how this paper's theoretical results would imply that adaptive methods enjoy better convergence rates compared to regular Fed-Avg. In particular, all the algorithms achieve the rate of $O(1/\\sqrt{T})$, and do not differ from each other in that aspect. To better highlight my point, for instance, let's consider the convex setting. There, Duchi et al. show that while AdaGrad has $O(1/\\sqrt{T})$ convergence rate in the worst case (same as SGD), it can achieve better rates in certain scenarios, such as gradients being sparse over coordinates, etc. Hence, this way, the theoretical result shows the advantage of using adaptive methods over the standard SGD method. From my understanding, this paper does not provide that kind of results, and so one important question remains unanswered: Why switching to adaptive methods after all? I would appreciate it if the authors provide further discussions on this matter.\n\n-----\n\nPost rebuttal comment: I appreciate the authors' responses, especially on highlighting the theoretical challenges. I have raised my score accordingly.", "extra_scores": {}, "extra_text": {}}, {"review_id": "JP2YOWdjHj", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper extends the server model averaging step in FedAvg to a more general adaptive optimization step on the global model, specifically, by writing the model averaging as a gradient descent step using a pseudo gradient. Three variants of this scheme (FedOpt) are presented, based on three adaptive optimizers, including AdaGrad, ADAM, and YOGI. While there exist works applying the server-side momentum method, the paper argues that the proposed framework is more general since any adaptive optimizer can be applied to extend FedAvg. Under three assumptions (Lipschitz gradient, bounded gradients, bounded variances), the paper provides a local convergence analysis with nonconvex objectives. The achieved bounds match the best-known convergence rate for FL under reasonable conditions (i.e., T is sufficiently large compared to K). They also provide guidance for how to decay the local learning rate to avoid client drift, and show that increasing local steps can help to reduce the communication rounds (under certain conditions). They also reflect the dependency of these bounds on client heterogeneity measured by global variance and indicates how to mitigate the problem in FedOpt. Experimental comparison to FedAvg, FedAvgM, and SCAFFOLD on seven benchmark FL tasks show that the proposed three adaptive FL methods are better or comparable on early-stage convergence and final validation-set performance. A study about tuning the local and global learning rate is also presented.\n\nPros:\n\n(1) The paper is well written and easier to understand for most parts. The authors moved the proof and experiment details to the appendix, which largely improves the readability. \n\n(2) The proposed FedOpt, as an extension of FedAvg, can incorporate adaptive optimizers for the server-side model update, and thus provide a general framework that can potentially cover a family of FL algorithms, which can be an important contribution to the FL community. \n\n(3) The interpretation of the bounds is also informative: showing several expected or preferred trade-offs in FL and sheds light on how to reduce the communication costs and the negative effects of client heterogeneity. \n\n(4) The experiments show the compelling performance of the proposed methods, especially on Stack Overflow LR and EMNIST AE, on which the improvements over the other two baselines are remarkable.\n\nCons and suggestions:\n\n(1) One main novelty and focus of this paper is the server-side adaptive optimization, and the purpose is to combat client drift and reduce the variance of stochastic gradient, especially in the cross-device FL setting. However, is it the case that improving client-side optimization can be more effective in resolving these two issues? FedOpt covers the client-side optimization in its framework but did not expand it on possible options, while other works like FedProx focus more on it. It would be helpful to provide a discussion and comparison of the two strategies and their advantages/disadvantages on the two issues.\n\n(2) The theoretical bounds show the trade-off between T and K, the client and server learning rates, and the client heterogeneity. In experiments, only different options of client and server learning rates have been studied. It would be more sound to show the other trade-offs indicated by the bounds.\n\n(3) In experiments, FedAvgM is a very competitive baseline since on most tasks, its performance is very close to the proposed methods, except on SO LR and EMNIST AE. Considering its similarity in methodology with FedOpt, it reduces the novelty of the paper.\n\n(4) There are only three FL baselines that have been compared in the experiments. Is there a specific reason for not including other FL methods into the comparison? In addition, it is interesting to compare with (i) FedOpt with both adaptive client optimizer (or FedProx) and adaptive server optimizer and (2) FedOpt with only adaptive client optimizer.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes adaptive optimization algorithms for federated learning that are federated versions of existing adaptive algorithms such as Adam, Adagrad, and Yogi. The paper establishes convergence guarantees for the proposed algorithms and performs an extensive experimental evaluation. Following the discussion, the reviewers were positive about the paper and felt that the author responses addressed their concerns. I recommend accept.
Decision: Accept (Poster)