Mime: Mimicking Centralized Stochastic Algorithms in Federated Learning
Authors: Sai Praneeth Karimireddy
Martin Jaggi
Satyen Kale
Mehryar Mohri
Sashank J. Reddi
Sebastian U Stich
Ananda Theertha Suresh
Venue:  (2021)
Keywords: Federated learning
Federated optimization
Adaptive optimization
Adam
Variance Reduction
Distributed optimization
Decentralized optimization

Abstract: Federated learning is a challenging optimization problem due to the heterogeneity of the data across different clients. Such heterogeneity has been observed to induce \emph{client drift} and significantly degrade the performance of algorithms designed for this setting.
In contrast, centralized learning with centrally collected data does not experience such drift, and has seen great empirical and theoretical progress with innovations such as momentum, adaptivity, etc.
In this work, we propose a general framework {\sc Mime} which mitigates client-drift and adapts arbitrary centralized optimization algorithms (e.g. SGD, Adam, etc.) to federated learning.
{\sc Mime} uses a combination of \emph{control-variates} and \emph{server-level statistics} (e.g. momentum) at every client-update step to ensure that each local update mimics that of the centralized method. Our thorough theoretical and empirical analyses strongly establish \mime's superiority over other baselines.

Under review as a conference paper at ICLR 2021
MIME: MIMICKING CENTRALIZED STOCHASTIC AL- GORITHMS IN FEDERATED LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
Federated learning (FL) is challenging setting for optimization due the het- erogeneity of the data across different clients. This heterogeneity has been shown cause client drift, which can significantly degrade the performance of algo- rithms designed for the FL setting, In contrast, centralized learning with centrally collected data not affected by such drift and has seen great empirical and theo- retical progress with innovations such as momentum and adaptivity. In this work, we propose general algorithmic framework, MIME, which mitigates client drift and adapts arbitrary centralized optimization algorithms such as SGD) and Adam to the federated learning setting. MIME uses a combination of control-variates and server-level statistics (e.g. momentum) at every client-update step to ensure that each local update mimics that the centralized method run on iid data. Our thorough theoretical and empirical analyses strongly establish MIME'S superiority over other baselines.
INTRODUCTION
Federated learning has become an important paradigm in large-scale machine learning where the training data remains distributed over large number of clients, which may be mobile phones or network sensors (Konečný al., 2016b;a; McMahan al., 2017: Mohri al., 2019; Kairouz al., 2019). A centralized model, here referred to server model, is then trained, without ever transmit- ting client data over the network, thereby providing some basic levels of data privacy and security.
Two important settings are distinguished in Federated learning (Kairouz et al., 2019, Table 1): the cross-device and the cross-silo settings. The cross-silo setting corresponds to a relatively small num- ber of reliable clients, typically organizations, such as medical or financial institutions. In contrast, in the cross-device federated learning setting, the number of clients may be extremely large and include, for example, all 3.5 billion active android phones (Holst, 2019). Thus, in that setting, may never make even a single pass over the entire clients' data during training. The cross-device setting is further characterized by resource-poor clients communicating over highly unreliable net- work. Together, the essential features of this setting give rise to unique challenges not present in the cross-silo setting. Here, we interested in the cross-device setting, for which we will formalize and study stochastic optimization algorithms.
The de facto standard algorithm for this setting is FEDAVG (McMahan al., 2017), which performs multiple SGD updates on the available clients, before communicating to the server. While this approach can reduce the total amount of communication required, performing multiple steps on the same client can lead to "over-fitting" to its atypical local data, phenomenon known as client drift (Karimireddy et al., 2020). Furthermore, algorithmic innovations such as momentum (Sutskever et al., 2013; Cutkosky and Orabona, 2019), adaptivity (Kingma and Ba, 2014; Zaheer et al., 2018; Zhang al., 2019), and clipping (You al., 2017; 2019; Zhang et al., 2020) are critical to the success of deep learning applications and need to be incorporated into the client updates, replacing the SGD update of FEDAVG. Perhaps due to such deficiencies, there exists a large gap in performance between the centralized setting, where data is centrally collected the server, and the federated setting (Zhao et al., 2018; Hsieh et al., 2019; Hsu et al., 2019; 10; Karimireddy et al., 2020).
To overcome such deficiencies, propose new framework, MIME, that mitigates client drift and adapts arbitrary centralized optimization algorithms, e.g. SGD with momentum or Adam, to the federated setting. In each local client update, MIME uses global statisties, e.g. momentum, and an

---

Under review as a conference paper at ICLR 2021
SVRG-style correction mimic the updates of the centralized algorithm run iid. data. These global statistics are computed only at the server level and kept fixed throughout the local steps, thereby avoiding bias due to the atypical local data of any single client.
Contributions. We summarize our main results below.
We formalize the cross-device federated learning problem, and propose a new framework MIME that can adapt arbitrary centralized algorithms to this setting.
We prove that incorporating server momentian into each local client update reduces client drift and leads to optimal statistical rates.
Further, we quantify the usefulness of performing multiple local updates on a single client by carefully tracking the bias (client-drift) introduced. This is the first analysis showing improved rates by taking additional multiple steps for general smooth functions.
Finally, we also propose a simpler variant, MIMELITE, with an empirical performance similar to MIME. We report the results of thorough experimental analysis demonstrating that both MIME and MIMELITE are faster than FEDAVG.
Related work. Analysis of FedAvg: Much of the recent work in federated learning has focused on analyzing FEDAVG. For identical clients, FEDAVG coincides with parallel SGD, for which Zinkevich et al. (2010) derived an analysis with asymptotic convergence. Sharper and more refined analyses of the same method, sometimes called local SGD, were provided by Stich (2019), and more recently by Stich and Karimireddy (2019), Patel and Dieuleveut (2019), Khaled et al. (2020), and Woodworth et al. (2020b), for identical functions. Their analysis was extended to heteroge- neous clients in (Wang et al., 2019; Yu et al., 2019b; Karimireddy et al., 2020; Khaled et al., 2020; Koloskova et al., 2020). Charles and Konečný (2020) derived a tight characterization of FedAvg with quadratic functions and demonstrated the sensitivity of the algorithm both client and server step sizes. Matching upper and lower bounds were recently given by Karimireddy et al. (2020) and Woodworth al. (2020a) for general functions, proving that FEDAVG can be slower than even SGD for heterogencous data, due to the client-drift.
Comparison SCAFFOLD: For the cross-silo setting where the number of clients is relatively low, Karimireddy et al. (2020) proposed the SCAFFOLD algorithm, which uses control-varien variates (similar SVRG) to correct for client drift. However, their algorithm crucially relies on stateful clients which repeatedly participate in the training process. In contrast, we focus on the cross-device setting where clients may be visited only once during training and where they are stateless. This akin to the difference between the finite-sum and stochastic settings in traditional centralized optimization.
Improvements to FedAvg: Hsu et al. (2019) and Wang et al. (2020c) observed that using server momentum significantly improves over vanilla FEDAVG. This idea was generalized by Reddi et al. (2020), who replaced the server update with an arbitrary optimizer, e.g. Adam. However, these methods only modify the server update while using SGD for the client updates. MIME, other hand, ensures that every local client update resembles the optimizer e.g. MIME would apply momentum in every client update and not just the server level. Beyond this, Li et al. (2018) proposed to add a regularizer to ensure client updates remain close. However, its usefulness is unclear (cf. Fig. Karimireddy et al., 2020; Wang et al., 2020b). Other orthogonal directions which can be combined with MIME include tackling computation heterogeneity, where some clients perform many more updates than others (Wang al., 2020b), improving fairness by modifying the objective (Mohri et al., 2019; Li et al., 2019), incorporating differential privacy (Geyer et al., 2017; Agarwal et al., 2018; Thakkar et al., 2020), Byzantine adversaries (Pillutla et al., 2019; Wang al., 2020a; He et al., 2020a), secure aggregation (Bonawitz et al., 2017: He et al., 2020b), etc. We refer the reader to the extensive survey by Kairouz al. (2019) for additional discussion.
the
2 PROBLEM SETUP
This section formalizes the problem of cross-device federated learning. We first examine some key challenges of this setting (cf. Kairouz et al., 2019) to ensure our formalism captures the difficulty:
1. Communication cost between the server and the clients is major concern and the source of bottleneck in federated learning; thus, a key metric for optimization in this setting is the number of communication rounds.
2

---

