Breaking the centralized barrier for cross-device federated learning∗
Year: 2022

Abstract:
Federated learning (FL) is a challenging setting for optimization due to the heterogeneity of the data across different clients which can cause a client drift phenomenon. In fact, designing an algorithm for FL that is uniformly better than simple centralized training has been a major open problem thus far. In this work, we propose a general algorithmic framework, MIME, which i) mitigates client drift and ii) adapts an arbitrary centralized optimization algorithm such as momentum and Adam to the cross-device federated learning setting. MIME uses a combination of control-variates and server-level optimizer state (e.g. momentum) at every client-update step to ensure that each local update mimics that of the centralized method run on i.i.d. data. We prove a reduction result showing that MIME can translate the convergence of a generic algorithm in the centralized setting into convergence in the federated setting. Moreover, we show that, when combined with momentum-based variance reduction, MIME is provably faster than any centralized method–the first such result. We also perform a thorough experimental exploration of MIME’s performance on real world datasets (implemented here).

heading: 1 Introduction
text: Federated learning (FL) is an increasingly important large-scale learning framework where the training data remains distributed over a large number of clients, which may be mobile phones or network sensors [38, 37, 43, 44, 28]. A server then orchestrates the clients to train a single model, here referred to as a server model, without ever transmitting client data over the network, thereby providing some basic levels of data privacy and security.
Two important settings are distinguished in FL [28, Table 1]: the cross-device and the cross-silo settings. The cross-silo setting corresponds to a relatively small number 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 [25]. Thus, in that setting, we may never make even a single pass over
∗This work was also appears under the alternative title “Mime: Mimicking Centralized Stochastic Algorithms in Federated Learning” [31].
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
the entire clients’ data during training. The cross-device setting is further characterized by resourcepoor clients communicating over a highly unreliable network. Together, the essential features of this setting give rise to unique challenges not present in the cross-silo setting. In this work, we are interested in the more challenging cross-device setting, for which we will formalize and study stochastic optimization algorithms. Importantly, recent advances in FL optimization, such as SCAFFOLD [32] or FedDyn [1], are not anymore applicable since they are designed for the cross-silo setting.
The problem. The de facto standard algorithm for the cross-device setting is FEDAVG [43], which performs multiple SGD updates on the available clients before communicating to the server. While this approach can reduce the frequency of communication required, performing multiple steps on the same client can lead to ‘over-fitting’ to its atypical local data, a phenomenon known as client drift [32]. This in turn leads to slower convergence and can, somewhat counter-intuitively, require larger total communication [69]. Despite significant attention received from the optimization community, the communication complexity of heterogeneous cross-device has not improved upon that of simple centralized methods, which take no local steps (aka SERVER-ONLY methods). Furthermore, algorithmic innovations such as momentum [59, 14], adaptivity [35, 75, 77], and clipping [71, 72, 76] are critical to the success of deep learning applications. The lack of a theoretical understanding of the impact of multiple client steps has also hindered adapting these techniques in a principled manner into the client updates, in order to replace the vanilla SGD update of FEDAVG.
To overcome such deficiencies, we propose a new framework, MIME, that mitigates client drift and can adapt an arbitrary centralized optimization algorithm, e.g. SGD with momentum or Adam, to the federated setting. In each local client update, MIME uses global optimizer state, e.g. momentum or adaptive learning rates, and an SVRG-style correction to mimic the updates of the centralized algorithm run on i.i.d. data. This optimizer state is computed only at the server level and kept fixed throughout the local steps, thereby avoiding overfitting to the atypical local data of any single client.
Contributions. We summarize our main results below. • MIME framework. We formalize the cross-device federated learning problem, and propose a
new framework MIME that can adapt arbitrary centralized algorithms to this setting. • Convergence result. We prove a result showing that MIME successfully reduces client drift.
We also prove that the convergence of any generic algorithm in the centralized setting translates convergence of its MIME version in the federated setting. • Speed-up over centralized methods. By carefully tracking the bias introduced due to multiple local steps, we prove that MIME with momentum-based variance reduction (MVR) can beat a lower bound for centralized methods, thus breaking a fundamental barrier. This is the first such result in FL, and also the first general result showing asymptotic speed-up due to local steps. • Empirical validation. We 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 indeed converge 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 [78] derived an analysis with asymptotic convergence. Sharper and more refined analyses of the same method, sometimes called local SGD, were provided by [56], and more recently by [57], [47], [34], and [70], for identical functions. Their analysis was extended to heterogeneous clients in [68, 74, 32, 34, 36]. [11] derived a tight characterization of FedAvg with quadratic functions and demonstrated the sensitivity of the algorithm to both client and server step sizes. Matching upper and lower bounds were recently given by [32] and [69] for general functions, proving that FEDAVG can be slower than even SGD for heterogeneous data, due to the client-drift.
Comparison to SCAFFOLD: For the cross-silo setting where the number of clients is relatively low, [32] proposed the SCAFFOLD algorithm, which uses control-variates (similar to SVRG) to correct for client drift. However, their algorithm crucially relies on stateful clients which repeatedly participate in the training process. FedDyn [1] reduces the communication requirements, but also requires persistent stateful clients. In contrast, we focus on the cross-device setting where clients may be visited only once during training and where they are stateless (and thus SCAFFOLD and FedDyn are inapplicable). This is akin to the difference between the finite-sum (corresponding to cross-silo) and stochastic (cross-device) settings in traditional centralized optimization [39].
Comparison to FedAvg and variants: [26] and [67] observed that using server momentum significantly improves over vanilla FEDAVG. This idea was generalized by [49], 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. We henceforth refer to this meta algorithm as FedAvg. FedAvgSGD, FedAvgMom, FedAvgAdam denote specific instantiations of the server optimizer in FedAvg with SGD, Momentum or Adam. MIME, on the other hand, ensures that every local client update resembles the optimizer e.g. MIME would apply momentum in every client update and not just at the server level. Beyond this, [40] proposed to add a regularizer to ensure client updates remain close. However, this may slow down convergence (cf. Fig. 5 and [32, 66]). Other orthogonal directions which can be combined with MIME include tackling computation heterogeneity, where some clients perform many more updates than others [66], improving fairness by modifying the objective [44, 41], incorporating differential privacy [20, 2, 61], Byzantine adversaries [48, 65, 30], secure aggregation [8, 24], etc. We defer additional discussion to the extensive survey by [28].
Momentum based variance reduction. Initial optimal methods for stochastic non-convex optimization like SPIDER [17] and SARAH [46] required intermittently computing very large batch gradients. Subsequently, it was shown that momentum based variance reduction (MVR) methods obtained a similar optimal rate without needing such large batch gradient computations [62, 14]. Momentum is an exponential moving average of many stochastic gradients and so it has much smaller variance than the stochastic gradients themselves. However, because these gradients are computed at different parameters it also has a bias. MVR adds a small additional correction term which significantly reduces this bias and provides improved rates.

