M IME : M IMICKING C ENTRALIZED S TOCHASTIC A LGORITHMS IN F EDERATED L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
Federated learning (FL) is a challenging setting for optimization due to the heterogeneity of the data across different clients.
This heterogeneity has been shown to cause a client drift , which can signiﬁcantly degrade the performance of algorithms designed for the FL setting.
In contrast, centralized learning with centrally collected data is not affected by such a drift and has seen great empirical and theoretical progress with innovations such as momentum and adaptivity.
In this work, we propose a general algorithmic framework, M IME , which mitigates client drift and adapts arbitrary centralized optimization algorithms such as SGD and Adam to the federated learning setting.
M IME 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 of the centralized method run on iid data.
Our thorough theoretical and empirical analyses strongly establish M IME ’s superiority over other baselines.
I NTRODUCTION
Federated learning has become an important paradigm in large-scale machine learning where the training data remains distributed over a large number of clients, which may be mobile phones or network sensors (Koneˇcn`y et al., 2016b;a; McMahan et al., 2017; Mohri et al., 2019; Kairouz et al., 2019).
A centralized model, here referred to as a server model, is then trained, without ever transmitting 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 number of reliable clients, typically organizations, such as medical or ﬁnancial 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, we 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 a highly unreliable network.
Together, the essential features of this setting give rise to unique challenges not present in the cross-silo setting.
Here, we are 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 F ED A VG (McMahan et 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-ﬁtting’ to its atypical local data, a 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 et al., 2019), and clipping (You et 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 F ED A VG .
Perhaps due to such deﬁciencies, there exists a large gap in performance between the centralized setting, where data is centrally collected on the server, and the federated setting (Zhao et al., 2018; Hsieh et al., 2019; Hsu et al., 2019; Karimireddy et al., 2020).
To overcome such deﬁciencies, we propose a new framework, M IME , 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, M IME uses global statistics, e.g. momentum, and an
SVRG-style correction to mimic the updates of the centralized algorithm run on i.i.d. data.
These global statistics are computed only at the server level and kept ﬁxed throughout the local steps, thereby avoiding a 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 M IME that can adapt arbitrary centralized algorithms to this setting.
• We prove that incorporating server momentum 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 ﬁrst analysis showing improved rates by taking additional multiple steps for general smooth functions.
• Finally, we also propose a simpler variant, M IME L ITE , with an empirical performance similar to M IME .
We report the results of thorough experimental analysis demonstrating that both M IME and M IME L ITE are faster than F ED A VG .
Related work.
Analysis of FedAvg: Much of the recent work in federated learning has focused on analyzing F ED A VG .
For identical clients, F ED A VG coincides with parallel SGD, for which
Zinkevich et al.
(2010) derived an analysis with asymptotic convergence.
Sharper and more reﬁned 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 heterogeneous 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 ˇ cn ` y (2020) 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 Karimireddy et al.
(2020) and
Woodworth et al.
(2020a) for general functions, proving that F ED A VG 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,
Karimireddy et al.
(2020) proposed the S CAFFOLD 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.
In contrast, we focus on the cross-device setting where clients may be visited only once during training and where they are stateless .
This is akin to the difference between the ﬁnite-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 signiﬁcantly improves over vanilla F ED A VG .
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.
M IME , on the other hand, ensures that every local client update resembles the optimizer e.g. M IME would apply momentum in every client update and not just at 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.
5, Karimireddy et al., 2020; Wang et al., 2020b).
Other orthogonal directions which can be combined with M IME include tackling computation heterogeneity, where some clients perform many more updates than others (Wang et 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 et 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 et al.
for additional discussion.
P ROBLEM SETUP
This section formalizes the problem of cross-device federated learning.
We ﬁrst examine some key challenges of this setting (cf.
Kairouz et al., 2019) to ensure our formalism captures the difﬁculty:
Communication cost between the server and the clients is a 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.
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.
There may be a wide heterogeneity or non-i.i.d.-ness due to the difference of data distributions for the clients.
Thus, our objective will be to minimize the following quantity within the fewest number of clientserver communication rounds:
[Equation]
Here, f i denotes the loss function of client i and { ζ i, 1 , . . .
, ζ i,n i } 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 as a ﬁnite sum.
In each round, the algorithm samples a 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: if there exists a client with arbitrarily different data whom we may never encounter during training, then there is no hope to actually minimize f .
Thus for (1) to be tractable, it is necessary to assume bounded dissimilarity between different f i .
(A1) G 2 -BGD or bounded gradient dissimilarity: there exists G ≥[-
-]
such that
[Equation]
Next, we also characterize the variance in the Hessians.
Note that if f i ( · ; ζ ) is L -smooth, (A2) is always satisﬁed with δ ≤ 2 L and hence is more of a deﬁnition rather than an assumption.
Note that however, in realistic examples we expect the clients to be similar and hence that δ (cid:28) L .
(A2) δ -BHD or bounded Hessian dissimilarity: Almost surely, f is δ -weakly convex
i.e.
∇ 2 f i ( x ) (cid:23) − δI and the loss function of any client i satisﬁes
[Equation]
In addition, we assume that f ( x ) is bounded from below by f (cid:63) and is L -smooth, as is standard.
U SING 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 a single client is 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 a client i ∼ D and run SGDm with momentum parameter β and step size η :
[Equation]
Here, the gradient ∇ f i ( x t ) is unbiased i.e. E [ ∇ f i ( x t )]
= ∇ f ( x t ) and hence we are guaranteed convergence.
However, this strategy can be communication-intensive 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 in each round of communication.
Starting from y
= x t − 1 , F ED A VG (McMahan et al., 2017) runs multiple SGD steps on the sampled client i ∼ D y
[Equation]
) for k ∈ [ K ] , (3) and then a pseudo-gradient ˜ g t = − ( y K − x t ) replaces ∇ f i ( x t − 1 ) in the SGDm algorithm (2).
This is referred to as server-momentum since it 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 is because by using multiple local updates, (3) starts overﬁtting to the local client data, optimizing f i ( x ) instead of the actual global objective f ( x ) .
The net
[Figure]
effect is that F ED A VG moves towards an incorrect point (see Fig 1, left).
If K is sufﬁciently large, approximately
[Equation] [Equation]
x
[Equation]
arg min x
[Equation]
Further, the server momentum is based on ˜ g t and hence is also biased.
Thus, it cannot correct for the client drift.
We next see how a different way of using momentum can mitigate client drift.
Mime approach.
F ED A VG experiences client drift because both the momentum and the client updates are biased.
To ﬁx the former, we compute momentum using only global statistics as in (2):
[Equation]
To reduce the bias in the local updates, we will apply this unbiased momentum every step:
[Equation]
[Equation]
[Equation]
Note that the momentum term is kept ﬁxed during the local updates i.e. there is no local momentum used, only global momentum is applied locally.
Since m t − 1 is a moving average of unbiased gradients computed over multiple clients, it intuitively is a good approximation of the general direction of the updates.
By taking a convex combination of the local gradient with m t − 1 , the update (5) is potentially also less biased.
In this way M IME combines the communication beneﬁts of taking multiple local steps and prevents client-drift (see Fig 1, right).
Sec.
C makes this intuition precise.
M IME 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 1 describes two variants M IME and M IME L ITE , which consists of three components i) a 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 U which updates the parameters x , and a statistics step V ( · ) which keeps track of global statistics s .
Each step of the base algorithm B =
( U , V ) uses a gradient g to
[Equation]
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, etc.
input: initial x and s , learning rate η and base algorithm 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| (cid:80) j ∈S ∇ f j ( x ) (only for Mime) on client i ∈ S in parallel do initialize local model y
i ← x for k = 1 , · · · , K do sample mini-batch ζ from local data y
i ← y
i − η U ( ∇ f i ( y i ; ζ )
− ∇ f i ( x ; ζ ) + c , s ) ( Mime ) y
i ← y
i − η U ( ∇ f i ( y i ; ζ ) , s ) ( MimeLite ) end for compute full local-batch gradient ∇ f i ( x ) communicate ( y i , ∇ f i ( x )) end on client s ← V (cid:16) 1 |S| (cid:80)
i ∈S ∇ f i ( x ) , s (cid:17) (update optimization statistics)
x ← 1 |S| (cid:80) i ∈S y i (update server parameters) 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 ﬁxed 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 s ← V (cid:16) 1 |S| (cid:80)
i ∈S ∇ f i ( x ) , s (cid:17) , where ∇ f i ( x ) = 1 n i (cid:80) n i ν =1 ∇ f i ( x ; ζ i,ν ) .
(S TATS )
Note that we use full-batch gradients computed at the server parameters x , not client parameters
y i .
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 M IME and M IME L ITE differentiated using colored boxes.
Starting from y i ← x , repeat the following K times y i ←
y i − η U ( g , s ) , where
[Equation]
M IME L ITE simply uses the local minibatch gradient whereas M IME uses an SVRG style correction (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 M IME L ITE closely matches the performance of M IME .
Finally, there are two modiﬁcations made in practical FL: we weigh all averages across the clients by the number of datapoints n i , and we perform K epochs instead of K steps (McMahan et al., 2017).
The former modiﬁes the objective (1) with f i being weighted by n i , and the latter has been empirically observed to perform better, but lacks strong justiﬁcation (Wang et al., 2020b).
T HEORETICAL ANALYSIS OF M IME
Table 1 summarizes the rates of server-only methods, F ED A VG , S CAFFOLD and M IME (new results highlighted in blue).
Our theory focuses on M IME with base algorithm of SGD with momentum based variance reduction 1 (MimeMVR) since it obtains optimal rates.
[Table] [Table]
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 ∞ , making the bounds vacuous.
Requires K ≥ L/δ number of local updates.
Theorem I.
For L -smooth f with G 2 gradient dissimilarity (A1) , δ Hessian dissimilarity (A2) and
F :=
( f ( x 0 ) − f (cid:63) ) , let us run MimeMVR for T rounds and generate an output x out . This output satisﬁes E (cid:107)∇ f ( x out ) (cid:107) 2 ≤ (cid:15) under the following conditions
• PL-Strongly convex without momentum: for η = ˜ O (cid:16) min (cid:0) 1 δK + µK + L , 1 µT (cid:1)(cid:17) , β = 0 , and
[Equation]
(cid:17)(cid:19)
[Equation]
• Non-convex without momentum: for η = O (cid:16) min (cid:0) 1 δK + L , ( SF G 2 TK 2 ) 1 / 2 (cid:1) (cid:17) , β = 0 , and T
[Equation]
{+O (cid:18) LG 2 F S(cid:15) 
 + ( L + δK ) F (cid:15)K (cid:19) ,+} • Non-convex with momentum: for η = O (cid:16) min (cid:0) 1 δK + L , ( SF
[Equation]
2 TK 3 ) 1 / 3 (cid:1)(cid:17) , β = 1 − O (
[Equation]
2 ( TG 2 ) 2 / 3 ) ,
[Equation]
O (cid:18)(cid:16) (1 + δ ) G 2 F S(cid:15)(cid:17) 3 / 4 + ( L +
[Equation]
) F (cid:15)K (cid:19) The expectation in E (cid:107)∇ f ( x out ) (cid:107) 2 ≤ (cid:15) 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 x out (which is chosen randomly from the client iterates y i as described in the Appendix).
Table 1 shows that the rate (ignoring constants) of F ED A VG on non-convex functions is (cid:0) G 2 S(cid:15) 2 +
G(cid:15) 3 / 2 + L(cid:15) (cid:1) .
This is slower than simply running SGD which obtains a rate of (cid:0) G 2 S(cid:15) 2 + L(cid:15) (cid:1) .
In contrast,
MimeSGD obtains a rate of (cid:0) G 2 S(cid:15) 2 + δ(cid:15) (cid:1) where δ (cid:28) L thus improving upon both SGD and FedAvg.
While asymptotically these three rates may seem equivalent, in machine learning we care about low accuracy settings where (cid:15) is not too small (Bottou, 2010) and so the lower order terms matter, as [Figure] also additionally evidenced by our experiments (Sec. 6).
We can also compare with SCAFFOLD (Karimireddy et al., 2020) which obtains a rate of (cid:0) NS (cid:1) 23 L(cid:15) where N is the total number of clients.
While asymptotically this is a faster rate, N in the cross-device setting is potentially inﬁnite or at least comparable to the total number of training rounds, making these bounds vacuous.
This too is reﬂected in our experiments (Fig. 5).
Finally, by incorporating momentum based variance reduction
[Figure]
MimeMVR matches the lower bound of Ω (cid:0) G/ √ S(cid:15) (cid:1) 32 by Arjevani et al. The momentum β used in this case is of the order of (1 − O ( TG 2 ) − 2 / 3 ) i.e. as T increases, our momentum parameter asymptotically approaches 1. In contrast, previous analyses of distributed momentum (e.g. Yu et al. (2019a))
prove rates of the form G 2 S (1 − β ) (cid:15) 2 , which are worse than that of standard SGD by a factor of 1 1 − β .
Thus, ours is the ﬁrst 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. as the clients become more heterogeneous, there is stronger client-drift and hence we need more momentum to compensate. Our analysis is is highly non-trivial and involves three 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, ii) carefully keeping track of the bias introduced via additional local steps, and iii) an SVRG correction to allow using mini-batches.
Our experiments (Sec. 6) verify that the ﬁrst two theoretical insights are indeed applicable in deep learning settings as well, whereas the latter seems to 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.
E XPERIMENTAL ANALYSIS
We run experiments on simulated and real datasets to conﬁrm our theory.
Our main ﬁndings are i)
M IME outperforms F ED A VG across all settings, ii) its SVRG correction is useful for convex problems, and iii) momentum signiﬁcantly improves performance for non-convex problems.
We consider four algorithms: S ERVER -O
NLY , F ED A VG , M IME , and M IME L ITE .
Each of these adapt base optimizers SGD, SGDm, and Adam.
The S ERVER -O
NLY method computes a full batch gradient on each of the sampled clients and uses their aggregate directly in the base optimizer (akin to (2)).
For F ED A VG , 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 a pseudogradient in the base optimizer (called server optimizer).
The learning rate for the server optimizer is ﬁxed to 1 as in (Wang et al., 2020c).
This is done to ensure all algorithms have the same number of hyper-parameters.
Finally, M IME and M IME L ITE follow Algorithm 1 and also run a ﬁxed number of epochs on the client.
Aggregation is weighted by the number of samples on the clients. [Figure]
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. S IMULATED CONVEX EXPERIMENTS
Our simulated experiments use two clients each with a simple scalar quadratic loss, as in (Karimireddy et al., 2020).
We use full-batch gradients with both clients participating every round.
The simulated data has Hessian dissimilarity δ = 1 (A2) and smoothness L = 2 .
We vary the gradient dissimilarity (A1) as G ∈ [1 , 10 , 100] .
All the algorithms use momentum with β = 0 .
5 and their learning rates were tuned up to a tolerance of 5E-3 to ensure lowest loss after 60 rounds.
The results are collected in Fig.
When G is small, we see that F ED A VG can outperform the S ERVER -O
NLY (SGDm) baseline, though its loss quickly plateaus.
On increasing G , F ED A VG becomes even slower.
M IME L ITE differs from F ED A VG only in how the momentum is used.
In all settings, it slightly outperforms F ED A VG though even it sees a substantial slow down as we increase G .
This reﬂects our theory which predicts that for convex cases, momentum does not give signiﬁcant gains.
M IME , 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. R EAL WORLD EXPERIMENTS
We run real world deep learning experiments on EMNIST62 with a 2 layer MLP model and on
Cifar100 with ResNet20, both accessed through Tensorﬂow Federated (TFF, 2020a).
All methods run 10 local epochs, batch size 20, and the learning rates for all methods were individually tuned.
We refer to Appendix A for additional details and results.
Fig 3 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 our 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 signiﬁcant impact on the performance of Mime.
This is also in line with our theory and conﬁrms 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 ﬁxed throughout a 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.1 shows that this consistently worsens the performance, conﬁrming that it is better to keep the statistics ﬁxed.
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.
C ONCLUSION
Our work initiated a formal study of the cross-device federated learning problem.
We argued that the natural heterogeneity among the clients gives rise to client drift and signiﬁcantly hampers the performance of approaches such as F ED A VG .
We then showed how momentum can be an excellent tool to overcome this client drift if used correctly.
Based on this observation, we introduced a new framework M IME 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 M IME via strong convergence guarantees and empirical evaluations.
Appendix A E XPERIMENTAL SETUP A. D ESCRIPTION OF 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 ResNet20 on CIFAR100.
In all cases we report the top-1 test accuracy in our experiments.
EMNIST uses the metadata indicating the original author of the characters to separate them into multiple clients yielding naturally partitioned 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 × number of rounds / number of clients) provides an indication of how much of the training data is seen with S HAKESPEARE being closest to the cross-silo setting and S TACK O VERFLOW representing the most cross-device in nature.
[Table]
We use Tensorﬂow federated datasets (TFF, 2020a) to generate the datasets.
Our federated learning simulation code is written in Jax (Frostig et 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 2 groups.
Black and white was reversed in EMNIST62 (i.e. subtracted from 1) to make them similar to MNIST.
CIFAR100 used the usual pre-processing (normalization and centering), and data augmentation (random crop and horizontal ﬂipping) following (kuangliu,(accessed June 4, 2020).
A. P RACTICALITY OF EXPERIMENTS
In the experiments we only cared about the number of communication rounds, ignoring that M IME actually needs twice the number of bits per round and that the S ERVER -O NLY methods have a much smaller computational requirement.
This is standard in the federated learning setting as introduced by McMahan et al.
and is justiﬁed 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 not true. Incorporating communication compression strategies (Suresh et al., 2017; Alistarh et al., 2017; Karimireddy et al., 2019) or clientmodel compression strategies (Caldas et al., 2018a; Frankle and Carbin, 2019; Hamer et al., 2020) into our M IME 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 in 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 as have very easy to set default hyper-parameters.
Thus in our experimental setup we keep 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 was made both keeping in mind the limited memory available to each client as well as to match prior work.
Finally, while we limit ourselves to sampling 20 workers per round due to computational constraints, in real world FL thousands of devices are often available for training simultaneously each round (Bonawitz et al., 2019).
They also note that the probability of each of these devices being available has clear patterns and is far from uniform sampling.
Conducting a large scale experimental study which mimics these alternate forms of heterogeneity is an important direction for future work.
A. H YPERPARAMETER SEARCH
We use the EMNIST62 with MLP model as a ‘test-bed‘ for exploring different algorithms given it being both a representative task of cross-device FL as well as being computationally efﬁcient.
All plots reported are for this setting.
A more ﬁne-grained search over hyperparameters to report the ﬁnal test accuracies is made over the rest of the tasks/datasets.
For all SGDm methods, we pick momentum β = 0 .
For Adam methods, we ﬁx β 1 = 0 .
9 , β 2 = 0 .
99 , and ε = 1 × 10 − 3 similar to (Reddi et al., 2020).
None of the algorithms use weight decay, clipping etc.
The learning rate is then tuned to obtain the best test accuracy.
For all experiments, unless explicitly mentioned otherwise, the learning rate is searched over a grid of
[Equation]
[Equation] [Equation]
A C OMPARISON 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 similar 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 (cid:15) parameter for Adam.
As we explain in Section A.2, we chose to keep these parameters to some default values to compare methods in the ‘low-tuning’ setting.
We also point that while FedAdam struggles to perform in this setup, MimeAdam and MimeLiteAdam are very stable and even often outperform their SGDm counterparts.
This is also the default/recommended behavior in TensorFlow
Federated (TFF, 2020b) and (Wang et al., 2020c).
B A DDITIONAL EXPERIMENTS B. E FFECT OF CHANGING STATISTICS
[Figure]
Clearly, C-Mime is always worse than Mime.
This shows that adapting statistics during local client updates makes them too biased, and it best to keep them ﬁxed during each round like Mime does.
[Figure]
[Table]
For F ED A VG we searched over both the client and server learning rates, whereas for M IME and
M IME L ITE , we search only over client (base) learning rate. This search is performed over a 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 β 1 in (0.9, 0.99), β 2 in (0.99, 0.999) and (cid:15) in (0.01, 0.001, 0.0001). In this setting as well M IME and M IME L ITE outperform F ED A VG but with a smaller margin.
[Table]
x − η ((1 − β ) g + β m ) m = (1 − β ) g + β m
[Table]
v x − η (cid:15) + √ v g v = (1 − β ) g 2 + β v
[Table]
m , v x − η (cid:15) + √ v ((1 − β 1 )
g + β 1 m ) m
{+= (1 − β 1 ) g + β 1 m v = (1 − β 2 ) g 
 + β 2 v+} C P ROOF OVERVIEW
In this section, we give proof sketches of the main components of Theorem I: i) how momentum reduces the effect of client drift, ii) how local steps can take advantage of Hessian similarity, and iii) why the SVRG correction improves constants.
Improving the statistical term via momentum.
Note that the statistical (ﬁrst) term in Theorem I without momentum ( β = 0) for the convex case is LG 2 µS(cid:15) .
This is (up to constants) optimal and cannot be improved.
For the non-convex case however using β = 0 gives the usual rate of LG 2 S(cid:15) 2 .
However, this can be improved to (cid:16) (1+ δ ) G 2 F S(cid:15)(cid:17) 3 / 4 using momentum.
This matches a 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 we sample a single client i t in round t and that we use full-batch gradients.
Also let the local client update at step k round t be of the form
[Equation]
The ideal choice of update is of course d (cid:63)k = ∇ f ( y )
but however this is unattainable.
Instead,
M IME with momentum β = 1 − a uses d SGDm k = ˜ m
k ← a ∇ f i ( y )
+ (1 − a ) m t − 1 where m t − 1 is the momentum computed at the server.
The variance of this update can then be bounded as
[Equation]
The last step follows by unrolling the recursion on the variance of m .
We also assumed that η is small enough that y ≈ x t − 2 .
This way, momentum can reduce the variance of the update from
G 2 to ( aG 2 ) by using past gradients computed on different clients.
To formalize the above sketch requires slightly modifying the momentum algorithm similar to (Cutkosky and Orabona, 2019), and is carried out in Appendix G.
Improving the optimization term via local steps.
The optimization (second) term in Theorem I for the convex case is δK + L µK and for the non-convex case (with or without momentum) is δK + L (cid:15)K .
In contrast, the optimization term of the server-only methods is L/µ and L/(cid:15) respectively.
Since in most cases δ (cid:28) L , the former can be signiﬁcantly smaller than the latter.
This rate also suggests that the best choice of number of local updates is L/δ i.e. 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 to a careful analysis of the bias in the gradients computed during the local update steps.
Note that for client parameters y k − 1 , the gradient E [ ∇ f i t ( y k − 1 )] (cid:54) =
E [ ∇ f ( y k − 1 )] since y k − 1 was also computed using the same loss function f i t .
In fact, only the ﬁrst gradient computed at x t − 1 is unbiased.
Dropping the subscripts k and t , we can bound this bias as:
[Equation]
since unbiased
[Equation]
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
[Equation]
Here, we used the smoothness of f i ( · ; ζ ) and assumed that y i
≈ x
since we don’t move too far within a single round.
Thus, the SVRG correction allows us to use minibatch gradients in the local updates while still ensuring that the variance is of the order G 2 /S .
D T ECHNICALITIES
We examine some additional deﬁnitions and introduce some technical lemmas.
D. A SSUMPTIONS AND DEFINITIONS
We make precise a few deﬁnitions and explain some of their implications.
We ﬁrst 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 2 (gradient dissimilarity) is large while δ (Hessian dissimilarity) is small, or vice-versa.
Our assumption about the bound on the G gradient dissimilarity can easily be extended to ( G, B ) gradient dissimilarity used by (Karimireddy et al., 2019):
[Equation]
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 sufﬁciently over-parameterized model, it is possible to make G = 0 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 L ).
The second crucial deﬁnition 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 as 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 x and y the following holds: [Equation]
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 deﬁne Ψ
i ( z ; ζ ) =
f i ( z ; ζ ) − f ( z ) .
Then, we can rewrite (A2) and (8) respectively as
[Equation]
and
[Equation]
Thus (8) and (A2) are both different notions of smoothness of Ψ i ( x ; ζ ) (formal deﬁnition of smoothness will follow soon).
The latter deﬁnition closely matches the notion of squared-smoothness used by Arjevani et al.
and is a promising relaxation of (A2).
However, we run into some technical issues since in our case the variable y can also be a random variable and depend on the choice of the client i .
Extending our results to this weaker notion of Hessian-similarity and proving tight non-convex lower bounds is an exciting theoretical challenge.
Finally note that if the functions f i ( x ; ζ ) are assumed to be smooth as in (Shamir et al., 2014;
Arjevani and Shamir, 2015; Karimireddy et al., 2020), then Ψ i (( x ; ζ ) is 2 L -smooth.
Thus, we always have that δ ≤ 2 L .
But, as Shamir et al.
(2014) show, it is possible to have δ (cid:28) L if the data distribution amongst the clients is similar.
Further, the lower bound from Arjevani and
Shamir (2015) proves that Hessian-similarity is the crucial quantity capturing the number of rounds of communication required for distributed/federated optimization.
We next deﬁne the terms smoothness and strong-convexity which we repeatedly use in the paper.
(A3) f is L-smooth and satisﬁes:
[Equation]
The assumption (A3) also implies the following quadratic upper bound on f
[Equation]
Further, if f is twice-differentiable, (A3) implies that (cid:107)∇ 2 f ( x ) (cid:107) ≤ β for any x .
(A4) f is µ -PL strongly convex (Karimi et al., 2016) for µ > 0 if it satisﬁes:
[Equation]
Note that PL-strong convexity is much weaker than the standard notion of strong-convexity (Karimi et al., 2016).
D. S OME TECHNICAL LEMMAS
Now we cover some technical lemmas which are useful for computations later on.
First, we state a relaxed triangle inequality true for the squared (cid:96) 2 norm.
Lemma 1 (relaxed triangle inequality) .
Let { v 1 , . .
. , v τ } be τ vectors in R d .
Then the following are true:
[Equation] any
[Equation]
Proof.
The proof of the ﬁrst statement for any c > 0 follows from the identity:
[Equation]
For the second inequality, we use the convexity of x → (cid:107) x (cid:107) 2 and Jensen’s inequality
[Equation]
Next we state an elementary lemma about expectations of norms of random vectors.
Lemma 2 (separating mean and variance) .
Let { Ξ 1 , . . .
, Ξ τ } be τ random variables in R d which are not necessarily independent.
First suppose that their mean is E [Ξ i ] = ξ i and variance is bounded as E [ (cid:107) Ξ i − ξ i (cid:107) 2 ] ≤ σ 2 .
Then, the following holds
[Equation]
Now instead suppose that their conditional mean is E [Ξ i | Ξ i − 1 , . . .
Ξ 1 ] = ξ i i.e. the variables { Ξ i − ξ i } form a martingale difference sequence, and the variance is bounded by E
[ (cid:107) Ξ
i − ξ i (cid:107) 2 ] ≤ σ 2 as before.
Then we can show the tighter bound
[Equation]
Proof.
For any random variable X , E [ X 2 ] =
( E [ X − E
[ X ]])
+ ( E [ X ]) 2 implying
[Equation]
Expanding the above expression using relaxed triangle inequality (Lemma 1) proves the ﬁrst claim:
[Equation]
For the second statement, ξ i is not deterministic and depends on Ξ i − 1 , . . .
Ξ 1 .
Hence we have to resort to the cruder relaxed triangle inequality to claim
[Equation]
and then use the tighter expansion of the second term:
[Equation]
The cross terms in the above expression have zero mean since { Ξ i − ξ i } form a martingale difference sequence.
E P ROPERTIES OF FUNCTIONS WITH δ BOUNDED H ESSIAN DISSIMILARITY
We now study two lemmas which hold for any functions which satisfy (A2).
The ﬁrst is closely related to 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 :
[Equation]
Proof.
Consider the function Ψ( z ) :
= f i ( z ; ζ ) − f ( z ) .
By the assumption (A2), we know that [Equation] ≤ δ for all z i.e. Ψ is δ -smooth.
By standard arguments based on taking limits (Nesterov, 2018), this implies that
[Equation]
Plugging back the deﬁnition of Ψ into the above inequality proves the lemma.
Next, we see how weakly-convex functions satisfy a weaker notion of “averaging does not hurt”.
Lemma 4 (averaging) .
Suppose f is δ -weakly convex.
Then, for any γ ≥ δ , and a sequence of parameters { y i } i ∈S and x :
[Equation]
, where ¯
[Equation]
i ∈S y
[Equation]
Proof.
Since f is δ -weakly convex, Φ( z ) :
= f ( z ) + γ 2 (cid:107)
z − x (cid:107) 2 is convex.
This proves the claim since 1 |S| (cid:80) i ∈S Φ( y i ) ≤ Φ(¯ y ) .
F A NALYSIS OF M IME SGD ( WITHOUT MOMENTUM )
Let us rewrite the MimeSGD update using notation convenient for analysis.
In each round t , we sample clients S t such that |S t | = S .
The server communicates the server parameters x t − 1 as well as the average gradient across the sampled clients c t − 1 deﬁned as
[Equation]
Note that computing c t − 1 itself requires two rounds of communication.
But from a theoretical viewpoint, this only changes the communication rounds required by a constant factor and hence we ignore this issue.
Practically, we recommend using MimeLite if this additional rounds of communication are an issue.
Then each client i ∈ S t makes a copy y ti,
= x t − 1 and perform K local client updates.
In each local client update k ∈ [ K ] , the client samples a dataset ζ ti,k and
[Equation]
After K such local updates, the server then aggregates the new client parameters as
[Equation]
Variance of update.
Consider the local update at step k on client i , dropping superscript t y i,k
= y
i,k − 1 − η d i,k , where d i,k := ∇ f i ( y i,k − 1 ; ζ i,k ) − ∇ f i ( x ; ζ i,k )
+ c .
Lemma 5.
Given that assumptions (A1) and (A2) are satisﬁed, each client update satisﬁes
[Equation]
Proof.
Starting from the deﬁnition of d i,k and the relaxed triangle inequality,
[Equation]
We used Lemma 3 to bound the ﬁrst term.
Taking expectations on both sides to bound the second term via (A1) yields the lemma.
Distance moved in each round.
We show that the distance moved by a 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 y i,k − 1 simply as y and y i,k as y + .
Lemma 6.
For update following (12) for η ≤ 1 4 Kδ satisfying (A1) and (A2) , we have at any step k ,
[Equation]
Proof.
Starting from the update (12) and the relaxed triangle inequality Lemma 1 with c = K ≥ 1 ,
[Equation]
The second to last step used the variance bound in Lemma 5.
The proof now follows from the restriction on step-size since 16 K 2 η 2 δ 2 ≤ 1 .
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 µ ≥
and each step of MimeSGD with step size η ≤ min (cid:16) 118 L ,
[Equation]
, 1 42 µK (cid:17) , and given that (A1) – (A3) hold, we have
[Equation]
where we deﬁne
[Equation]
Proof.
The assumption that f is L -smooth implies a quadratic upper bound (10).
Using this in our case, we have
[Equation]
Let us examine the terms T 1 and T 2 separately.
By our variance bound Lemma 5, we have that
[Equation]
To simplify T 1 , the biggest obstacle is that E [ ∇ f i ( y ; ζ )]
(cid:54) =
∇ f ( y ) since y itself depends on the sampling of the client i .
Only the server gradient is unbiased and E [ c ] = ∇ f ( x ) .
Instead we will use the similarity of the functions as in Lemma 3:
[Equation]
The ﬁrst inequality above used that for any a, b , the following holds − 2 ab
= ( a − b ) 2 − a 2 − b 2 ≤ ( a − b ) 2 − a 2 .
The second used the similarity Lemma 3.
Combining the terms T 1 and T 2 together, we have
[Equation]
To bound the distance between y and x , we use Lemma 6 multiplied on both sides by δ (cid:0)
+ 3 K (cid:1)
K − k .
Note that δ ≤ δ (cid:0) 1 + 3 K (cid:1) K − k ≤ 21 δ .
This gives us for any constant µ ≥δ (cid:16) 1 + 3 K (cid:17) K − k E (cid:107) y
+ − x (cid:107) 2 ≤ δ (cid:16) 1 + 3 K (cid:17) K − k (cid:16)
+ 2 K (cid:17) (cid:107)
y − x (cid:107) 2 + 6 Kδ (cid:16) 1 + 3 K (cid:17) K − k η 2 G 2 S
+ 6 Kδ (cid:16) 1 + 3 K (cid:17) K − k η 2 (cid:107)∇ f ( y ) (cid:107)≤ δ (1 − µη ) (cid:16) 1 + 3 K (cid:17) K − ( k − 1) (cid:107)
y − x (cid:107) 2 + 6 Kδ (cid:16) 1 + 3 K (cid:17) K − k η 2 G 2 S
+ 6 Kδ (cid:16) 1 + 3 K (cid:17) K − k η 2 (cid:107)∇ f ( y ) (cid:107)
+ (cid:16) 1 + 3 K (cid:17) K − k
( µηδ − δ K ) (cid:107)
y − x (cid:107)≤ δ (cid:0) 1 + 3 K (cid:1) K − ( k − 1) (cid:107)
y − x (cid:107)
+ 126 Kδη 2 G 2 S + 126 Kδη 2 (cid:107)∇ f ( y ) (cid:107)+ (21 µηδ − δ K ) (cid:107)
y − x (cid:107) 2 .
The second inequality from the last used that 1+2 /K <
(1+3 /K )(1 − µη )
= (1+2 /K )
+(1 /K − (1 + 3 /K ) µη ) .
This is true by our restriction that η < 1 42 µK , which implies (1 + 3 /K ) µη < 4 µη < 1 / (10 K ) and so that (1 /K − (1 + 3 /K ) µη ) > 0 .
Adding the two bounds, we get the following recursion
[Equation]
Now, note that our constraint on the step-size η ≤ min( 118 L , 1 756 δK ) implies that 252 Kδη 2 +3 Lη 2 ≤ η 2 and K ( ηδ 2 +3 Lη 2 δ 2 +42 µηδ ) ≤ 2 δ .
Plugging this into the above bound and recalling that δ ≤ L ﬁnishes the proof.
Convergence for PL strongly-convex functions.
We will unroll the one step progress Lemmato compute a linear rate.
Theorem II.
Suppose that (A1) – (A4) are satisﬁed for µ > 0 .
Then the updates of MimeSGD with step-
[Equation]
where we deﬁne F := f ( x 0 ) − f (cid:63) , ¯ y tk is chosen to be y ti,k for i ∈ S t uniformly at random, and the output x out to be ¯ y tk with probability proportional to (1 − ηµ 4 ) KT − kt .
Proof.
Note that by PL strong convexity (A4), we have
[Equation]
Using this, we can tighten the one step progress Lemma 7 as η
[Equation]
Now take a weighted sum over the steps k using weights (1 − ηµ 4 ) K − k η
[Equation]
k ∈
[Equation]
{+By the initialization y ti, 
+}
= x t − 1 and hence Φ ti, 0 = E f ( x t − 1 ) − f (cid:63) and further by the averaging
Lemma 4, we have
[Equation]
Hence, on averaging over the clients we get the one round progress lemma
[Equation]
Now further taking a weighted average over the rounds t ∈
[ T ] with weights proportional to (1 − ηµ 4 ) tK gives
[Equation]
Finally [Equation] choosing the right step size, similar to Lemma 23 of (Karimireddy et al., 2020) yields the desired rate.
Convergence for general functions.
We will unroll the one step progress Lemma 7 to compute a sublinear rate.
Theorem III.
Suppose that (A1) – (A3) are satisﬁed.
Then the updates of MimeSGD with step-size η = min (cid:16) η max ,
[Equation]
(cid:17) for η max = min (cid:0) 118 L , [Equation] (cid:1) satisfy
[Equation]
where we deﬁne F := f ( x 0 ) − f (cid:63) .
Proof.
By summing over the equations from Lemma 7 for all local steps in one round we obtain
[Equation]
By the initialization y ti,
= x t − 1 , hence A ti,
=
A tj, 0 = E
[ f ( x t − 1 )] for all i, j ∈ S t .
Furthermore, by Lemma
[Equation]
This means that we can keep unrolling over all rounds, obtaining
[Equation]
By noting A 1 i, 0 − A Ti,K = ( f ( x 0 ) − f (cid:63) ) − ( E [ f ( x T ) − f (cid:63) ) ≤ F and the choice of the stepsize the theorem follows.
G A NALYSIS OF M IME MVR ( 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 ﬁt the M IME framework (B ASE A LG ) since it requires computing gradients at two points on the same batch.
However, it is straightforward to extend the idea of M IME to MVR as we will now do.
We use MVR as a theoretical justiﬁcation 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 ﬁt the framework of M IME .
MimeMVR algorithm.
Now, we formally describe the MimeMVR algorithm.
In each round t , we sample clients S t such that |S t | =
S .
The server communicates the server parameters x t − 1 , the momentum m t − 1 and the average gradient across the sampled clients c t − 1 deﬁned as
[Equation]
Note that both c t − 1 and m t − 1 use gradients and parameters from previous rounds (different from the previous section).
Then each client i ∈ S t makes a copy y ti,
= x t − 1 and perform K local client updates.
In each local client update k ∈ [ K ] , the client samples a dataset ζ ti,k and y ti,k
= y ti,k − 1 − η d ti,k , where
[Equation] [Equation]
After K such local updates, the server then aggregates the new client parameters as
[Equation]
The momentum term is updated at the end of the round for a ≥ 0 as
[Equation]
As we can see, the momentum update of MVR can be broken down into the usual SGDm update, and a correction.
Intuitively, this correction term is very small since f i is smooth and x t −
≈ x t − 2 .
Another way of looking at the update (17) is to note that if all functions are identical i.e. f j = f k 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.
Momentum variance bound.
We compute the variance of the server momentum m t − 1 .
Deﬁne the variance term V t
= m t − ∇ f ( x t − 1 ) .
Then its expected norm can be bounded as follows.
Lemma 8.
For the momentum update (17) , given (A1) and (A2) , the following holds for any a ∈ [0 , 1] and V t :
= m t − ∇ f ( x t − 1 )
[Equation]
Proof.
Starting from the momentum update (17),
[Equation]
Now
[Equation]
the term V t − 1 does not have any information from round t and hence is statistically independent of the rest of the terms.
Further
[Equation]
the rest of the terms have mean 0.
Hence
[Equation]
we can separate out the zero mean noise terms from the V t − 1 following Lemma 2 and then the relaxed triangle inequality Lemma 1 to claim
[Equation]
The inequality used the Hessian similarity Lemma 3 to bound the second term and the heterogeneity bound (A1) to bound the last term.
Finally, note that (1 − a ) 2 ≤ (1 − a ) ≤ 1 for a ∈ [0 , 1] .
Update variance bound.
Now we examine the variance of our update in each local step d ti,k .
Lemma 9.
For the client update (15) , given (A1) and (A2) , the following holds for any a ∈ [0 , 1]
[Equation]
Proof.
Starting from the client update (15), we can rewrite it as
[Equation]
We can use the relaxed triangle inequality Lemma 1 to claim
E
[Equation]
d ti,
[Equation]
The last inequality used the Hessian similarity Lemma 3 to bound the second term and the heterogeneity bound (A1) to bound the last term.
Also, (1 − a ) 2 ≤ 1 since a ∈ [0 , 1] .
Distance moved in each step.
We show that the distance moved by a client in each step during the client update can be controlled.
Lemma 10.
For MimeMVR updates (15) with η ≤ 1 6 Kδ and given (A1) and (A2) , the following holds
[Equation]
where we deﬁne
[Equation]
Proof.
Starting from the MimeMVR update (15) and the relaxed triangle inequality with c = 2 K
[Equation]
The last inequality used the update variance bound Lemma 9.
We can simplify the expression further since η ≤ 1 6 Kδ implies 18 Kη 2 δ 2 ≤ 12 K .
Similar computations for E (cid:107) y ti,k − x t − 1 (cid:107) 2 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 MimeMVR with step size η ≤ min (cid:0) 1 L ,
[Equation]
and momentum parameter a = 1536 η 2 δ 2 K 2 .
Then, given that (A1) – (A3) hold, we have
[Equation]
Proof.
The assumption that f is L -smooth implies a quadratic upper bound (10).
[Equation]
The second equality used the fact that for any a, b , − 2 ab = ( a − b ) 2 − a 2 − b 2 .
The second term can be removed since η ≤ 1 L .
Taking expectation on both sides and using the update variance bound
Lemma 9,
[Equation]
Multiplying the momentum variance bound Lemma 8 by 3 ηa , we have
[Equation]
We will also multiply the distance bound Lemma 10 by 8 ηδ 2 K a (cid:0)
+ 2 K (cid:1)
K − k .
Note that for any
K ≥ 1 and k ∈ [ K ] , we have 1 ≤ (cid:0)
+ 2 K (cid:1)
K − k ≤ 8 .
Then we get
[Equation]
where recall that we deﬁned ∆ ti,k :=
max
[Equation]
Combining the three inequalities together, we get
[Equation]
Note that 1152 η 3 δ 2 K 2 a = 3 η 4 since we deﬁned a = 1536 η 2 δ 2 K 2 .
Further, a ≤ 1 when deﬁned this way since we assumed η ≤ 1 40 δK .
Similarly, the deﬁnition of a implies that 384 η 3 δ 2 K 2 a = η 4 .
Thus, we can simplify the above expression as
[Equation]
This proves the lemma.
Progress in one round.
Let us sum over all the steps within a round to compute the progress made in a full round.
Lemma 12.
For any round of MimeMVR with step size η ≤ min (cid:0) 1 L ,
[Equation]
and momentum parameter a = 1536 η 2 δ 2 K 2 .
Then, given that (A1) – (A3) hold, we have
[Equation]
where we deﬁne the sequence
[Equation]
Proof.
We start by summing Lemma 11 over the client updates
[Equation]
Recall that we deﬁned ∆ ti,k := max (cid:16) E (cid:107) y ti,k − x t − 2 (cid:107) 2 , E (cid:107) y ti,k − x t − 1 (cid:107) 2 , E (cid:107) x
t − 1 − x t − 2 (cid:107) 2 (cid:17) .
Because y ti, 0 = x t − 1 , we can simplify
[Equation]
Then by the averaging Lemma 4, we have
[Equation]
So by averaging our recursion over the sampled clients, and diving our summation over the updates by K , we get
[Equation]
Theorem IV (non-convex convergence of MimeMVR)
[Equation]
Let us run MimeMVR with step size η ≤ min (cid:16) 1 15 K (cid:0) FSTδ 2 G 2 (cid:1) 1 / 3 , 1 L , 1 40 δK (cid:17) and momentum parameter a =
η 2 δ 2 K 2 .
Then, given that (A1) – (A3) hold
[Equation]
we have
[Equation]
where we deﬁne F := f ( x 0 ) − f (cid:63) .
Proof.
Unroll the one round progress Lemma 12 and average over T rounds to get
KST
[Equation]
t ∈
[Equation]
Our choice of step size now yields the desired rate
[Equation]