Under review as a conference paper at ICLR 2021
2. Each client is likely to participate at most once, due to the extremely large number of clients; furthermore, each individual client may have very little data of its own.
3. There may be a wide heterogeneity non-i.i.d.-ess due to the difference of data distributions for the clients.
Thus, our objective will be minimize the following quantity within the fewest number of client- server communication rounds:
n4 fa Si,u)
f(x) - [D)) (i
= The Σ
Here, fi denotes the loss function of client and (Gr. Stn.. its local data. Since the number of clients is extremely large, while size of each local data is rather modest, we represent the former as an expectation and the latter finite sum. each round, the algorithm samples subset of clients (of size S) and performs some updates to the server model. There is some inherent tension between the second and the third challenge outlined above: there exists a client with arbitrarily different data whom may never encounter during training, then there is no hope to actually minimize f. Thus for (1) to be tractable, is necessary to assume bounded dissimilarity between different ft-
(A1) G2-BGD or bounded gradient dissimilarity: there exists G such that EnDlll -V G², Vx.
Next, we also characterize the variance in the Hessians. Note that if fal:s () is L-smooth, (A2) is always satisfied with 5 < 2L and hence is more of a definition rather than assumption. Note that however, in realistic examples we expect the clients to be similar and hence that 6 < L.
(A2) 8-BHD bounded Hessian dissimilarity: Almost surely. f is ò-weakly convex i.c. v f()) = -ô1 and the loss function of any client i satisfies
| fi(x; ()
v2 [(=) < 6 Væ:
In addition, we assume that f(x) is bounded from below by f" and is L-smooth, as is standard.
USING MOMENTUM TO REDUCE CLIENT DRIFT
In this section we examine the tension between reducing communication by running multiple client updates each round, and degradation in performance due to client drift (Karimireddy et al., 2020). To simplify the discussion, we assume single client sampled each round and that clients use full-batch gradients.
Server-only approach. A simple way to avoid the issue of client drift is to take no local steps. We sample client ~ D and run SGDm with momentum parameter and step size y
21-1-
- ((1 100-33)f(ee (); + 3m,-1),
m, -1=
8m--1
Here, the gradient (z(;)) is unbiased i.e. [[V)((;;) (=x;;) and hence guaranteed convergence. However, this strategy can be communication-inenssii and we are likely to spend all our time waiting for communication with very little time spent on computing the gradients.
FedAvg approach. To reduce the overall communication rounds required, we need to make more progress each round of communication. Starting from yo = æt-1, FEDAvG (MeMahan et al., 2017) runs multiple SGD steps on the sampled client i ~ DD
y = yk v f. ( for l E [K],
and then pseudo-gradien 9. - x.,) replaces V (x( in the SGDm algorithm (2). This referred to as server-momentun since is computed and applied only at the server level (Hsu et al., 2019). However, such updates give rise to client-drift resulting in performance worse than the naive server-only strategy (2). This because by using multiple local updates, (3) starts over- fitting to the local client data, optimizing fi(x) instead of the actual global objective f(x).. The net
3

---

Under review as a conference paper at ICLR 2021
FEDAVG updates
MIMR updates
Et
222
Figure 1: Client-drift in FEDAVG (left) and MIME (right) is illustrated for 2 clients with local steps and momentum parameter = 0.5. The local SGD updates of FEDAvG (shown using arrows for client I and client2) move towards the average of client optima "i t which can be quite different from the true global optimum x". Server momentum only speeds up the convergence to the wrong point in this case. In contrast, MIME uses unbiased momentum and applies it locally every update. This keeps the updates of MIME closer to the true optimum x².
effect is that FEDAVG moves towards incorrect point (see Fig left). If K is sufficiently large, approximately
yk 2. where af arg min (1x))
+ Eslllg ~ (z, - E>>D(>;])
Further, the server momentum is based on gr and hence also biased. Thus, cannot correct for the client drift. We next see how a different way of using momentum can mitigate client drift.
Mime approach. FEDAVG experiences client drift because both the momentum and the client updates are biased. To fix the former, we compute momentum using only global statistics in (2):
m: = (1 - f. + 8m,-1.
(4)
To reduce the bias in the local updates, will apply this unbiased momentum every step:
yk yk-1 (7(11 B) V (gx Bm;-1) for k E [K]-
(5)
Note that the momentum term is kept fixed during the local updates i.e. there is no local momentum used, only global momentum is applied locally. Since m;-1 is a moving average of unbiased gra- dients computed over multiple clients, it intuitively is good approximation of the general direction of the updates. By taking convex combination of the local gradient with m,-1, 1. the update (5) potentially also less biased. In this way MIME combines the communication benefits of taking multiple local steps and prevents client-drift (see Fig 1, right). Sec. C makes this intuition precise.
4 MIME FRAMEWORK
In this section we describe how to adapt arbitrary centralized algorithms (and not just SGDm) to the federated learning problem (1) while ensuring there is no client-drift. Algorithm describes two variants MIME and MIMELITE, which consists of three components i) base algorithm we are trying to mimic. ii) how we compute the global statistics, and iii) the local client updates.
Base algorithm. We assume the centralized base algorithm we are imitating can be decomposed into two steps: an update step ii which updates the parameters æ, and a statistics step V(-) which keeps track of global statistics Each step of the base algorithm B (H1, V) uses gradient
x f x (2(g s).
(BASEALG)
Vgg s).
V may track multiple statistics which we represent collectively as s. While SGDm (2) is clearly of this form, Appendix ?? shows this for other algorithms like Adam,
4

---

Under review as a conference paper at ICLR 2021
Algorithm 1 Mime and MimeLite
input:
initial and s, learning rate 1j and base algorithm B (11, V)
for each round ( (= 1...., do
sample subset S of clients
communicate (ar, 8) to all clients E S
communicate e - [S EjEs (z) (only for Mime)
client i E S in parallel do initialize local model yı f- æ for k ..... K do
sample mini-batch Ç from local data
y: y: n V f(( (z) + c, s) (Mime)
y: f y/ (x (MimeLite)
end for
compute full local-batch gradient Va(()) communicate (y:, (((())) end on client
8 + 3 Eies Vf((at), a) (update optimization statistics)
x - [S] Eies (update server parameters)
1
end for
Compute statistics globally, apply locally. When updating the statistics of the base algorithm, we use only the gradient computed at the server parameters. Further, they remain fixed throughout the local updates of the clients. This ensures that these statistics remain unbiased and representative of the global function f-)). At the end of the round, the server performs
( Eies fi(()) 8)
where (f(
1 Σn R2
(STATS)
Note that we use full-batch gradients computed at the server parameters æ, not client parameters y:-
Local client updates. Each client E S performs K updates using U of the base algorithm and minibatch gradient. There variants possible corresponding MIME and MIMELITE differentiated using colored boxes. Starting from +- æ, repeat the following K times
Yi y: (g s),
where
(CLTSTEP)
g Vi(ln: Vf( + [S] Ejes (z)
MIMELITE simply uses the local minibatch gradient whereas MIME tion (Johnson and Zhang, 2013). This is done to reduce the noise from sampling a local mini-batch. While this correction yields faster rates in theory (and in practice for convex problems), in deep learning applications we found that MIMELITE closely matches the performance of MIME.
SVRG style correc-
Finally, there by the number of datapoints n4, and we perform K epochs instead of К steps (McMahan et al., 2017). The former modifies the objective (I) with being weighted by n4, and the latter has been empirically observed to perform better, but lacks strong justification (Wang et al., 2020b).
modifications made in practical FL.: weigh all averages across the clients
5 THEORETICAL ANALYSIS OF MIME
Table summarizes the rates of server-only methods, FEDAVG, SCAFFOLD and MIME (new results highlighted in blue). Our theory focuses MIME with base algorithm of SGD with momentum based variance reduction (MimeMVR) since it obtains optimal rates.
'The momentum based variance reduction (MVR). introduced by Cutkosky and Orabona (2019), is modi- fication of the standard SGDm algorithm to make it amenable to analysis. All our theory uses MVR, while our experiments SGDm.
5

---

Under review as a conference paper at ICLR 2021
Table Number of communication rounds required to reach f for L-smooth functions (log factors are ignored) with S clients sampled each round. C bounds the gradient dissimilarity (AI), and bounds the Hessian dissimilarity (A2). FEDAVG is slower than the server-only methods due to additional drift terms. Convergence SCAFFOLD depends on the total number of clients N which is potentially infinite. MIME matches the optimal statistical rates (first term in the rates) of the server-only methods while improving the optimization (second) term (typically 8 < L).
Algorithm
Non-convex
μ-- p-Strongly convex
SERVER-ONLY
G²
G MSe
SGD (Ghadimi and Lan, 2013)
+
+
Se
G
MVR (Cutkosky and Orabona, 2019)
+
V
FEDAVG
FedSGD (Karimireddy et al., 2020)
C2
+
G
+
L
G²
+
+
L
Sea
3/2
uSe
uyEe
SCAFFOLD2 (Karimireddy et al., 2020)
N
L
N
+
L
S
8
1
МIME
G2
G 6
MimeSGD
+
+
pSe
MimeMVR
VSc
Lower bound (Arjevani et al., 2019)
Ω
2 G2
vSe
Se
G Requires K 2 e² number of local updates with within-client variance of o².
2 In cross-device FL, the total number of clients (N) can be of the same order as number of rounds (since we only make few passes over all clients), or even co, making the bounds vacuous.
3 Requires K > L/ō number of local updates.
Theorem I. For L-smooth f with G2 gradient dissimilarity (A1), 8 Hessian dissimilarity (A2) and F = (f(x²)) - f*),, , let us run MR for T rounds and generate an output xu This output satisfies EllV f(n"")" S under the following conditions
PI-SStonnly convex without momentum: for n = Ö(min(
(mmin(XXIKKKLE 1
OKRUK+L ) = 0. and
L.G
T++ 8K Flog (2)).
==Ö
+
USE
μK
Non-convex without momentum: for O
minn (ann (TTKKz) )1/2)) and
L
(L+
T=
Se2
cK
Non-convex with momentum: for 17 = 0
|mmin
OK-L (GTKK3 (13))
1
SF
O
(TG7)77))
82
8)
+
3/4
(L (L+8K)F
T= 0
Sea
ek
The expectation in F 01l < is taken both the sampling of the clients during the running of the algorithm, the sampling of the mini-batches in local updates, and the choice of out (which is chosen randomly from the client iterates y; as described the Appendix).
Table shows that the rate (ignoring constants) of FEDAVG on non-convex functions is ( + G ( 4). In contrast,
G 4) This is slower than simply running SGD which obtains rate of
MiD obtains rate of $2 While asymptotically these three rates may seem equivalent, in machine learning we care about low accuracy settings where is not small (Bottou, 2010) and so the lower order terms matter, as
G² 6) where 5 < L thus improving upon both SGD and FedAvg.

---

Under review as a conference paper at ICLR 2021
G 10
G 100
201
10
l05s
107
107
18"
10t4
20
40
60
20
40
20
60
204
10
loss
20"
17'
Scom Km
10-4
40
60
20
60
loss
10a1
20-3
12:10
10 00
80 rounds
rounds
rounds
Figure 2: SGDm (dashed black), FedSGDm (top), MimeLiteSGDm (middle), and MimeSGDm (bottom) on simulated data, all with momentum (3 0.5). FedAvg gets slower as the gradient- dissimilarity (G) increases (to the right). MimeLite shows a similar pattern, but consistently better than FedAvg. Mime is significantly faster than both and is unaffected by heterogeneity (G).
also additionally evidenced by experiments (Sec. 6). We can also compare with SCAFFOLD $ where N is the total number of clients.
(Karimireddy et al., 2020) which obtains a rate of ( (z) While asymptotically this is a faster rate, N the cross-device setting potentially infinite or at least comparable to the total number of training rounds, making these bounds vacuous. This too is reflected our experiments (Fig. 5).
Finally, by incorporating momentum based variance reduction, MimeMVR matches the lower bound 3 of ((G/vSc)) by Arjevani et al. (2019). The momentum 3 used in this case is of the order of (1 O(TG))-2/3) i.e. T increases, momentum parameter asymptotically approaches 1. In contrast, previous analyses of distributed momentum (e.g. Yu et al. (2019a)) prove rates of the form
5(1-3)22 which are worse than that of standard SGD by factor of T-3 Thus, ours is the first result which theoretically showcases the usefulness of using large momentum in distributed and federated learning. Our theory suggests that the momentum parameter should be increased if G increases i.e. the clients become more heterogeneous, there is stronger client-drift and hence we need more momentum to compensate.
G²
Our analysis is is highly non-trivial and involves three crucial ingredients: i) computing the momen- tum at the server level to ensure that it remains unbiased and then applying it locally during every client update to reduce variance, carefully keeping track of the bias introduced via additional lo- cal steps, and iii) an SVRG correction to allow using mini-batches. Our experiments (Sec. 6) verify that the first two theoretical insights are indeed applicable deep learning settings as well, whereas the latter seems matter more in convex settings. See App. C where we make this discussion more concrete and Appendices F-G for detailed proofs and theorem statements.
6 EXPERIMENTAL ANALYSIS
We run experiments simulated and real datasets to confirm our theory. Our main findings are i) MIME outperforms FEDAVG across all settings, ii) its SVRG correction is useful for convex problems, and iii) momentum significantly improves performance for non-convex problems.
We consider four algorithms: SERVER-ONLY, FEDAVG, MIME, and MIMELITE. Each of these adapt base optimizers SGD, SGDm, and Adam. The SERVER- -ONLY method computes full batch gradient each of the sampled clients and uses their aggregate directly in the base optimizer (akin to (2)). For FEDAVG, we follow Reddi et al. (2020) who run multiple epochs of SGD on each client sampled, and then aggregate the net client updates. This aggregated update is used as pseudo- gradient in the base optimizer (called server optimizer). The learning rate for the server optimizer is fixed to 1 as in (Wang et al., 2020c). This is done to ensure all algorithms have the same number of hyper- prperpparameter.. Finally, MIME and MIMELITE follow Algorithm and also run a fixed number of epochs on the client. Aggregation is weighted by the number of samples the clients.

---

Under review as a conference paper at ICLR 2021
Momnettum methods on EMNIST62
Adam methods on EMNIST62
Advantage of momentum on EMNIST62
accuracy
0.78
SG0m
Adam FedAdam Mimeecaa MimeUltaAdam
FedSCimm, 0,9
FedSCiD
Mime 500m, B 0.9
MimesGDm
Fed5Gi
MimeLtteGDm
MimeGGD, p=
0.70
200
400
60n
800
1000
200
400
600
800
1000
200
400
600
800
1000
Momentum methods
Cifar100
Adam methods
Cifar100
Advantage of momentum CIFARI00
0.35
Fed56Dm, f- 0.9
Mina5CDm, A-000
FedsGD, A
accuracy
MimesGo, 8
0.20
SGDm
Adam
Fed
FedAdam
Mimessom
Mimstdam
MimeLiteeddam
0.05
200
400
600
800
1000
200
400
600
800
1000
200
400
1000
rounds
rounds
rounds
Figure 3: Server-only, FedAvg, Mime, and MimeLite with SGDm (left) and Adam (middle) run on (top) EMNIST62 and 2 hidden layer (300u-100) MLP and (bottom) Resnet20 run on Cifar 100. Mime and MimeLite have very similar performance and are consistently the best. FedAvg is even worse than the server-only baselines. Also, Mime makes better use of momentum than FedAvg, with a large increase in performance (right).
6.1 SIMULATED CONVEX EXPERIMENTS
Our simulated experiments use two clients each with simple scalar quadratic loss, as in (Karim- ireddy et al., 2020). We use full-batch gradients with both clients participating every round. The simulated data has Hessian dissimilarity 5 (A2) and smoothness L 2. We vary the gradient dissimilarity (A1) as G E [1, 10, 100]. All the algorithms use momentum with 3 = 0.5 and their learning rates were tuned up to tolerance of 5E-3 to ensure lowest loss after 60 rounds. The results are collected in Fig. 2. When G is small, we see that FEDAVG can outperform the SERVER-ON -ONLY (SGDm) baseline, though its loss quickly plateaus. On increasing G, FEDAVG becomes even slower. MIMELITE differs from FEDAVG only in how the momentum is used. In all settings, it slightly out- performs FEDAVG though even it sees a substantial slow down as we increase G. This reflects our theory which predicts that for convex cases, momentum does not give significant gains. MIME, on the other hand, is substantially faster than all other methods and is even unaffected by changing G. Thus, in this simple convex setting, the SVRG correction completely eliminates client drift.
6.2 REAL WORLD EXPERIMENTS
We run real world deep learning experiments on EMNIST62 with a 2 layer MLP model and on C0 with ResNet20, both accessed through Tensorflow Federated (TFF, 2020a). All methods run 10 local epochs, batch size 20, and the learning rates for all methods were individually tuned. We refer Appendix A for additional details and results. Fig shows the results.
Mime > MimeLite Server-only > FedAvg. Mime and MimeLite have the best performance. FedAvg is slower than even the naive server-only methods which make no local updates. This perfectly mirrors theory that Mime > server-only FedAvg. The performance of MimeLite is because SVRG correction may not be necessary in deep learning (Defazio and Bottou, 2019).
With momentum without momentum. Fig. 3 (right) examines the impact of momentum on FedAvg and Mime. Momentum slightly improves the performance of FedAvg, whereas it has a significant impact on the performance of Mime. This is also in line with our theory and confirms that Mime's strategy of applying it locally at every client update makes better use of momentum.
Fixed statistics > updated statistics. Finally, we check how the performance of Mime changes if instead of keeping the momentum fixed throughout round, we let it change. The momentum is reset at the end of the round ignoring the changes the clients make to it. Appendix B. shows that this consistently worsens the performance, confirming that it is better keep the statistics fixed.
8

---

Under review as a conference paper at ICLR 2021
Together, the above observations validate all aspects of Mime (and MimeLite) design: compute statistics at the server level, and apply them unchanged at every client update.
7 CONCLUSION
Our work initiated formal study of the cross-device federated learning problem. We argued that the natural heterogeneity among the clients gives rise to client drift and significantly hampers the performance of approaches such as FEDAVG. We then showed how momentum can be excellent tool to overcome this client drift used correctly. Based on this observation, we introduced new framework MIME which not only overcomes client drift, but also adapts arbitrary centralized algorithms such as Adam to the federated setting without any additional hyper- parameters.... We demonstrated the superiority of MIME via strong convergence guarantees and empirical evaluations.
REFERENCES
Naman Agarwal, Ananda Theertha Suresh, Felix X. Yu, Sanjiv Kumar, and Brendan McMahan. private distributed SGD. In Proceedings of
cpSGD: Communication-efficie and NeurlPS, pages 7575-7586, 2018.
Dan Alistarh, Demjan Grubic, Jerry Li, Ryota Tomioka, and Milan Vojnovic. Qsgd: Communication-effiiient sgd via gradient quantization and encoding. In Advances in Neural Information Processing Systems (NeurIPS), 2017.
Yossi Arjevani and Ohad Shamir. Communication complexity of distributed convex learning and optimization. In Advances in neural information processing systems, pages 1756-1764, 2015.
Yossi Arjevani, Yair Carmon, John C Duchi, Dylan Foster, Nathan Srebro, and Blake Woodworth. Lower bounds for non-convex stochastic optimization. arXiv preprint arXiv: DeXivv/9122 1912.0235, 2019.
Keith Bonawitz, Vladimir Ivanov, Ben Kreuter, Antonio Marcedone, H. Brendan McMahan, Sarvar Patel, Daniel Ramage, Aaron Segal, and Kam Seth. Practical secure aggregation for privacy- preserving machine learning. In Proceedings of the 2017 ACM SIGSAC Conference on Computer and Communications Security, pages 117711991. ACM, 2017
Keith Bonawitz, Hubert Eichner, Wolfgang Grieskamp, Dzmitry Huba, Alex Ingerman, Vladimir Ivanov, Chloe Kiddon, Jakub Konečný, Stefano Mazzocchi, H Brendan McMahan, et al. Towards federated learning at scale: System design. arXiv preprint arXiv:1902.01046. 2019.
Léon Bottou. scale machine learning with stochastic gradient descent. In Proceedings of COMPSTAT'2010, pages 177-186. Springer, 2010.
Sebastian Caldas, Jakub Konečny, H Brendan McMahan, and Ameet Talwalkar, Expanding the reach of federated learning by reducing client resource requirements. arXiv preprint arXiv:1812.0722 2018a.
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 arXiv: 18/2.01097, 2018b.
Zachary Charles and Jakub Konečný, On the outsized importance of learning rates in local update methods. arXiv preprint arXiv:2007.00878, 2020.
Ashok Cutkosky and Harsh Mehta. Momentum improves normalized SGD. arXiv:2002.03305, 2020.
arXiv preprint
Ashok Cutkosky and Francesco Orabona. Momentum-based variance reduction non-convex SGD. In Advances in Neural Information Processing Systems, pages 15210-15219,200
Aaron Defazio and Léon Bottou. On the ineffectiveness of variance reduced optimization for deep learning. Advances in Neural Information Processing Systems, pages 1753-1763, 2019.
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. International Conference Learning Representations (ICLR), 2019.

---

Under review as a conference paper at ICLR 2021
Roy Frostig, Matthew James Johnson, and Chris Leary. Compiling machine learning programs via high-level tracing. Systems for Machine Learning, 2018.
Robin C Geyer, Tassilo Klein, and Moin Nabi. Differentially private federated learning: A client level perspective. arXiv preprint crXiv: 2017.
Saeed Ghadimi and Guanghui Lan. Stochastic first-and zeroth-order methods for nonconvex stochas- tic programming. SIAM Journal on Optimization, 23(4):2341-2368, 2013.
Haiku. JAX Resnet model in Haiku, 2020. URL https://gitthu. com com/deepmind/ dm http///ooommasteeee// httpaaaawwo..eeeee/s/eeettttttte
Jenny Hamer, Mehryar Mohri, and Ananda Theertha Suresh. FedBoost: Communication-effie effiicient algorithms for federated learning. In 37th International Conference on Machine Learning (ICML), 2020.
Andrew Hard, Kurt Partridge, Cameron Nguyen, Niranjan Subrahmanya, Aishanee Shah, Pai Zhu, Ignacio Lopez Moreno, and Rajiv Mathews. Training keyword spotting models on non-iid data with federated learning. arXiv preprint 200:20005.10066 2020.
Lie Hc, Sai Praneeth Karimireddy, and Martin Jaggi. Byzantine- robust learning on heterogencous datasets via resampling. arXiv preprint arXiv:2006.0900 003665, 2020a.
Lie He, Sai Praneeth Karimireddy, and Martin Jaggi. Secure byzantine-robust machine learning. arXiv preprint arXiv:2006.0474: 2020b.
Arne Holst. Smartphone users worldwide 2016-2021, 2019.
Kevin Hsich, Amar Phanishayee. Onur Mutlu, and Phillip B Gibbons. The Non-IID data quagmire of decentralized machine learning. arXiv preprint 19000:1910.00199. 2019.
Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the effects of non-identical data distribution for federated visual classification. arXiv preprint wwXivv:909.0633:0 2019.
Rie Johnson and Tong Zhang. Accelerating stochastic gradient descent using predictive variance reduction. In Advances in neural information processing systems, pages 315-323, 2013.
Peter Kairouz, 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:1912.04977 2019.
Hamed Karimi, Julie Nutini, and Mark Schmidt. Linear convergence of gradient and proximal- gradient methods under the polyak-tojasiewve condition. Joint European Conference on Ma- chine Learning and Knowledge Discovery in Databases, pages 795-81 811. Springer, 2016.
Sai Praneeth Karimireddy, Quentin Rebjock, Sebastian U. Stich, and Martin Jaggi. Error feedback fixes SignSGD and other gradient compression schemes. In 36th International Conference on Machine Learning (ICML), 2019.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian U Stich, and Ananda Theertha Suresh. SCAFFOLD: Stochastic controlled averaging for on-device federated learning. 37th International Conference on Machine Learning (ICML), 2020.
Ahmed Khaled, Konstantin Mishchenko, and Peter Richtárik. Tighter theory for local SGD on indentical and heterogeneous data. Proceedings of AISTATS, 2020.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Anastasia Koloskova, Nicolas Loizou, Sadra Boreiri, Martin Jaggi, and Sebastian U Stich. A unified theory of decentralized SGD with changing topology and local updates. In 37th International Conference Machine Learning (ICMI), 2020.
10

---

Under review as a conference paper at ICLR 2021
Jakub Konečný, H. Brendan McMahan, Daniel Ramage, and Peter Richtárik. Federated optimiza- tion: Distributed machine learning for on-device intelligence. arXiv preprint arXiv: arXiv:1610.02527, 2016a.
Jakub Konečný, H. Brendan McMahan, Felix X. Yu, Peter Richtárik, Ananda Theertha Suresh, and Dave Bacon. Federated learning: Strategies for improving communication efficiency. arXiv preprint trriiv:/600.055 2016b.
kuangliu. Train CIFARIO with PyTorch, 2020 (accessed June 4, 2020). URL htt.httts:////tthee com/ buangliu/ http:l/ipppptoocc-hhaeeea
Tian Li, Anit Kumar Sahu, Maziar Sanjabi, Manzil Zaheer, Ameet Talwalkar, and Virginia Smith. On the convergence of federated optimization in heterogeneous networks. arXiv preprint arXiv:1811..06127, 2018.
Tian Li. Maziar Sanjabi, and Virginia Smith. Fair resource allocation in federated learning. arXiv preprint arXiv:/905. /0497, 2019.
Brendan McMahan, Fider Moore, Daniel Ramage, Seth Hampson, and Blaise Agüera y Arcas. Communication-effcee learning of deep networks from decentralized data. In Proceedings of AISTATS, pages 1273-1282, 2017.
Mehryar Mohri, Gary Sivek, and Ananda Theertha Suresh. Agnostic federated learning. arXiv preprint arXiv:1902.00146, 2019.
Yurii Nesterov. Lectures on convex optimization, volume 137. Springer, 2018.
Kumar Kshitij Patel and Aymeric Dieuleveut. Communication trade-off's for synchronized dis-
tributed SGD with large step size. In 33rd Conference on Neural Information Processing Systems (NeurIPS), 2019.
Krishna Pillutla, Sham M Kakade, and Zaid Harchaoui. Robust aggregation for federated learning. arXiv preprint arXiv: 19/2./344/3 2019.
Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and H Brendan McMahan. Adaptive federated optimization. arXiv preprint 0rXiv:2003.00295. 2020.
Sashank J. Reddi, Jakub Konečný, Peter Richtárik, Barnabás Póczós, and Alex Smola. Aide: Fast and communication efficient distributed optimization. arXiv preprint
2016.
Ohad Shamir, Nati Srebro, and Tong Zhang. Communication-effiiine distributed optimization using an approximate newton-type method. In International conference on machine learning, pages 1000--008, 2014.
Sebastian U. Stich. Local SGD converges fast and communicates little. International Conference on Learning Representations (ICLR), 2019.
Sebastian U. Stich and Sai Praneeth Karimireddy. The error-feedback framework: Better rates for SGD with delayed gradients and compressed communication. arXiv preprint urXiv:1909.06. 2019.
Ananda Theertha Suresh, Felix X. Yu, Sanjiv Kumar, and H. Brendan MeMahan. Distributed mean estimation with limited communication. In Proceedings of the 34th International Conference Machine Learning-Vain Volume 70, pages 3329-3337. JMLR. org, 2017.
Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initial- ization and momentum in deep learning. In International conference on machine learning, pages 1139-1147, 2013.
TFF. Tensorflow Federated Datasets, 2020a. URL wtw.pp///o/w..ttoorooooooooo.... mwwaammo/daaaaateaa datanndaaattsststsa
federated/age appi http:/appioocspytttaaaa

---

Under review as a conference paper at ICLR 2021
TFF. Tensorflow Federated Averaging Implementation, 2020b.
URL https://ww.. tersorflow.org/ wawgg//eeeatteedeeee attpa//cc/ppottooo///////aaaaaaa process.
Om Thakkar, Swaroop Ramaswamy, Rajiv Mathews, and Françoise Beaufays. Understanding unin- tended memorization in federated learning. arXiv preprini arXiv:2006.074990 2020.
Quoe Tran-Dini Nhan H. Pham, Dzung T. Phan, and Lam M. Nguyen. Hybrid stochastic gradi- ent descent algorithms for stochastic nonconvex optimization. arXiv preprint 17Xiv:1005.0535 2019.
Sharan Vaswani, Francis Bach, and Mark Schmidt. Fast and faster convergence of SGD for over- parameterized models and an accelerated perceptron. arXiv preprint arXiv:1880.0703 2018.
Hongyi Wang, Kartik Sreenivasan, Shashank Rajput, Harit Vishwakarma, Saurabh Agarwal, Jy- yong Sohn, Kangwook Lee, and Dimitris Papailiopoulos. Attack of the tails: Yes, you really can backdoor federated learning. arXiv preprint ar Xrvi::2007. 005044, 2020a.
Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and H Vincent Poor. Tackling the objective in- consistency problem in heterogeneous federated optimization. arXiv preprint arXiv:2007.07481, 2020b.
Jianyu Wang, Vinayak Tantia, Nicolas Ballas, and Michael Rabbat. SlowMo: Improving communication-effcient distributed sgd with slow momentum. International Conference on Learning Representations (ICLR), 2020c.
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 Communications, 37(6): 20122025222 2019.
Blake Woodworth, Kumar Kshitij Patel, and Nathan Srebro. Minibatch vs local SGD for heteroge- neous distributed learning. arXiv preprint arXiv:2006.0477: 04735, 2020a.
Blake Woodworth, Kumar Kshitij Patel, Sebastian U Stich, Zhen Dai, Brian Bullins, H Brendan McMahan, Ohad Shamir, and Nathan Srebro. Is local SGD better than minibatch SGD? In 37th International Conference on Machine Learning (ICML), 2020b.
Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv preprint arXiv: /708.0000000 038888 2017.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsich. Large batch optimization for deep learning: Training bert in 76 minutes. In International Conference on Learning Representations, 2019.
Hao Yu, Rong Jin, and Sen Yang. On the linear speedup analysis of communication efficient mo- mentum sgd for distributed non-convex optimization. arXiv preprint /005.03/90003///// 2019a.
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, pages 5693-5700, 2019b.
Manzil Zaheer, Sashank Reddi, Devendra Sachan, Satyen Kale, and Sanjiv Kumar. Adaptive meth- ods for nonconvex optimization. In Advances in neural information processing systems, pages 9793-9803, 2018
Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, San- jiv Kumar, and Suvrit Sra. Why ADAM beats SGD for attention models. arXiv preprint arXiv:1912.03199 2019.
Jingzhao Zhang, Tianxing He, Suvrit Sra, and Ali Jadbabaie. Why gradient clipping accelerates training: A theoretical justification for adaptivity. In International Conference on Learning Rep- resentations, 2020.
12

---

Under review as a conference paper at ICLR 2021
Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. Federated learning with non-iid data. arXiv preprint arXi:: 00582, 2018.
Martin Zinkevich, Markus Weimer, Lihong Li, and Alex J Smola. Parallelized stochastic gradient descent. In Advances in neural information processing systems, pages 2595-2603, 2010.
13

---

Under review as a conference paper at ICLR 2021
Appendix
A EXPERIMENTAL SETUP
A.1 DESCRIPTION DATASETS AND TASKS
We perform 4 tasks over 3 datasets: i) On the EMNIST62 (extended MNIST) dataset (Caldas et al., 2018b) we run a convex logistic regression model, a fully connected MLP with 2 hidden layers (300u- 100), and convolution model with two CNN layers and two dense layers with dropout. ii) we also run a ResNer20 on CIFART00.
In all cases we report the 1 test accuracy in our experiments. EMNIST uses the metadata indicat- ing the original author of the characters to separate them into multiple clients yielding naturally par- titioned dataset. Table 2 summarizes the statistics about the different datasets. Note that the average number of rounds a client participates in (computed as sampled clients xumbee of rounds / nnmbber of clients) provides an indication of how much of the training data is seen with SHAKESPEARE being closest to the cross-silo setting and STACKOVERFLOW representing the most cross-device in nature.
Table 2: Details about the datasets used and experiment setting.
EMNIST62
CIFAR100
Clients
3,400
500
Examples
671,585
50,000
Sampled clients
20
20
Batch size
20
20
Number of epochs
10
10
Avg. rounds each client participates
5.9
40
We use Tensorflow federated datasets (TFF, 2020a) to generate the datasets. Our federated learning simulation code is written in Jax (Frostig al., 2018). Our Resnet18 model is based off of (Haiku, 2020) (Resnet v2), and following (Hsieh et al., 2019; Reddi et al., 2020) we replace batch norm with group norm with groups. Black and white was reversed in EMNIST62 (i.e. subtracted from 1) to make them similar to MNIST. CIFAR 100 used the usual pre-processing (normalization and centering). and data augmentation (random crop and horizontal flipping) following (kuangliu, 2020 (accessed June 4, 2020).
A.2 PRACTICALITY OF EXPERIMENTS
In the experiments only cared about the number of communication rounds, ignoring that MIME actually needs twice the number of bits per round and that the SERVER- -ONLY methods have a much smaller computational requirement. This is standard in the federated learning setting introduced by McMahan et al. (2017) and is justified because most of the time in cross-device FL is spent in establishing connections with devices rather than performing useful work such as communication or computation. In other words, latency and not bandwidth or computation are critical in cross device FL. However, one can certainly envision cases where this is true. Incorporating communication compression strategies (Suresh et al., 2017; Alistarh et al.. 2017; Karimireddy et al., 2019) or client- model compression strategies (Caldas al., 201 Sa; Frankle and Carbin, 2019; Hamer et al., 2020) into our MIME framework can potentially address such issues and are important future research directions.
As we already discussed previously, we believe both the datasets and the tasks being studied here are quite realistic in nature. We now discuss our choice of other parameters the experiment setup (number of training rounds, sampled clients, batch-size, etc.) Each round of federated learning takes 2-3 mins in the real world and is relatively independent of the size of communication (Bonawitz et al., 2019) implying that training 1000 rounds takes 1.4-2 days. This underscores the importance of ensuring that the algorithms for federated learning converge in as few rounds as possible, as well have very easy to set default hyper-parameters, Thus in experimental setup keep
14

---

Under review as a conference paper at ICLR 2021
all parameters other than the learning rate to their default values. In practice, this learning rate can be set by set using a small dataset on the server (as in (Hard et al., 2020)). The choice of batch size being 10 made both keeping in mind the limited memory available to each client well as to match prior work. Finally, while we limit ourselves to sampling 20 workers per round due computational constraints, in real world FL thousands of devices are often available for training simuitaneously each round (Bonawitz et al., 2019). They also note that the probability of each of these devices being available has clear patterns and far from uniform sampling. Conducting large scale experimental study which mimics these alternate forms of heterogeneity is an important direction for future work.
A.3 HYPERPARAMETER SEARCI
We the EMNIST62 with MLP model 'test-bed' for exploring different algorithms given it being both a representative task of cross-device FL as well as being computationally efficient. All plots reported for this setting. A more fine- ene-ggained search over hyperparameters report the final test accuracies is made over the rest of the tasks/datasets.
For all SGDm methods, we pick momentum
0.9. For Adam methods, fix 31 0.9, 82 = 0.99, and = 1 x 10-3 similar to (Reddi et al., 2020). None of the algorithms use weight decay, clipping etc. The learning rate is then tuned obtain the best accuracy.
For all experiments, unless explicitly mentioned otherwise, the learning rate is searched over grid of
E [1 x 10', ,,,1, x 10-", 10 2, 10 1x 10 4 1x 10 %.
A.4 COMPARISON WITH PREVIOUS RESULTS
As far as we are aware, (Reddi et al., 2020) is the only prior work which conducts a systematic experimental study of federated learning algorithms over multiple realistic datasets. The algorithms comparable across the two works (e.g. FedSGD, FedSGDm, and FedAdam) have qualitatively sim- ilar performance except with one exception: FedAdam consistently underperforms FedSGDm. We believe this difference is because (Reddi et al., 2020) additionally tune the server learning rate and the parameter for Adam. As we explain in Section A.2, chose keep these parameters some default values to compare methods in the 'low-tuning" setting. We also point that while FedAdam struggles perform in this setup, MimeAdam and Admm are very stable and even often outperform their SGDm counterparts. This is also the default/recommendee behavior in TensorFlow Federated (TFF, 2020b) and (Wang et al., 2020c).
B ADDITIONAL EXPERIMENTS
B. EFFECT OF CHANGING STATISTICS
Momentum methods EMNIST62
Adam methods EMNIST62
0.86
0.86
accuracy
accuracy
0.78
SGDm
Adam
MimesGiDm MimesGDD
MimeAdimm
C-Mimeddam
0.70
0.70
800
1000
600
rounds
rounds
Figure 4: Server-only, FedAvg, Mime, and C-Mime with SGDm (left) and Adam (right) run on EMNIST62 with 2 hidden layer (300u-100) MLP. C-Mime changes the statistics (momentum for SGDm, and first two moments for Adam) using the local client updates. These changes are discarded the end of the round and the statistics are reset using only the server level gradients as in Mime. Clearly, C-Mime is always worse than Mime. This shows that adapting statistics during local client updates makes them too biased, and best keep them fixed during each round like Mime does.
15

---

Under review as a conference paper at ICLR 2021
Momentum methods on EMNIST62
0.86
accuracy
Mime Scaffold FedProx
200
1000
Figure 5: Comparison with Scaffold and FedProx for cross-device FL: Mime, SCAFFOLD and FedProx with SGDm run on EMNIST62 with a hidden layer (300u-100) MLP. For FedProx and SCAFFOLD, in addition to tuning the learning rate, we search for the best server momentum B [0, 0,.0990..k 0.99]. FedProx uses an additional regularizer μe which we search over 10.1, 0.5, 1] (note that FedProx with μ 0 the same FedAvg). The best accuracy (which are plotted here) was by 3 = 0 for both and y = 0.1 for FedProx. Note that FedProx the slowest method here (in fact it is even slower than FedAvg). The additional regularizer does not seem to reduce client drift while still slowing down convergence (Karimireddy et al., 2020; Wang et al., 2020b). SCAFFOLD also slower than Mime in this setup. This because SCAFFOLD was designed for the cross- silo setting and not corss-device setting. The large number of clients (N = 3.4k) means that each client on averaged visited less than times during the entire training (20 clients per round for Ik rounds). Hence, the client control variate stored is quite stale (from about 200 rounds ago) which slows down the convergence. This perfectly reflects our theoretical understanding that when the number of clients N is large relative to training rounds (which is true in the cross-device setting) SCAFFOLD is outperformed by MIME.
FEDAVG
MIME
MIMELITE
SGD
SGDm Adam
SGD
SGDm Adam SGD
SGDm Adam
EMNIST62 logistic
66.0
66.8
67.5
67.2
67.3
68.0
66.1
67.3
68.0
EMNIST62 CNN
84.9
85.4
85.7
86.0
85.6
86.0
85.0
85.4
85.7
Table High tuning setting: final test accuracy (larger is better) with fully tuned hyppe--parmmetes. For FEDAVG we searched over both the client and server learning rates, whereas for MIME and MIMELITE, we search only over client (base) learning rate. This search performed over grid (0.002, 0.005, 0.01, 0.02, 0.05, 0.1, 0.2, 0.4, 0.8). For momentum, we chose best of (0.9, 0.99) and for Adam, we varied 31 in (0.9, 0.99), 82 in (0.99, 0.999) and in (0.01, 0.001, 0.0001). In this setting well MIME and MIMELITE outperform FEDAVG but with smaller margin.
Table 4: Additional algorithmic details: Decomposing base algorithms into parameter update (U) and statistics tracking (V).
Algorithm SGD
Tracked statistics s
Update step L
Tracking step V
#9
SGDm m
x (111 - -B)g + +8m)
m = (1 3)g + 3m
RMSProp
x
7 ct
v = (1 8)g9 + 300
m (1 -31g + p.m
Adam
m.
((1 -B1)g + B.m)
V = (1 }2)g² + +B2
C PROOF OVERVIEW
In this section, we give proof sketches of the main components of Theorem I: how momentum reduces the effect of client drift, ii) how local steps can take advantage Hessian similarity. and iii) why the SVRG correction improves constants.
16

---

Under review as a conference paper at ICLR 2021
Improving the statistical term via momentum. Note that the statistical (first) term in Theorem I without momentum (3 = 0) for the convex case pSe
LG This is (up constants) optimal and cannot gives the usual rate of LG"
be improved. For the non-convex case however using
However,
"e
this can be improved (1+8)GFF) 3/1
using momentum. This matches similar improvement in the centralized setting (Cutkosky and Orabona. 2019; Tran-Dinh et al., 2019) and is in fact optimal (Arjevani et al., 2019). Let us examine why momentum improves the statistical term. Assume that sample single client it in round and that full-batch gradients. Also let the local client update at step l round be of the form
Se:
+ ndy
(6)
The ideal choice of update is of course di = f but however this is unattainable. Instead, MIME with momentum B = a uses dsGDDm = m + (V(y + (1 a)m;-1 where m is the momentum computed the server. The variance of this update can then be bounded as Vf( 11 S a E (v) - V +
El mg
El -Vf(y))
d² + +1 (1-a) (l-=a) (1-a))\mm---V((aa--))) (f(
aG²
The last step follows by unrolling the recursion on the variance of m. We also assumed that v is small enough that y Φt-2. This way, momentum reduce the variance of the update from G to (aG²) by using past gradients computed on different clients. To formalize the above sketch requires slightly modifying the momentum algorithm similar (Cutkosky and Orabona, 2019), and carried out in Appendix G.
Improving the optimization term via local steps. The optimization (second) term in Theorem I for the convex case is dK+L and for the non-convex case (with without momentum) is 8K+L, uk
<K
In contrast, the optimization term of the server-only methods is L/P μ and L/ce respectively. Since in most cases 5 < L. the former can be significantly smaller than the latter. This rate also suggests that the best choice of number of local updates is L/ō ie. we should perform more client updates when they have more similar Hessians. This generalizes results of (Karimireddy et al., 2020) from quadratics to all functions.
This improvement is due careful analysis of the bias in the gradients computed during the local update steps. Note that for client parameters yk-1, the gradient F fi, (yk-1)) o F(V(( since Yk-1 also computed using the same loss function fi.- In fact, only the first gradient computed at - is unbiased. Dropping the subscripts & and f, we can bound this bias as:
((V)(t)) ((())) ((V()()) - V((()) V - V((y)) +
=0 since unbiased
2 [[(V?
V V f())) x)) ~ El(y - x]]
Thus, the Hessian dissimilarity (A2) control the bias, and hence the usefulness of local updates. This intuition can be made formal using Lemma 3.
Mini-batches via SVRG correction. In our previous discussion about momentum and local steps, we assumed that the clients compute full batch gradients and that only one client is sampled per round. However, in practice a large number (S) of clients are sampled and further the clients use mini-batch gradients. The SVRG correction reduces this within-client variance since
G2
G2
Var V ()) V Vi(x;; C) + [S] Eies V)((m) & T.2 2012
S
S
Here, we used the smoothness of fe(:: () and assumed that ys = x since we don't 't move too far within single round. Thus, the SVRG correction allows updates while still ensuring that the variance of the order G²/S.
minibatch gradients in the local
D TECHNICALITIES
We examine additional definitions and introduce some technical lemmas.
17

---

and proving tight Under review as a conference paper at ICLR 2021
D.1 ASSUMPTIONS AND DEFINITIONS
We make precise a few definitions and explain some of their implications. We first discuss the two assumptions on the dissimilarity between the gradients (A1) and the Hessians (A2). Loosely, these two quantities are an extension of the concepts of variance and smoothness which occur in centralized SGD analysis to the federated learning setting. Just as the variance and smoothness are completely orthogonal concepts, we can have settings where G² (gradient dissimilarity) is large while 8 (Hessian dissimilarity) is small, or vire-versa.
Our assumption about the bound on the G gradient dissimilarity can easily be extended to (G, gradient dissimilarity used by (Karimireddy et al., 2019):
E Vf(x))) - G + B² B
All the proofs in the paper extend in a straightforward manner to the above weaker notion. Since this notion does not present any novel technical challenge, we omit it in the rest of the proofs. Note however that the above weaker notion can potentially capture the fact that by increasing the model capacity, we can reduce G. In the extreme case, by taking a sufficiently over-parameterrieed model, it is possible to make G O in certain settings (Vaswani et al., 2018). However, this comes both at a cost of increased resource requirements (i.e. higher memory and compute requirements per step) but can also result in other constants increasing (e.g. B and T.).
The second crucial definition we use in this work is that of ō bounded Hessian dissimilarity (A2). This has been used previously in the analyses of distributed (Shamir et al., 2014; Arjevani and Shamir, 2015; Reddi et al., 2016) and federated learning (Karimireddy et al., 2020), but has been restricted to quadratics. Here, we show how to extend both the notion well as the analysis to general smooth functions. The main manner we will use this assumption is in Lemma 3 to claim that for any æ and y the following holds;
(8) Here the expectation is both over Ç as well as the choice of client i. To understand what the above condition means, it is illuminating to define (, fi(2; f(z). Then, we can rewrite (A2) and (8) respectively as
((r: - Vi((:; Vf(x)- - f 1 S 8 x1
and V - VV, C)12
Thus (8) and (A2) are both different notions of smoothness of ( () (formal definition of smooth- ness will follow soon). The latter definition closely matches the notion of squared-smootheg by Arjevani et al. (2019) and is a promising relaxation of (A2). However, we run into some techni- cal issues since in our case the variable can also be random variable and depend on the choice of the client i. Extending our results to this weaker notion of Hessian-similne non-convex lower bounds is an exciting theoretical challenge.
used
Finally note that if the functions f: <) are assumed to be smooth as in (Shamir et al., 2014; Arjevani and Shamir, 2015; Karimireddy al., 2020), then (( <) is 2L-smooth. Thus, always have that 5 < 2L. But, as Shamir et al. (2014) show, it is possible to have 5 < L if the data distribution amongst the clients is similar. Further, the lower bound from Arjevani and Shamir (2015) proves that Hessian-simillrity is the crucial quantity capturing the number of rounds of communication required for distributeddfederraee optimization.
We next define the terms smoothness and strong-convexity which we repeatedly use in the paper.
(A3) f is L-smooth and satisfies:
[V - f S L| - -yl. for any y.
(9)
The assumption (A3) also implies the following quadratic upper bound on f
L
f(y) f(x) + (Vf(z), y x)) + プ yl ==12
(10)
Further, if f is twice-differentiaall, (A3) implies that || f())l S 3 for any æ.
(A4) f is p-PL strongly convex (Karimi et al., 2016) for μ. > 0 if it satisfies:
[-V 2 (n(f(x) f*)-
Note that PL-strong convexity is much weaker than the standard notion of ctrong- (Karimi et al., 2016).
18

---

Under review as a conference paper at ICLR 2021
D.2 SOME TECHNICAL LEMMAS
Now we cover some technical lemmas which are useful for computations later on. First, we state relaxed triangle inequality for the squared &g norm.
Lemma (relaxed triangle inequality). Let (U1...., 4t be vectors in R4, Then the following are me:
0 + wyll $ (1 c)|l=,1 (1+ + ,) f for any e 0, and
2. |E vill T Σ 11:11.
Proof. The proof of the first statement for any > 0 follows from the identity:
||e + wyll (1 + c)]l0;11 (1 + - o,1 Ilven, + Jz
For the second inequality, we use the convexity -> ||a||| and Jensen's inequality
ト
2
Next we state an elementary lemma about expectations of norms of random vectors.
Lemma (separating mean and variance). Let , E. be random variables in R4 which are necessarily independent. First suppose that their mean is [[E,] E and variance is bounded as E - E S o². Then, the following holds
F
T i=1
Σ
E1l + +
Now instead suppose that their conditional mean is I...E. i.e. the variables { & form martingale difference sequence, and the variance is bounded by E[I| E&ll S 0² as before. Then we can show the tighter bound
<
1 2TO2.
Proof. For any random variable X, E[X*] = (E[[ E[X]])+ (E[X))² implying
Σ E
|> E
E(I Σ E. 1
Expanding the above expression using relaxed triangle inequality (Lemma 1) proves the first claim:
Elll Σ> - &ll T Σ ElIE, <
= 7
For the second statement, & is not deterministic and depends on E.. Ei. Hence we have to resort to the cruder relaxed triangle inequality to claim
Σ
E
2
E.1l
2E
> E, - 1
iml
and then use the tighter expansion the second term:
E|| -
- - - =
-3)] =
Σ E [| &11 To²
The cross terms the above expression have zero mean since sequence.
(E,-6,} } form martingale difference
19

---

Under review as a conference paper at ICLR 2021
E PROPERTIES OF FUNCTIONS WITH ó BOUNDED HESSIAN DISSIMILARITY
We now study two lemmas which hold for any functions which satisfy (A2). The first is closely related the notion of smoothness (A3).
Lemma 3 (similarity). The following holds for any two functions f((:i)) and f() satisfying (A2), and any x. y:
6 f.l( Vfi(=: () Vf(x)) 1 a112
Proof. Consider the function (=)) (: f()). By the assumption (A2), know that [[]]l & 5 for all i.. v is 6-smooth By standard arguments based on taking limits (Nesterov,
2018), this implies that
(v(y)) - (=) < ö1yy - æi|-
Plugging back the definition of into the above inequality proves the lemma.
Next, we see how weakly-convex functions satisfy a weaker notion of "averaging does hurt". Lemma 4 (averaging). Suppose f is ō-weakly convex. Then, for any Y 2 ō, and sequence of parameters {y: lies and x:
S Σ fly)) 2 211 1-
1
uill2 2 f()) 11:1-
| --ll where y
Σ yi.
S
Proof. Since ʃ is 6-weakly convex, o(z) == f(z) + 111 0 is convex. This proves the claim 1 [S] Eies (zy:) I (()).
since
L
20

---

Under review as a conference paper at ICLR 2021
F ANALYSIS OF MIMESGD (WITHOUT MOMENTUM)
Let us rewrite the MimeSGD update using notation convenient for analysis. each round t, we sample clients St such that |S*| S. The server communicates the server parameters x' as well as the average gradient across the sampled clients c²-l defined as
d'' S
Σ fi(a"'1). iES'
(11)
Note that computing c/-l itself requires two rounds of communication, But from a theoretical view- point, this only changes the communication rounds required by constant factor and hence ignore this issue. Practically, we recommend using MimeLite if this additional rounds of communication
issue.
Then each client i E St makes a copy y = x''1 and perform K local client updates. In each local client update & E [ the client samples dataset Çi and
v = yiAA (V (i,+) Vi((a' + ). After K such local updates, the server then aggregates the new client parameters as
(12)
x
Σ
vl,k
(13)
S
KES
Variance of update. Consider the local update at step k on client i, dropping superscript
Yo.k yi.k - ndd,, where d = Vi(( - ffi( (..)) + c.
Lemma 5. Given that assumptions (A1) and (A2) are satisfied, each client update satisfies
3G2 E| El < S + [y
=11 33llff((.
Proof. Starting from the definition of d and the relaxed triangle inequality,
||d,, 2== = V (i - fz (r,k) c|l
(yt. v (fi - (e-
3| (3 V.f (x,) + (z) Vf(y,--))l +
<55 yi, 2=|2 3|/e (3]]=)))(+++3//vv...)>)
We used Lemma 3 to bound the first term. Taking expectations both sides to bound the second term via (A1) yields the lemma.
Distance moved in each round. We show that the distance moved by client in each round during the K updates can be controlled. To further reduce the burden of notation, we will drop he subscript i, k and refer - simply as y and yi as yt.
Lemma 6. For update following (12) for n I 4кσ satisfying (A1) and (A2), we have any step k,
1
G2 + 6kn'
E Ely - x < (1 + 1 21| 6Kn
S
Proof. Starting from the update (12) and the relaxed triangle inequality Lemma with K K21 1.
1y 2 2
Ely bd
x||
2) Fl
x 1| c)n Eldd1
G²
< (1 + b Elly 1012 + 3(1 K)n
+3 + K) y 2122+30 + 3(1 {)
S
G
< π + 6 E y 2 + 6K,,
+ 6kT)
5
The second to last step used the variance bound in Lemma 5. The proof now follows from the restriction step-size since 2, 1.
21

---

Under review as a conference paper at ICLR 2021
Progress in one client update. We now have the tools required to keep track of the progress made in one round.
Lemma 7. For any constant u 2 O and each step of MimeSGD with step size π s
min
T2pK) and given that (A1)-A3) hold, we have
I8I 7565K
(255K LI
E S A A/.E
28
where we define
1 E[s(((x) 5 ((1+ + 3. K-k
Ellyi, - -xx- 1x2-113 and
A Elf(y( 8(1 - (un) (1) + k
K-++1
E| - 11-
Proof. The assumption that f is L-smooth implies quadratic upper bound (10). Using this in our case, we have
Lnj²
E] - d)} +
El|d]2
2
Li
=<EE(V)() V)i(y: Vii( + c))
Eld|
T2
Let us examine the terms Ti and To separately. By our variance bound Lemma 5, we have that
3Tm/202
3ln²
T2
ly
2
[Vf()))
28
2
2
To simplify Ti, the biggest obstacle is that f O) + Vyly) since y itself depends on the sampling of the client i. Only the server gradient is unbiased and Elcl v (()). Instead we will use the similarity of the functions as in Lemma 3:
T = 7 EK fzy) Vf () ()
< 2 [(V](1)) + - [V(l: Vii V( (
no²
- EIVS(1))
+ 2 Cly =
The first inequality above used that for any a, b, the following holds -2ab (a - b) a²
- b² <
(a - b)2 The second used the similarity Lemma 3. Combining the terms Ti and To together,
have
Elf(y+)] [lf())) (3Lm² 7) E(V (186 + +L
Ely y x 12 + 3L
25
To bound the distance between y and æ, we use Lemma 6 multiplied on both sides by 0(1 + xx K-k
Note that 8 < 5(1 ) K- < 21d. This gives for any constant μ 2 0
5(1
3 K
K Ely x1|² ك 6 61 + xx
K-k
x|| =ll-=11++66(( 66666 K) n
K-k G2
5 5
K-k
6K6
1+
3
(1)))1
K)
K K-(k-1)
K-k 02
G2 5
S 611 -mm)
1 1+
3
1g - xi|2 + 6K6 [1
6Kō
3
11
k
K
K
K-k
8
6K9 1 +
x
3
7/ ((]((( 2 + (I (1 x))
(umo
1 x
G2
<6(1+ + 8 K-
(K- K-
|18 - x + 126K 8n1
+ 126K
5
+ (21mm8
1) 1
22

---

Under review as a conference paper at ICLR 2021
The second inequality from the last used that 1+2/K < (1+3/K)((- (142/K)+
(1 + (/33))mm) This is true by our restriction that 17 < 42μK which implies (1+3/K) < 4pm <
1/(10K) and so that (1/K - (1 + 3/K)jnn) > 0. Adding the two bounds, we get the following recursion
Elf(y 5
K-k
y
x S [(f(v)) +0(1+ 3
K- K-(k-1)
(1 (1- (1-1m)llyy ==1
A
=:Ai,e-1
(252Kōr)²
3L1 -1)
2 + (118² + 3L + 442un1)
Flyy - x012
2
K
(3ln²² + 252
+
25
Now, note that our constraint on the step-size 17 S min( 18T 7566K implies that +3 <
and K (18²² +3Ln2 42pm6)) < 26. Plugging this into the above bound and recalling that 8 < L 2 finishes the proof.
Convergence for PL strongly-convex functions. We will unroll the one step progress Lemma to compute linear rate.
Theorem Suppose that (AI)-(A) are satisfied for μ > Then the updates of MimeSGD with satisfy
step-size 1j (nin(//maxx: ①
w lTTK)
for Mhuax min
1 187.
7568K 42μk
T.G2
F
1 18L + 7568K + 42μK TK) )
exp
HTS
Mmax
where we define F = f(x)) fr, V is chosen to be v St uniformly at random, and the output n be y with probability proportional (1 2y K
Proof. Note that by PL strong convexity (A4), we have
HIVS())
2
| + THe ((()) - (
4
Using this, we can tighten the one step progress Lemma 7 as
8 ElIVS(via 1)12
(1-14) 1m77
1)
W1)
(1 + K
K-k+1
E Ellyii
x
212
4
(11-147) (1) )
3 K-k
(255KLT)))
E(S(vi,) 1r]] +
El
-112
25
=:01
Now take a weighted sum over the steps & using weights (1- - 2 K
EI < T - (1 - 42) K
D4 K-E
K.K
+
25
kE[K]
kE[K]
By the initialization vi, -11 and hence $ E [f(z==1)= f and further by the averaging Lemma 4, we have
M k Ef(x') ft
S
iES
23

---

Under review as a conference paper at ICLR 2021
Hence, averaging over the clients get the round progress lemma
1
Σ
S - - (1--4) (E f(x') - f*)
85 REIK iES"
L2
(1 (1- - -
28
KE[K]
Now further taking a weighted average over the rounds / E [7] with weights proportional to (1
yk gives
7
(1- (OI--))KT-A E C =
8$
(E|T] kE[K] FES!
EE MMyKT (K (255K
25
EE[T] KE[K]
Finally, choosing the right step size, similar to Lemma 23 of (Karimireddy et al., 2020) yields the desired rate.
Convergence for general functions. sublinear rate.
We will unroll the one step progress Lemma 7 to compute
Theorem III. Suppose that (A1)-(A3) are satisfied. Then the updates of MimeSGD with step-size satisfy
7 min Mmax V255K-LGTT
VFS
for Tamax
min 181 7568K)
GVLF
(L+8K)F
ΣΣ KTS
El EII
VTS
TK
KE[R] iESt
where we define F = [(x=)) l
Proof. By summing over the equations from Lemma 7 for all local steps in one round we obtain
K
255K Lo
1
EIV
< A - 1 +
2
28
A=1
By the initialization v x hence A = Al.0 E[f(('-1)) for all i,j E St, Furthermore,
by Lemma 4
AL.x 2 E[f(x")] + - E[f(x') A/
S
iES.
This means that we can keep unrolling over all rounds, obtaining
255T 125TK2LLn2
17
ΣΣ Σ
E|V < 41,0 -A,,x +
25
25
1=1
iES
By noting 4/.0 theorem follows.
AT M (s(x9) (*) (F[[((²² f*) s F and the choice of the stepsize the
24

---

Under review as a conference paper at ICLR 2021
G ANALYSIS OF MIMEMVR (WITH MOMENTUM BASED VARIANCE REDUCTION)
In this section we see how to use momentum based variance reduction (Cutkosky and Orabona, 2019) to reduce the variance of the updates and improve convergence. It should be noted that MVR does not exactly fit the MIME framework (BASEALG) since it requires computing gradients at two points the same batch. However, it is straightforward to extend the idea of MIME to MVR will now do. We use MVR as a theoretical justification for why the usual momentum works well in practice. An interesting future direction would be to adapt the algorithm and analysis of (Cutkosky and Mehta, 2020), which does fit the framework of MIME.
MimeMVR algorithm. Now, we formally describe the MimeMVR algorithm. each round t, sample clients 8 such that |S = S. The server communicates the server parameters xt-1, the momentum f-1 and the average gradient across the sampled clients cn defined as
d = Σ 3) S iES
(14)
Note that both c² and m use gradients and parameters from previous rounds (different from the previous section).
Then each client E S makes a copy v = x'~1 and perform local client updates. In each local client update h E [K] the client samples dataset Ci and
yi,x vi ndh,e where
d = a al................... -i Vf((2''' V c ((-aa))))
(15)
+ (1- (1-a)VVV())r-ii
vf(x'=1;;,)))
After K such local updates, the server then aggregates the new client parameters as
x²
S Σ J
(16)
The momentum term updated at the end of the round for 2 as
m' E Vfz(x(-1)) + (1 a)m''! (1-a){ Ejess (f
SGDm
correction
(17)
As see, the momentum update of MVR can be broken down into the usual SGDm update, and correction. Intuitively, this correction term is very small since fi is smooth and '--1 Rs x² 1-2 Another way of looking the update (17) to note that if all functions for any j, k, then (17) just becomes the usual gradient descent. Thus MimeMVR tries to maintain an exponential moving average of only the variance terms, reducing its bias. We refer to (Cutkosky and Orabona, 2019) for more detailed explanation of MVR.
identical i.e. ʃq - fx
We compute the variance of the server momentum m Define the variance term Vt m Vf(w Then its expected norm be bounded follows.
Momentum variance bound.
Lemma 8. For the momentum update (17), given (A1) and (A2), the following holds for any [0, 1] and V == m V -1)
E < (1 - E 112 + 28 El|z-~1 2 2a²0²
s
25

---

Under review as a conference paper at ICLR 2021
Proof. Starting from the momentum update (17),
V (1 - a)Vi=1
+ (1-a)
1 Σ
([ -V V/(00'1)
+ V
S
JES
Σ
(V v
+ S
JES
Now, the term V does not have any information from round and hence is statistically indepen- dent of the rest of the terms. Further, the rest of the terms have mean O. Hence, we can separate out the zero mean noise terms from the V following Lemma and then the relaxed triangle inequality Lemma to claim
ElIV* S (1 a)² E||
+ 2(1 a)²
Σ (Vij(=''1) -V5((((-2) Vs(x'11) + V5(x
S
JESL
1
+ 2a
[Vfs(x -
s
jESt
20² (1 a) [E| + 2(1 a)2 |=:" 100--1--'200 2112 +
S
The inequality used the Hessian similarity Lemma to bound the second term and the heterogeneity bound (AI) to bound the last term. Finally, note that (1 a)² < (1 a) I 1 for E [0, 1].
Update variance bound. Now we examine the variance our update each local step di
Lemma 9, For the client update (15), given (A1) and (A2), the following holds for any E [0, 1]
Elidi V s 3 3 38 El 21-2112 + 3a²
5
Proof. Starting from the client update (15), we can rewrite it as
d -f (1 a)Vi-1
+ (V.i Vii(a' 2; -.,)) -d.))-V((iix-- V/
Σ Vf - vf
+ a
JES
We can use the relaxed triangle inequality Lemma to claim
El|di
3(1 a)² V
+ 3(1 - -3(1-a) (V - - - v)(x=-2)) ||²
+ 3a
Σ Vf -V
JESt
Ga²
< /3
38
21-2112
S
The last inequality used the Hessian similarity Lemma 3 to bound the second term and the hetero- geneity bound (A1) bound the last term. Also, a)² since E [0, 1].
n
26

---

Under review as a conference paper at ICLR 2021
Distance moved in each step. We show that the distance moved by client in each step during the client update can be controlled.
Lemma 10. For MimeMVR updates (15) with v holds
S 6kō and given (A1) and (A2), the following
G2 Lin I (+ { K A -i + Ka
Kr²
+ 187 K K
V
S
where we define 1i,x max [Ell3'. x² 211 Elly
11 - -11))
Proof. Starting from the MimeMVR update (15) and the relaxed triangle inequality with e 2K,
El x x222112
El
ndi
21
1
+
E|| Ellyi, - x 212 + (2K + 1)n² El|d
2K
1
+
EllyiR-11 - x' 21 + 6KT Elldi - Vf(v(
2K
6Ki
FIIVS
1 +
1
+
18K 882
El - x.' 212
2K
18Kv 18K E||"
+ GKTKK x-11
S
The last inequality used the update variance bound Lemma We can simplify the expression further since I 6Ka implies 18K s 2k Similar computations for Elly xx yield the lemma.
Progress in one step. Now we have all the tools required to compute the progress made in each round.
Lemma 11. For any step of MV with step size n S min(t T and momentum parameter 153 K2, Then, given that (A1)-A3) hold, have
[[f(v(,)] + 37 ElIV*122
81102 K
(
2
11
A L
a
K
37
8n82K K
2
K-(k-1)
s <E(f(g)R-1))
E||v-1112
1 +
1
K
111361 K2
7 "IVVf 11) 11
S
Proof. The assumption that f is L-smooth implics a quadratic upper bound (10).
Ln
f(v(,)) (00-f((/aaa)) - 7. VV wtw.fyl.ee11)) d
+
|d 1|?
2
2 v /f )1
Laj²
|
ブ ||d -
!
2
The second equality used the fact that for any a, b. -2ab = (a - 6)2 a² b². The second term be removed since 77 t. Taking expectation both sides and using the update variance bound Lemma 9,
3na2
E E((( - E. ( S - E [EV[(((
25
3n
E|
3155 2
El
El
x 2222
27

---

Under review as a conference paper at ICLR 2021
Multiplying the momentum variance bound Lemma 8 by 3n have
3 31 El| 6nd2
E| x 212 GnaG²
E a v a
-3n E1|tt11111
Sn K (1+ ))*²**
K-
We will also multiply the distance bound Lemma 10 by
Note that for any
K and [K], have 1 1 (1
x) K-k
I 8. Then get
8n6² K
(1
2 K-k
8n0²K K
(1 1+
2 K-(k-1) 1)
8n62
1
L
A k-
A k-1
K
G
1152 K2
|| 2 384 [V
+ 1152n
S
a
a
where recall that we defined A == max Ellui, Combining the three inequalities together, we get
21
E V
1|-
x
x 211)
E ((( + 31 E|IV* 811²2 1
2
K-k
A i
a
K
37
8168 K
(
2 K-(K-1)
<Efly 1) +
1
a
a
K
3841 384 K2
!
E
2
3a
on
1152 28 k2
+ +6+
2
s
31
1152m
822
3n/
E||V--112
2
3
6
8
+
n
A
1
2
Note that x2
3r since we defined = 15 K2 Further, S 1 when defined this
way since we assumed S Tol Similarly, the definition of implies that 32 we can simplify the above expression as
7 Thus,
4
IL (( + 37 E|Iv*12 8n0²K K
2 K
K-k
1+
A/K
37
81682 K E
2
K-(k-1)
S E EE
+
E
1
Dix-1
a
K
111361 K22
4
S
This proves the lemma.
Progress in one round. in full round.
Let us sum over all the steps within round compute the progress made
Lemma 12. For any round of Mi with step size 7 S min( +
1
and momentum param-
T0ok
eter a = 2 82 K?. Then, given that (A1)-(A3) hold, we have
El & Φc-1 T.. 11136 82 K2G²
1
Σ
4KS
S
RE(KT,JES
where we define the sequence
+ 3nK Φh = > E[f(x'))
8n62
E||v*|2+
E
El
x 111
a
28

---

Under review as a conference paper at ICLR 2021
Proof. We start by summing Lemma 11 over the client updates
1)I s E(S()())) + 3nK
+ 8n6²K K
K
Σ
EIVVS(((
E
El
1 K
Ai,
kE[K]
37
8168
E||
A K
a
111361 111361782 G²
5
Recall that we defined A == max EEllvi. xx" cause yi,o - x1-1, we can simplify
x² E| - 211) Be-
Elf(y(,)] 8n82 K
8n02K K
A S [[f(x
E| - -xx--xx"11 21
Then by the averaging Lemma 4, we have
1
Σ Elf(y))))
S1z6 K
A j.
E[ff(U), + El x1,
5
s
JES
jES
Snjō K
2
E[f(z')]
El Fl x²
So by averaging our recursion over the sampled clients, and diving our summation over the updates by K, we get
7
Σ
S # Elf(x'
37
Ellv- 81662
112
El|0
x 21
4KS kE[K],jES
a
=
F[((a')) 3nK
8115
FI|V*122
F F|
2x. 112
a
a
11
S
Theorem IV (non-convex convergence of MimeMVR). Let us run MimeMVR with step size S
1/3 + wook and momentum parameter a 1536n/ 22 82. Then, given that
min
I5K
TF
(A3) hold, we have
((
(1+8)) F 2/3 SSTTT
(L (L+8K)F
E
KST
KT
[E[T]KE[K] KE(K)JES
where define F = - ft
Proof. Unroll the one round progress Lemma 12 and average over T rounds to get
4(40 Φ7)
111367788 R??
E Σ El](
KST
nKT
S
(E[T] JES
A(f(x²) ff)) nKT
11136 K2
S
Our choice of step size now yields the desired rate.
29


Reviews:
[{"review_id": "wxO2rsE1uJ4", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "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": "The paper proposes a new framework for solving federated learning. The authors consider a specific setting that there are many clients, and each client is allowed to compute the full gradient. The authors claim that the current setting’s main issue is the client drift, and the proposed framework can reduce such an issue and thus achieve a faster convergence rate. Here are my main concerns of the current paper:\n1. The client drift issue is not clearly defined in this paper.\n2. The comparisons in Table 1 are not fair since different methods have different assumptions. For example, the authors claim that the proposed MimeSGD algorithm can reduce communication costs compared with FedSGD. However, MimeSGD depends on the extra Hessian assumptions (A2).\n3. Assumption A2 seems to be a very strong assumption. Intuitively, it contains two variances: one from the gradient dissimilarity and another from the within-client variance. Therefore, why can we expect that delta is very small? In addition, is the improvement of your framework comes from this assumption?\n4. It seems that the algorithm with theoretical guarantees is different from the one implemented in experiments. To validate your results, the experiments should use the same algorithm.\n5. The parameter beta in Theorem 1 seems to be a very small value. However, the experiments suggest a very large beta. This seems to be inconsistent with your theorem. Is this due to the different algorithms you implement in experiments?\n6. What is the definition of x^out? \n7. The proposed framework requires additional communications due to the correction term c. Did you consider this extra communication cost in your experiments?\n8. In the experiments, since each client performs several epochs of updates, the comparison with the server-only method is not fair since it only performs one update each communication.\n9. In Figure 2, why is there a spike for your Mime method?\n10. Typos: relative slow-> relative low. vour-> our.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "n0g9Zn_XBm0", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "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 of the paper:\nThis paper introduces the MIME framework which can adopt standard centralized SGD methods in federated learning environments and this framework handles the well known \"client drift\" problem. It is done by cleverly using server statistics such as momentum in local devices. Ample empirical evaluations are provided in order to validate the theoretical claims.\n\nQuality:\nAlthough I did not check the proofs thoroughly, the presented theoretical results looks natural and believable.  Ample experiments are provided to validate the claims. Clarity of the paper is good in general but there are places that can be improved. The quality and the presentation of the paper can be improved significantly. \n- Most of the remarks under theorem 1 about the table 1 are redundant. The same remarks are given before the table and sometimes multiple times. For example, the claims about the case where $\\delta << L $ and claims about SCAFFOLD. These can be clearly stated in contributions.\n- In theorem 1, $\\mathbb{E}[||\\nabla  f (x^{out})||^2] \\le \\epsilon$, expectation over what? I believe it is $i$?\n- In assumptions A1, A2, the bound on gradient is defined over the expectation of local datasets $i$ but the Hessian bound is for all $i$. Is there any specific reason for this? Or is this a standard assumption?\n- Typos: Last line in the first paragraph of Related work, there is an extra \"?P:\". Page 3 before equation 1, \"vour\" -> \"our\".  In equation 5, nothing written after \"where\".\n\n\nOriginality and significance:\nI believe this paper adds decent contributions toward understanding of federated learning and and tackles the \"client drift\" problem in a simple way.\n\n\nPros:\n- Nice to have a flexible framework that can adopt standard centralized methods with similar guarantees.\n- Extensive experiments validates the provided theoretical claims.\n\nCons:\n- This work is incremental on the paradigm of using momentum in order to improve distributed SGD methods.\n\nOther comments:\nThe paper titled \"On the Linear Speedup Analysis of Communication Efficient Momentum SGD for Distributed Non-Convex Optimization\" seems to be discussing a similar approach that uses momentum updates in locally. Can the authors distinguish between these two results? \n\n\n===========================================================================================================\n\nAdded after author response:\n\n--------------------------------------------------------------------------------------------\n\nI have read the author responses and other reviews. I believe authors have a sufficiently addressed my concerns regarding the quality of the paper. I understand the improvement in using momentum in this way and I think it is decent contribution. However, considering this improvement and the concerns raised by other reviewers, I maintain my score.", "extra_scores": {}, "extra_text": {}}, {"review_id": "hFbsU8M2glW", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 4, "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": "In the paper, authors proposed to utilize the global statistics in local client updates to mitigate the client drift problem. The algorithm mine and minelite are very interesting and look reasonable to me. They also prove the convergence rate of the proposed algorithm and evaluate its empirical performances on synthetic and real federated learning simulations.\n\nThe following are my concerns:\n1) At page 2, analysis of fedavg, the last sentence \"technically, it was a slightly modified FEDAVG version? P: FEDAVG as it is understood encompasses methods with server lr.\" looks very strange to me.\n2) typo in the sentence of eq (1) \"vour\".\n3) A more straightforward compared algorithm is fedavg using sgdm on clients. \n4) I am confused about the performance of fedsgdm and fedadam in the paper. According to \"Adaptive federated optimization\" paper, the accuracy of fedsgdm and fedadam achieve about 85%. However, in the paper, the results are far from that. \n5) The same question also applies to cifar100 experiments. Besides, it looks like that mimeadam stops converging after 600 rounds, however, fedadam will continue converging to about 50% according to the  \"Adaptive federated optimization\" paper.\n\n====\nAfter rebuttal:\n\nI read all reviews and rebuttals, especially the argument about Lemma 7. I agree with reviewer 2 that the Lemma 7 is not correct. The assumption of $\\mu\\eta < \\frac{1}{42k}$ contradicts with the requirement $\\mu\\eta > \\frac{1}{k}$ of term 4.  This error needs to be resolved before acceptance. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "QtYl72JOtKf", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 4, "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 proposes a way to apply various variance reduction/momentum based method to the federated learning scenario, especially when there is distribution drift among the clients. The main claim of this paper is that the global statistics (momentum, control variance et al) should be update at the server side only, which is helpful to reduce the bias of these terms. The paper also provides convergence analysis for their methods and attains the best convergence result for their MimeMVR method.\nHowever, I find a key wrong in theoretical analysis. Specifically, in the proof of Lemma 7, the last second inequality in page 21 is wrong.  Thus I have difficulty judging the contribution of this paper.  \n\nSome comments are given as follows:\n\n1.\tIn the proof of Lemma 7, when the authors bound the distance between y and x, how do you get at the last second inequality in page 21, it seems to rely on: 1 + 2/k < 1 + 2/k - mu*eta* 3/k, which obviously not hold.\n2.\tFor the MimeMVR, the update rule for d^t_{i,k} in the equality (14), there is a  -1 term, which looks strange.  Why does this term exit?\n3.\tIn the proof of Lemma 8, the second equality at page 25 should be inequality.\n4.\tIn the proof of Lemma 9, there are typos using the brackets in the third equality at page 25.\n\n\n---------------------------------------------------------------------------------------------------\nAfter Rebuttals\n---------------------------------------------------------------------------------------------------\nThanks for your responses. I still have doubts about the proof of lemma 7. You not only change from '2' to '3', you also add a new term (the fourth term). I admit that $1 + 2/k < (1 + 3/k)(1 – \\mu*\\eta)$ holds under your assumption, but you also need: $\\mu*\\eta> 1/k$, i.e. let the fourth term be greater than zero, to make the inequality hold, however, this contradicts the previous assumption. So I still believe this inequality does not hold. \n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The submission introduces a theoretically-justified solution to 'client-drift' in federated learning.  Generally, the reviewers agreed that this could be a strong paper, and several of their minor concerns have been addressed in the rebuttal. Unfortunately, a major issue raised during the discussion was the correctness of Lemma 7. Despite the extra clarifications provided by the authors, Reviewer2 still believes it is incorrect and R4 also sided with them.  As theoretical analysis is the major contribution of this work, we have to reject the submission. I would strongly encourage the authors to fix the issue (or clarify the proof), and resubmit it to one of the upcoming top ML conferences. 
Decision: Reject