heading: 2 Problem setup
text: This section formalizes the problem of cross-device federated learning [28]. Cross-device FL is characterized by a large number of client devices like mobile phones which may potentially connect to the server at most once. Due to their transient nature, it is not possible to store any state on the clients, precluding an algorithm like SCAFFOLD. Furthermore, each client has only a few samples, and there is wide heterogeneity in the samples across clients. Finally, communication is a major bottleneck and a key metric for optimization in this setting is the number of communication rounds.
Thus, our objective will be to minimize the following quantity within the fewest number of clientserver communication rounds:
f(x) = Ei∼C [ fi(x) := 1
ni ni∑ ν=1 fi(x; ζi,ν) ] . (1)
Here, fi denotes the loss function of client i and {ζi,1, . . . , ζi,ni} its local data. Since the number of clients is extremely large, while the size of each local data is rather modest, we represent the former as an expectation and the latter as a finite sum. In each round, the algorithm samples a subset of clients (of size S) and performs some updates to the server model. Due to the transient and heterogeneous nature of the clients, it is easy to see that the problem becomes intractable with arbitrarily dissimilar clients. Thus, it is necessary to assume bounded dissimilarity across clients. (A1) G2-BGV or bounded inter-client gradient variance: there exists G ≥ 0 such that
Ei∼C [‖∇fi(x)−∇f(x)‖2] ≤ G2 , ∀x .
Next, we also characterize the variance in the Hessians.
(A2) δ-BHV or bounded Hessian variance: Almost surely, the loss function of any client i satisfies
‖∇2fi(x; ζ)−∇2f(x)‖ ≤ δ , ∀x . This is in contrast to the usual smoothness assumption that can be stated as:
(A2*) L-smooth: ‖∇2fi(x; ζ)‖ ≤ L , ∀x , a.s. for any i. Note that if fi(x; ζ) is L-smooth then (A2) is satisfied with δ ≤ 2L, and hence (A2) is weaker than (A2*). In realistic examples we expect the clients to be similar and hence that δ L. In addition, we assume that f(x) is bounded from below by f? and is L-smooth, as is standard.

heading: 3 Mime framework
text: In this section we describe how to adapt an arbitrary centralized optimizer (referred to as the “base” optimizer) which may have internal state (e.g. momentum) to the federated learning problem (1) while ensuring there is no client-drift. Algorithm 4 describes our framework. We develop two variants, MIME and MIMELITE, which consist of three components i) a base optimizer we are seeking to mimic, ii) the global (server) optimizer state computation, and iii) the local client updates.
Algorithm 1 Mime and MimeLite
input: initial x and s, learning rate η and base optimizer B = (U ,V) for each round t = 1, · · · , T do
sample subset S of clients communicate (x, s) to all clients i ∈ S communicate c← 1|S| ∑ j∈S ∇fj(x) (only Mime)
on client i ∈ S in parallel do initialize local model yi ← x for k = 1, · · · ,K do
sample mini-batch ζ from local data gi ← ∇fi(yi; ζ)−∇fi(x; ζ) + c (Mime) gi ← ∇fi(yi; ζ) (MimeLite)
update yi ← yi − ηU(gi, s) end for compute full local-batch gradient∇fi(x) communicate (yi,∇fi(x))
end on client s ← V ( 1 |S| ∑ i∈S ∇fi(x), s ) (update optimizer state) x← 1|S| ∑ i∈S yi (update server parameters)
end for
Base optimizer. We assume the centralized base optimizer we are imitating can be decomposed into two steps: an update step U which updates the parameters x, and a optimizer state update step V(·) which keeps track of global optimizer state s. Each step of the base optimizer B = (U ,V) uses a gradient g to update the parameter x and the optimizer state s as follows: x← x− η U(g, s) , s← V(g, s) . (BASEOPT)
As an example, consider SGD with momentum. The state here is the momentum mt and uses the following update steps:
xt = xt−1 − η ((1− β)∇fi(xt−1) + βmt−1) ,
mt = (1− β)∇fi(xt−1) + βmt−1 . Thus, SGD with momentum can be represented in the above generic form with U(g, s) = (1 − β)g + βs and V(g, s) = (1 − β)g + βs. Table 5 in Appendix shows how other algo-
rithms like Adam, Adagrad, etc. can be represented in this manner. We keep the update U to be linear in the gradient g, whereas V can be more complicated. This implies that while the parameter update step U is relatively resilient to receiving a biased gradient g while V can be much more sensitive.
Compute optimizer state globally, apply locally. When updating the optimizer state 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 optimizer state remain unbiased and representative of the global function f(·). At the end of the round, the server performs
s← V ( 1 |S| ∑ i∈S ∇fi(x), s ) ,
∇fi(x) = 1ni ∑ni ν=1∇fi(x; ζi,ν) . (OPTSTATE)
Note that we use full-batch gradients computed at the server parameters x, not client parameters yi.
Local client updates. Each client i ∈ S performs K updates using U of the base algorithm and a minibatch gradient. There are two variants possible corresponding to MIME and MIMELITE differentiated using colored boxes. Starting from yi ← x, repeat the following K times
yi ← yi − ηU(gi, s) (CLIENTSTEP)
where gi ← ∇fi(yi; ζ) for MIMELITE, and gi ← ∇fi(yi; ζ)−∇fi(x; ζ) + 1|S| ∑ j∈S ∇fj(x)
for MIME. MIMELITE simply uses the local minibatch gradient whereas MIME uses an SVRG
style correction [27]. 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.
Finally, there are two modifications made in practical FL: we weight all averages across the clients by the number of datapoints ni [43], and we perform K epochs instead of K steps [66].

heading: 4 Theoretical analysis of Mime
text: Table 1 summarizes the rates of MIME (highlighted in blue) and MIMELITE (highlighted in green) and compares them to SERVER-ONLY methods when using SGD, Adam and momentum methods as the base algorithms. We will first examine the convergence of MIME and MIMELITE with a generic base optimizer and show that its properties are preserved in the federated setting. We then examine a specific momentum based base optimizer, and prove that Mime and MimeLite can be asymptotically faster than the best server-only method. This is the first result to prove the usefulness of local steps and demonstrate asymptotic speed-ups.

heading: 4.1 Convergence with a generic base optimizer
text: We will prove a generic reduction result demonstrating that if the underlying base algorithm converges, and is robust to slight perturbations, then MIME and MIMELITE also preserve the convergence of the algorithm when applied to the federated setting with additinoal local steps. Theorem I. Suppose that we have G2 inter-client gradient variance (A1), L-smooth {fi} (A2*), and σ2 intra-client gradient variance (A3). Further, suppose that the updater U of our baseoptimizer B = (U ,V) satisfies i) linearity for a fixed state s: U(g1 + g2; s) = U(g1; s) + U(g2; s), and ii) Lipschitzness: ‖U(g; s)‖ ≤ B‖g‖ for some B ≥ 0. Then, running MIME or MIMELITE with K local updates and step-size η is equivalent to running a centralized algorithm with step-size η̃ := Kη ≤ 12LB , and updates
xt ← xt−1 − η̃ U(gt + et , st−1) , and st ← V(gt, st−1) , where we have
an unbiased gradient Et[gt] = ∇f(xt−1), with variance bounded as
Et‖gt −∇f(xt−1)‖2 ≤
{ G2
S MIME , G2
S + σ2 KS MIMELITE .
and finally a small error bounded as
1 B2L2η̃2 Et‖ et ‖ 2 ≤ { Et‖gt‖2 MIME , Et‖gt‖2 +G2 + σ 2
K MIMELITE .
Here, we have proven that MIME and MIMELITE truly mimic the centralized base algorithm with very small perturbations—the magnitude of et is O(η̃2). The key to the result is the linearity of the parameter update step U( · ; s). By separating the base optimizer into a very simple parameter step U and a more complicated optimizer state update step V , we can ensure that commonly used algorithms such as momentum, Adam, Adagrad, and others all satisfy this property. Armed with this general reduction, we can easily obtain specific convergence results. Corollary II ((Mime/MimeLite) with SGD). Given that the conditions in Theorem I are satisfied, let us run T rounds withK local steps using SGD as the base optimizer and output xout. This output satisfies E‖∇f(xout)‖2 ≤ for F := f(x0)− f?, G̃2 := G2 + σ2/K and
• µ-PL inequality: η = Õ (
1 µKT
) , and
T = Õ ( LG2 µS + LF µ log ( 1 )) MIME , Õ ( LG̃2
µS + LG̃ µ √ + LFµ log ( 1 )) MIMELITE .
• Non-convex: for η = O (√
FS LG̃2TK2
) , and
T = O ( LG2F S 2 + LF ) MIME , O ( LG̃2F S 2 + L2G̃F 3/2 + LF ) MIMELITE .
Table 1: Number of communication rounds required to reach ‖∇f(x)‖2 ≤ (log factors are ignored) with S clients sampled each round. All analyses except SCAFFOLD assume G2 bounded gradient dissimilarity (A1). All analyses assume L-smooth losses, except MimeLiteMVR and MimeMVR, which only assume δ bounded Hessian dissimilarity (A2). Convergence of SCAFFOLD
depends on the total number of clientsN which is potentially infinite. FEDAVG and MIMELITE are slightly slower than the server-only methods due to additional drift terms in most cases. MIME is the fastest and either matches or improves upon the optimal statistical rates (first term in the rates). In fact, MimeMVR and MimeLiteMVR beat lower bounds for any server-only method when δ L.
Algorithm Non-convex µ-PL inequality SCAFFOLDa [32] ( N S ) 2 3 L N S + L µ
SGD SERVER-ONLY [21] LG 2
S 2 + L
G2 µS + L µ
MimeLiteSGD≡ FedAvgSGD c LG 2 S 2 + L 2G 3/2 + L
G2 µS + LG µ √ + L µ
MimeSGD LG 2 S 2 + L
G2 µS + L µ
ADAM SERVER-ONLY [75]b L
−G2/S –
MimeLiteAdambc L √ S
−G2/S –
MimeAdamb L −G2/S –
Momentum Variance Reduction (MVR) SERVER-ONLY [14] LG√
S 3/2 + G
2
S + L
–
MimeLiteMVRd δ(G+σ) 3/2 + G 2+σ2 + δ
–
MimeMVRd δG√ S 3/2 + G 2 S + δ
– SERVER-ONLY lower bound [5] Ω (
LG√ S 3/2 + G 2 S + L
) Ω ( G2
S ) a Num. clients (N ) can be same order as num. total rounds or even∞, making the bounds vacuous. b Adam requires large batch-size S ≥ G2/ to converge [50, 75]. Convergence of FedAdam with client sampling is unknown ([49] only analyze with full client participation). c RequiresK ≥ σ2/G2 number of local updates. Typically, intra-client variance is small (σ2 . G2). d RequiresK ≥ L/δ number of local updates. Faster than the lower bound (and hence any SERVERONLY algorithm) when δ L i.e. our methods can take advantage of Hessian similarity, whereas SERVER-ONLY methods cannot. In worst case, δ ≈ L and all methods are comparable.
If we take a sufficient number of local steps K ≥ G2/σ2, then we have G̃ = O(G) in the above rates. On comparing with the rates in Table 1 for SERVER-ONLY SGD, we see that MIME exactly matches its rates. MIMELITE matches the asymptotic term but has a few higher order terms. Note that when using SGD as the base optimizer, MIMELITE becomes exactly the same as FEDAVG and hence has the same rate of convergence. Corollary III ((Mime/MimeLite) with Adam). Suppose that the conditions in Theorem I are satisfied, and further |∇jfi(x)| ≤ H for any coordinate j ∈ [d]. Then let us run T rounds using Adam as the base optimizer withK local steps, β1 = 0, ε0 > 0, η ≤ ε20/KL(H+ε0), and any β2 ∈ [0, 1). Output xout chosen randomly from {x1, . . .xT } satisfies E‖∇f(xout)‖2 ≤ for
T =  O ( LF (H+ε0) 2 ε20( −G̃2/S) ) MIME Adam , O ( LF (H+ε0) 2 √ S
ε20( −G̃2/S)
) MIMELITE Adam .
where F := f(x0)− f?, G̃2 := G2 + σ2/K.
Note that here ε0 represents a small positive parameter used in Adam for regularization, and is different from the error . Similar to the SERVER-ONLY analysis of Adam [75], we assume β1 = 0
and that batch size is large enough such that S ≥ G2/ . A similar analysis can also be carried out for AdaGrad, and other novel variants of Adam [42].

heading: 4.2 Circumventing server-only lower bounds
text: The rates obtained above, while providing a safety-check, do not beat those of the SERVER-ONLY approach. The previous best rates for cross-device FL correspond to MimeLiteSGD which is O(LG 2
S 2 + L2G 3/2
) [34, 36, 69]. While, using a separate server-learning rate can remove the effect of the second term [33], this at best matches the rate of SERVER-ONLY SGD O(LG 2
S 2 ). This is significantly slower than simply using momentum based variance reduction (MVR) as in in the FL setting (SERVER-ONLY MVR) which has a communication complexity of O( LG√
S 3/2 ) [14]. Thus,
even though the main reason for studying local-step methods was to improve the communication complexity, none thus far show such improvement. The above difficulty of beating SERVER-ONLY may not be surprising given the two sets of strong lower bounds known.
Necessity of local steps. Firstly, [5] show a gradient oracle lower bound of Ω( LG√ S 3/2 ). This matches the complexity of MVR, and hence at first glance it seems that SERVER-ONLY MVR is optimal. However, the lower bound is really only on the number of gradients computed and not on the number of clients sampled (sample complexity) [18], or number of rounds of communication required. In particular, multiple local updates increases number of gradients computed without needing additional communication offers us a potential way to side-step such lower bounds. A careful analysis of the bias introduced as a result of such local steps is a key part of our analysis.
Necessity of δ-BHD. A second set of lower bounds directly study the number of communication rounds required in heterogeneous optimization [6, 69]. These results prove that there exist settings where local steps provide no advantage and SERVER-ONLY methods are optimal. This however contradicts real world experimental evidence [43]. As before, the disparity arises due to the contrived settings considered by the lower bounds. For distributed optimization (with full client participation) and convex quadratic objectives, δ-BHD (A2) was shown to be a sufficient [54, 51] and necessary [6] condition to circumvent these lower bounds and yield highly performant methods. We similarly leverage δ-BHD (A2) to design novel methods which significantly extend prior results to i) all smooth non-convex functions (not just quadratics), and ii) cross-device FL with client sampling.
We now state our convergence results with momentum based variance reduction (MVR) as the basealgorithm since it is known to be optimal in the SERVER-ONLY setting. Theorem IV. For L-smooth f with G2 gradient dissimilarity (A1), δ Hessian dissimilarity (A2) and F := (f(x0) − f?), let us run MVR as the base algorithm for T rounds with K ≥ L/δ local steps and generate an output xout. This output satisfies E‖∇f(xout)‖2 ≤ for • MimeMVR : η = O ( min ( 1 δK , ( SF G2TK3 ) 1/3 )) , momentum β = 1−O( δ 2S2/3 (TG2)2/3 ), and
T = O ( δGF√
S 3/2 + G2 S + δF
) .
• MimeLiteMVR : η = O ( min (
1 δK , ( F Ĝ2TK3
)1/3 )) , momentum β = 1−O( δ 2
(TĜ2)2/3 ), and
T = O (δĜF 3/2 + Ĝ2 + δF ) .
Here, we define Ĝ2 := G2 + σ2 and the expectation in E‖∇f(xout)‖2 ≤ is taken both over the sampling of the clients during the running of the algorithm, the sampling of the mini-batches in local updates, and the choice of xout (which is chosen randomly from the client iterates yi).
Remarkably, the rates of our methods are independent of L and only depend on δ. Thus, when δ ≤ L and δ ≤ L/S for MimeMVR and MimeLiteMVR, the rates beat the server only lower bound of Ω( LG√
S 3/2 ). In fact, if the Hessian variance is small and δ ≈ 0, our methods only needO(1/ ) rounds to communicate. Intuitively, our results show that local steps are very useful when heterogeneity (represented by δ) is smaller than optimization difficulty (captured by smoothness constant L).
MimeMVR uses a momentum parameter β of the order of (1 − O(TG2)−2/3) i.e. as T increases, β asymptotically approaches 1. In contrast, previous analyses of distributed momentum (e.g. [73]) prove rates of the form G 2
S(1−β) 2 , which are worse than that of standard SGD by a factor of 1 1−β .
Thus, ours is also the first result which theoretically showcases the usefulness of using large momentum in distributed and federated learning. While we only prove the utility of local steps for MimeMVR, we believe our theory can be extended to other local update methods as well.
Our analysis is highly non-trivial and involves two crucial ingredients: i) computing the momentum at the server level to ensure that it remains unbiased and then applying it locally during every client update to reduce variance, and ii) carefully keeping track of the bias introduced via additional local steps. Our experiments (Sec. 5) verify our theoretical insights are indeed applicable in deep learning settings as well. See App. B for a proof sketch and App. G–H detailed proofs.

heading: 5 Experimental analysis on real world datasets
text: We run experiments on natively federated datasets to confirm our theory and accurately measure real world performance. Our main findings are i) MIME and MIMELITE consistently outperform FEDAVG, and ii) momentum and adaptivity significantly improves performance.

heading: 5.1 Setup
text: Algorithms. We consider three (meta) algorithms: FEDAVG, MIME, and MIMELITE. Each of these adapt four base optimizers: SGD, momentum, Adam, and Adagrad. FEDAVG follows [49] who run multiple epochs of SGD on each client sampled, and then aggregate the net client updates. This aggregated update is used as a pseudo-gradient in the base optimizer (called server optimizer). The learning rate for the server optimizer is fixed to 1 as in [67]. This is done to ensure all algorithms have the same number of hyper-parameters. MIME and MIMELITE follow Algorithm 4 and also run a fixed number of epochs on the client. However, note that this requires communicating both the full local-batch gradient as well as the parameter updates doubling the communication required to be sent by the client. For a fairer comparison, we split the sampled clients in MIME and MIMELITE into two groups–the first communicates only full local-batch gradient and the latter communicates only parameter updates. Thus, all methods have equal client communication to the server. This variant retains the convergence guarantees up to constants (details in the Appendix). We also run Loc-MIME where instead of keeping the global optimizer state fixed, we update it locally within the client. The optimizer state is reset after the round finishes. In all methods, aggregation is weighted by the number of samples on the clients.
Datasets and models. We run five simulations on three real-world federated datasets: EMNIST62 with i) a linear classifier, ii) an MLP, and iii) a CNN, iv) a charRNN on Shakespeare, and v) an LSTM for next word prediction on StackOverflow, all accessed through Tensorflow Federated [60]. The learning rates were individually tuned and other optimizer hyper-parameters such as β for momentum, β1, β2, ε0 for Adam and AdaGrad were left to their default values, unless explicitly stated otherwise. We refer to Appendix C for additional setup details and discussion.

heading: 5.2 Ablation and comparative study
text: In order to study the different algorithms, we train a 2 hidden layer (300µ-100) MLP on EMNIST62 with 10 local epochs for 1k rounds and use SGD+momentum (with tuned β) as the base optimizer.
Mime ≈ MimeLite > FedAvg > SCAFFOLD > FedProx. Fig. 1 (left) shows MIME and MIMELITE have nearly identical performance, and are about 7× faster than FedAvg. This implies
our strategy of applying momentum to client updates is faster than simply using server momentum. FedProx [40] uses an additional regularizer µ tuned over [0.1, 0.5, 1] (µ = 0 is the same as FedAvg). Regularization does not seem to reduce client drift but still slows down convergence [66]. SCAFFOLD [32] is also slower than Mime and FedAvg in this setup. This is because in cross-device setting with a large number of clients (N = 3.4k) means that each client is visited less than 6 times during the entire training (20 clients per round for 1k rounds). This means that the correction term utilized by SCAFFOLD uses control-variates which are quite stale (computed about 200 rounds ago) which slows down the convergence. In contrast, the SVRG correction term in Mime is computed using clients sampled in the current or previous rounds, and so is much more accurate. With momentum > without momentum. Fig. 1 (center) 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 > locally updated optimizer state. Finally, we check how the performance of Mime changes if instead of keeping the momentum fixed throughout a round, we let it change. The latter is a way to combine global and local momentum. The momentum is reset at the end of the round ignoring the changes the clients make to it. Fig. 1 (right) shows that this worsens the performance, confirming that it is better to keep the global optimizer state fixed as predicted by our theory.
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.

heading: 5.3 Large scale comparison with equal server and client communication
text: We perform a larger scale study closely matching the setup of [49]. For both MIME and MIMELITE, only half the clients compute and transmit the updated parameters, and other half transmit the full local-batch gradients. Hence, client to server communication cost is the same for all methods for all clients. However, MIME and MIMELITE require sending additional optimization state to the clients. Hence, we also reduce the number of clients sampled in each round to ensure sum total of communication at each round is 40× model size for EMNIST and Shakespeare experiments, and 100× model size for the StackOverflow next word prediction experiment. Since we only perform 1 local epoch, the hyper-parameters (e.g. epsilon for adaptive methods) are more carefully chosen following [49], and MIME and MIMELITE use significantly fewer clients per round, the difference between FEDAVG and MIME is smaller here. Table 2 summarizes the results. For the image classification tasks of EMNIST62 logistic and EMNIST62 CNN, Mime and MimeLite with Adam achieve the best performance. Using momentum (both with SGD and in Adam) significantly improves their performance. In contrast, FedAvgAdam is more unstable with worse performance. This is because FedAvg is excessively sensitive to hyperparameters (cf. App. E).
We next consider the character prediction task on Shakespeare dataset, and next word prediction on StackOverflow. Here, the momentum based methods (SGD+momentum and Adam) are slower than their non-momentum counterparts (vanilla SGD and AdaGrad). This is because the mini-batch gradients in these tasks are sparse, with the gradients corresponding to tokens not in the mini-batch being zero. This sparsity structure is however destroyed when using momentum or Adam. For the same reason, Mime which uses an SVRG correction also significantly increases the gradient density.
Discussion. For traditional tasks such as image classification, we observe that Mime (especially with Adam) usually outperforms MimeLite which in turn outperforms FedAvg. These methods are able to successfully leverage momentum and adaptivity to improve performance. For tasks where the client gradients are sparse, the SVRG correction used by Mime hinders performance. Adapting our techniques to work with sparse gradients (à la Yogi [75]) could lead to further improvements. Also, note that we reduce communication by naı̈vely reducing the number of participating clients per round. More sophisticated approaches to save on client communication including quantization or sparsification [58, 3], or even novel algorithmic innovations [1] could be explored. Further, server communication could be reduced using memory efficient optimizers e.g. AdaFactor [55] or SM3 [4].

heading: 6 Conclusion
text: Our work initiated a formal study of the cross-device federated learning problem and provided theoretically justified algorithms. We introduced a new framework MIME which overcomes the natural client-heterogeneity in such a setting, and can adapt arbitrary centralized algorithms such as Adam without additional hyper-parameters. We demonstrated the superiority of MIME via strong convergence guarantees and empirical evaluations. Further, we proved that a particular instance of our method, MimeMVR, beat centralized lower-bounds, demonstrating that additional local steps can yield asymptotic improvements for the first time. We believe our analysis will be of independent interest beyond the federated setting for understanding the sample complexity of non-convex optimization, and for yielding improved analysis of decentralized optimization algorithms.

Review summaries:
The paper proposes MIME, a general algorithmic framework for federated learning (FL) that can adapt to (and match the performance of) arbitrary centralized algorithms like Adam, Momentum SGD, and SGD. The authors show that via the MIME framework the convergence of any centralized algorithm can be translated into the convergence of the algorithm in the FL setting. The authors also propose a momentum-based variance reduction (MVR) algorithm that the authors claim to be faster than any centralized algorithm. Finally, the authors conduct extensive numerical experiments to show that the algorithms under the MIME framework perform well on real datasets.

This paper studies cross-device federated learning setting where the number of clients is very large and most of the clients participate in only one round of communication. Because of this low client participation per round, it is hard to have any client state/memory. State is usually helpful in traditional federated learning setting (eg: in SCAFFOLD) in estimating client heterogeneity, which can then be used to correct for the bias of the client. To mitigate this, the authors propose two meta-algorithms: MIME and MIMELITE, to adaptive a general class of centralized algorithms to this cross-device setting. MIME uses variance reduction at the client (like SCAFFOLD) and MIMELITE does not.  The paper analyzes the effect of these meta-algorithms for three standard centralized (server-only) deep learning optimization methods: SGD, Adam, MomentumVarianceReduction (MVR). However, reviewer could not find the pseduo-code for the MVR in the paper. For SGD, and Adam authors show that MIME modification performs at least as good as their centralized base algorithms for nonconvex problems. For MVR, authors prove that the server-level momentum of MIME helps beat the performance of centralized baselines under low-heterogeneity.

The paper introduces a new generalized algorithmic optimization framework named MIME for adapting and analyzing the convergence of any optimization algorithm applied in centralized settings into cross-device federated learning environments. The algorithmic reduction is accomplished by decoupling the centralized optimization algorithm into optimizer (server-side) and parameter state (client-side) updates in federated settings. The convergence of the proposed adaptation is shown both theoretically and empirically over a range of challenging real-world datasets for different optimization algorithms. Interestingly, when applying the proposed federated adaptation for a specific momentum-based optimization algorithm, the adaptation is proven to be asymptotically faster than its centralized counterpart.

The paper proposes a framework MIME to convert centralized optimization algorithms to the federated learning setting. The key components are some control variates to reduce the effect of data distribution heterogeneity.  Contribution: 1. A framework to convert centralized algorithms to the federated learning setting. 2. Theoretical analysis to characterize the convergence of converted algorithms 3. Experiments show MIME framework can have better performance than FedAvg.