Why (and When) does Local SGD Generalize Better than SGD?
Authors: Xinran Gu
Kaifeng Lyu
Longbo Huang
Sanjeev Arora
Venue: ICLR 2023 poster (2023)
Keywords: local SGD
SDE
regularization
implicit bias
deep learning theory
optimization
distributed training

Abstract: Local SGD is a communication-efficient variant of SGD for large-scale training, where multiple GPUs perform SGD independently and average the model parameters periodically. It has been recently observed that Local SGD can not only achieve the design goal of reducing the communication overhead but also lead to higher test accuracy than the corresponding SGD baseline (Lin et al., 2020b), though the training regimes for this to happen are still in debate (Ortiz et al., 2021). This paper aims to understand why (and when) Local SGD generalizes better based on Stochastic Differential Equation (SDE) approximation. The main contributions of this paper include (i) the derivation of an SDE that captures the long-term behavior of Local SGD in the small learning rate regime, showing how noise drives the iterate to drift and diffuse after it has reached close to the manifold of local minima, (ii) a comparison between the SDEs of Local SGD and SGD, showing that Local SGD induces a stronger drift term that can result in a stronger effect of regularization, e.g., a faster reduction of sharpness, and (iii) empirical evidence validating that having a small learning rate and long enough training time enables the generalization improvement over SGD but removing either of the two conditions leads to no improvement.

Published as a conference paper at ICLR 2023
WHY (AND WHEn) DOES LOCAL SGD GENERALIZE BETTER THAN SGD?
Xinran Gu
Kaifeng Lyu"
Institute for Interdisciplinary Information Sciences Tsinghua University
Department of Computer Science Princeton University k klyuBcs, www.ppiiionnn.m.. edu
gxr21 @mai 1s. t.s inghua... edu. cn
Longbo Huang!
Sanjeev Arora
Institute for Interdisciplinary Information Sciences Tsinghua University longbohuang@tss Inghuaa edu.cn
Department of Computer Science Princeton University arora Bcs .cc.pp... pri nceton. edu
ABSTRACT
Local SGD is communication--ffiieenn variant of SGD for large-scale training, where multiple GPUs perform SGD independently and average the model param- eters periodically. It has been recently observed that Local SGD not only achieve the design goal of reducing the communication overhead but also lead higher test accuracy than the corresponding SGD baseline (Lin et al., 2020b), though the training regimes for this to happen are still in debate (Ortiz et al.) 2021). This paper aims to understand why (and when) Local SGD generalizes better based on Stochastic Differential Equation (SDE) approximation. The main contributions of this paper include (i) the derivation of an SDE that captures the long-term behavior of Local SGD in the small learning rate regime, showing how noise drives the iterate drift and diffuse after has reached close to the mani- fold of local minima, (ii) comparison between the SDEs of Local SGD and SGD, showing that Local SGD induces a stronger drift term that can result in a stronger effect of regularization, e.g., faster reduction of sharpness, and (iii) empirical ev- idence validating that having a small learning rate and long enough training time enables the generalization improvement over SGD but removing either of the two conditions leads to no improvement.
INTRODUCTION
As deep models have grown larger, training them with reasonable wall-clock times has led to new distributed environments and new variants of gradient- based training. Recall that Stochastic Gradi- Descent (SGD) tries to solve minger [e E)). where
Rd is the parameter vector of the model, f(0: <) is the loss function for data sample drawn from the training distribution D. e.g., the uniform distribution over the training set. SGD with learning rate v) and batch size does the following update at each step, using a batch of B independent & ... Er, D:
[<(@:
B Σ Ve.l..
1
0111 0: 79t, where gi
Parallel SGD tries improve wall-clock time when the batch size B is large enough. It distributes the gradient computation to K 2 2 workers, each whom focuses on a local batch of B := B K samples and computes the average gradient over the local batch. Finally, g: is obtained by averaging the local gradients over the workers.
However, large-batch training leads significant test accuracy drop compared small-batch
training baseline with the same number of training steps or epochs (Smith et al. 2020 Shallue et al.
Equaal contribution
Corresponding authors

---

Published as a conference paper at ICLR 2023
Top-1 Test Accuracy
Test Accuracy
0.7
0.70
0.90
H
H
0.75
0.6-
0.85
0.74
0.5
0
100
200
300
400
500
100
150
Epochs
Epochs
(a) CIFAR-10, B = -4996 ResNet-56.
(b) ImageNet, B 8192, ResNet-50.
Figure/ 1: Post-Local SGD (H 1) generalizes better than SGD (H 1). We switch to Local SGD at the first learning rate decay (epoch #250) for CIFAR-IO and the second learning rate decay (epoch #100) for ImageNet. See Appendix M. for training details.
2019 Keskar et al.
2017 Jastrzebski et al. 2017). D. Reducing this generalization gap is the goal of much subsequent research. It was suggested that the generalization gap arises because larger batches lead to reduction in the level of noise in batch gradient (see Appendix A for more discussion). The Linear Scaling Rule (Krizhevsky 2014 Goyal et al.) 2017 Jastrzębski et al. 2017) tries to fix this by increasing the learning rate in proportion to batch size. This is found to reduce the generalization gap for (parallel) SGD, but does not entirely eliminate it.
To reduce the generalization gap further, Lin et al. (020b) discovered that a variant of SGD, called Local SGD YYu et al.] 2019; Wang & Joshi 2019; Zhou & Cong 2018). can be used as a strong component. Perhaps surprisingly, Local SGD itself not designed for improving generalization, but for reducing the high communication cost for synchronization among the workers, which is another important issue that often bottlenecks large-batch training (Seide et al. 2014 Strom 2015 Chen iet al. 2016 Recht et al. 2011). Instead of averaging the local gradients per step as in parallel SGD, Local SGD allows K workers to train their models locally and averages the local model parameters whenever they finish H local steps. Here every worker samples a new batch at each local step, and in this paper we focus on the case where all the workers draw samples with or without replacement from the same training set. See Appendix for the pseudocode.
More specifically, Lin et al. 2020b) proposed Post-local SGD, a hybrid method that starts with parallel SGD (equivalent to Local SGD with H = in math) and switches to Local SGD with H > after fixed number of steps to- They showed through extensive experiments that Post- local SGD significantly outperforms parallel SGD in test accuracy when to is carefully chosen. In we reproduce this phenomenon on both CIFAR-10 and ImageNet.
Figure I
As suggested by the success of Post-local SGD, Local SGD can improve the generalization of SGD by merely adding more local steps (while fixing the other hyperparameters), at least when the train- ing starts from model pre-trained by SGD. But the underlying mechanism is not very clear, and there is also controversy about when this phenomenon can happen (see Section 2. for a survey). The current paper tries to understand: Why does Local SGD generalize better? Under what general conditions does this generulization benefit arise?
Previous theoretical research Local SGD is mainly restricted the convergence rate for minimiz- ing a convex or non-convex objective (see [opendix[AAoor a survey). A related line of works (Stich) 2018 Yu al.] 2019 Khaled al. 2020) showed that Local SGD has slower convergence rate compared with parallel SGD after running the same number of steps/epochs. This convergence re- sult suggests that Local SGD may implicitly regularize the model through insufficient optimization, but this does not explain why parallel SGD with early stopping. which may incur an even higher training loss, still generalizes worse than Post-local SGD.
Our Contributions. In this paper, we provide the first theoretical understanding on why (and when) switching from parallel SGD to Local SGD improves generalization.
1. In Section 2.2, we conduct ablation studies on CIFAR-10 and ImageNet and identify a clean setting where adding local steps to SGD consistently improves generalization: if the learning rate small and the total number of steps is sufficient, Local SGD eventually generalizes better than the corresponding (parallel) SGD baseline.
2. In Section 3.2
we derive special SDE that characterizes the long-term behavior of Local SGD in the small learning rate regime, as inspired by a previous work (Li et al. 2021b) that proposed this type of SDE for modeling SGD. These SDEs can track the dynamics after the iterate has reached close to a manifold of minima. In this regime, the expected gradient is near zero, but the gradient noise can drive the iterate wander around. In contrast to the conventional SDE 3 for
2

---

Published as a conference paper at ICLR 2023
SGD, where the drift and diffusion terms are connected respectively to the expected gradient and gradient noise, the SDE we derived for Local SGD has drift and diffusion terms both connected to gradient noise.
3. Section 3.3 explains the generalization improvement of Local SGD over SGD by comparing the corresponding SDEs: increasing the number of local steps H strengthens the drift term of SDE while keeping the diffusion term untouched. We hypothesize that having a stronger drift term can benefit generalization.
4. As a by-product, we provide new proof technique that can give the first quantitative approxi- mation bound for how well Li al. (2021b)'s 's SDE approximates SGD.
Back to the discussion the generalization gap between small- and large-batch training, we remark that this gap can occur early in training when the learning rate is very large (Smith et al. 2020) and Local SGD cannot prevent this gap in this phase. Instead, our theory suggests that Local SGD can reduce the gap in late training phases after decaying the learning rate.
2 WHEN DOES LOCAL SGD GENERALIZE BETTER?
In our motivating example of Post-local SGD,, switching from SGD to Local SGD can outperform running SGD alone (i.e., no switching) in test accuracy, but this improvement does not always arise and can depend on the choice of the switching time point. Because of this, a necessary first step for developing a theoretical understanding of Local SGD is to identify under what general conditions Local SGD can improve the generalization of SGD by merely adding local steps.
2.1 THE DEBATE ON LOCAL SGD
We first summarize debate in the literature regarding when to switch from SGD to Local SGD in running Post-local SGD, which hints the conditions so that Local SGD can improve upon SGD.
Local SGD generalizes better than SGD CIFAR-10. Lin al. (020b) empirically observed that Post-local SGD exhibits a better generalization performance than SGD. Most of their exper- iments are conducted on CIFAR-IO and CIFAR- 100 with multiple learning rate decays, and the algorithm switches from (parallel) SGD to Local SGD right after the first learning rate decay. We refer to this partieular choice of the switching time point as the first-decay switching strategy for short. To justify this strategy, they empirically showed that the generalization improvement can be less significant if starting Local SGD from the beginning or right after the second learning rate de- cay. It has also been observed by Wang & Joshi (2021) that running Local SGD from the beginning improves generalization, but the test accuracy improvement may not be large enough. A subsequent work by Lin et al.] (2020a) showed that adding local steps to Extrap-SGD, variant of SGD pro- posed therein, after the first learning rate decay also improves generalization, suggesting that the first-decay switching strategy can also be applied to the post-local variant of other optimizers.
Does Local SGD exhibit the same generalization benefit on large-scale datasets? Going be- yond CIFAR-10. 10, Lin et al.) (2020b) conducted a few ImageNet experiments and showed that Post- local SGD with first-decay switching strategy still leads to better generalization than SGD. However, the improvement is sometimes marginal, e.g., 0.1% for batch size 8192. For the general case, they suggested that the time of switching should be tuned aiming at "capturing the time when trajectory starts to get into the influence basin of local minimum" in a footnote, but no further discussion or experiments are provided to justify this guideline. Ortiz et al.] (202T) conducted a more extensive evaluation on ImageNet (with a different set of hyperparameters) and concluded with the opposite: the first-decay switching strategy can hurt the validation accuracy. Instead, switching at a later time, such as the second learning rate decay, leads to a better validation accuracy than SGD To explain this phenomenon, they conjecture that switching to Local SGD has regularization effect that is beneficial only in the short-term, so always better to switch as late as possible. They further conjecture that this discrepancy between CIFAR-10 and ImageNet is mainly due to the task scale. On TinyImageNes which is a spatially downscaled subset of ImageNet, the first-decay switching strategy indeed leads to better validation accuracy.
'This generalization improvement is not mentioned explicitly in (Ortiz et al. 2021 but can be clearly seen from Figures 7 and 8 in their paper.
3

---

Published as a conference paper at ICLR 2023
0.94
EOO
Topa1 Teat Accuacy
0.76-
Test Accurcyy
121
Test Accorse
0.94
0.00
0.92-
JVV
WW
0.24
0.00
0.50
0.58
1000
2000
3000
250
2000
-0030
00000
2000
100
150
200
250
Epochs
Epochs
Epechs
(a) CIFAR-I0, 10, start from random.
(b) CIFAR-I0 10, start from #250.
(c) ImageNet, start from #100.
Top-3 Test Acurasy
0.75-
0.50--9
0:15
0.76
Test Accuncy
0.5
Tixt Actunay
0.500
0.70
0.25
tram
0.64
0.16
0.32
0.00
50
100
150
0.91
30
150
750
2000
0.74
2
6
13
26
52
156
Epochs
Number of Local Steps H (log stale)
Number of Local Steps M (log stale)
(d) ImageNet, first phase = 3.2.
(c) CIFAR-10, test acc v.s. H.
(f) ImageNet, test acc v.s. H.
Figure 2: Ablation studies on v). H and training time in the same setting as Figure For (a)(d), we train from random initialization. For (f)c))e)f)) we start training from the checkpoints saved at the switching time points in Figure (ecure[[ppprr (epoch #250 for CIFAR-10 10 and epoch #100 for ImageNet). See Appendix M.2 for training details.
2.2
KEY FACTORS: SMALL LEARNING RATE AND SUFFICIENT TRAINING TIME
All the above papers agree that Post-local/Loca SGD improves upon SGD to some extent. How- ever, it is in debate under what conditions the generalization benefit can consistently occur. We now conduct ablation studies to identify the key factors so that adding local steps improves the general- ization of SGD. We run parallel SGD and Local SGD with the same learning rate 77, local batch size Bioe. and number of workers K. We start training from the same initialization and compare their generalization after the same number of epochs. As Post-local SGD can be viewed as Local SGD starting from SGD- D-D-ppetrained model, the initial point in experiments can be either random or a checkpoint of SGD training. See Appendix C for implementation details and Appendix M.2 for more details about the experimental setup.
The first observation we have is that the generalization benefits can be reproduced on both CIFAR- 10 and ImageNet in setting (see Fi. We remark that Post-local SGD and SGD in Lin et al. (2020)): Ortiz et al. (2021) are implemented with accompanying Nesterov momentum terms. The learning rate also decays couple of times in training with Local SGD. Nevertheless, our experi- ments show that the Nesterov momentum and learning rate decay are not necessary for Local SGD to generalize better than SGD. Our main finding after further ablation studies is summarized below:
Finding 2.1. Given sufficiently small learning rate and sufficiently long training time, Local SGD exhibits better generalization than SGD. if the number of local steps H per round is tuned properly according to the learning rate. This holds for both training from random initialization and from pre-trained models.
Now go through each point of main finding. See also Appendix for more plots.
(1). Pretraining is not necessary. In contrast to previous works claiming the benefits of Post-local SGD over Local SGD (Lin et al. [200b] Ortiz et al. (02T), initialization also generalizes significantly better than SGD, as long as the learning rate is small and the training time sufficiently long (Figure 2(a)). Starting from a pretrained model may shorten the time to reach this generalization benefit to show up (Figure 2(b)). but it is not necessary.
observe that Local SGD with random
(2). Learning rate should be small. We experiment with a wide range of learning rates con- clude that setting small learning rate is necessary. The learning rate is 0.32 for Figures ((a)
[[a]andd and 2(b) and is 0.16 for Figure 2(c) As shown in Figure 2(d) Local SGD encounters optimization difficulty in the first phase where is large (1) = 3.2),, resulting in inferior final test accuracy. Even for training
from a pretrained model, the generalization improvement of Local SGD disappears for large learn- ing rates (e.g., 17 1.6 in Figure In contrast, if a longer training time is allowed, reducing the learning rate of Local SGD does not lead to test accuracy drop (Figure 5(c)).
(3). Training time should be long enough. To investigate the effect of training time, in Fig- ures 2 and 2(c) we extend the training budget for the Post-local SGD experiments in Figure and observe that longer training time leads to greater generalization improvement upon SGD. On
the other hand, Local SGD generalizes worse than SGD in the first few epochs of Figures 2(a and 2(c) see Figures 5(a) and 5(b) for an enlarged view.
4

---

Published as a conference paper at ICLR 2023
(4). The number of local steps H should be tuned carefully. The number of local steps H has a complex interplay with the learning rate 17, but generally speaking, a smaller 1j needs a higher H to achieve consistent generalization improvement. For CIFAR-I0 with post-local training budget of 250 epochs (see Figure [[e)), the test accuracy first rises as H increases, and begins to fall as H exceeds some threshold for relatively large 1) (e.g., 1] 2 0.5) while keeps growing for smaller 7) (e.g., < 0.5). For ImageNet with post-local training budget of 50 epochs (see Figure 2()). the test accuracy first increases and then decreases in H for all learning rates.
Reconciling previous works. Our finding can help to settle the debate presented in Section 2.1 to a large extent. Simultancously requiring a small learning rate and sufficient training time poses a trade-off when learning rate decay is used with limited training budget: switching to Local SGD earlier may lead to a large learning rate, while switching later makes the generalization improve- ment of Local SGD less noticeable due to fewer update steps. It is thus unsurprising that first-decay switching strategy is not always the best. The need for sufficient training time does not contradict with Ortiz et al. (2021)'s is conjecture that Local SGD only has a "short-term" generalization benefit. In their experiments, the generalization improvement usually disappears right after the next learn- ing rate decay (instead of after fixed amount of time). We suspect that the real reason why the improvement vanishes is that the number of local steps H was kept as a constant, but our finding suggests tuning H after v) changes. In Figure 5(e) we reproduce this phenomenon and show that increasing H after learning rate decay retains the improvement.
Generalization performances at the optimal learning rate of SGD. In practice, the learning rate of SGD is usually tuned to achieve the best training loss/validation accuracy within a fixed training budget. Our finding suggests that when the tuned learning rate is small and the training time is sufficient, Local SGD can offer generalization improvement over SGD. As an example, in our experiments on training from an SGD-pretrainee model, the optimal learning rate for SGD is 0.5 on CIFAR-10 (Figure 2(e) and 0.064 on ImageNet (Figure 20) With the same learning rate as SGD, the test accuracy is improved by 1.1% on CIFAR-10 10 and 0.3% on ImageNet when using Local SGD with H = 750 and H = 26 respectively. The improvement could become even higher the learning rate of Local SGD is carefully tuned.
THEORETICAL ANALYSIS OF LOCAL SGD: THE SLOW SDE
In this section, we adopt an SDE-based approach to rigorously establish the generalization benefit of Local SGI in a general setting. Below, we first identify the difficulty of adapting the SDE framework to Local SGD. Then, we present our novel SDE characterization of Local SGD around the manifold of minimizers and explain the generalization benefit of Local SGD with our SDE.
Notations. We follow the notations in Section [ We denote by n the learning rate, K the number of workers, B the (global) batch size, Bio := B/K K the local batch size, H the number of local steps, {(0;c) <) the loss function for a data sample Ç. and D the training distribution. Furthermore, we define (())= = Bgw [e( E)] the expected loss, E(0)
(+((:)) (ovg-b/V((B:: as the noise covariance of gradients O. Let (Wi)200 denote the standard Wiener process. For mapping F Rd Rdd, denote by @F) the Jacobian at O and 82 F(O) the second order derivative O. Furthermore, for any matrix M E Raddd, 0² FF(()) M] = Eiela 08 M)ee where e; is the i-th vector of the standard basis. We write (VL)(O) as (3C(O) for short.
Local SGD, We use the following formulation of Local SGD for theoretical analysis. See also Oopendix[[for. the pseudocode. Local SGD proceeds in multiple rounds of model averaging, where each round produces global iterate Ø(s), In the (s + 1) round, every worker & [K] starts with its local copy of the global iterate o(:)) + g(s) and does H steps of SGD with local batches. In the f-th local step of the k-th worker, it draws a local batch of Bioe == B/K K independent samples from shared training distribution DD and updates as follows:
(()) (()
Sk.6.1.---; Btot,Broc -k.t.B..
(s) 0 t o(.)) 7 where (x) 1 Bise
Σ
V K.0i (l...) t t...... H (2)
I Bioc
The local updates on different workers are independent of each other as there is no communication. After finishing the H local steps, the workers aggregate the resulting local iterates k, and assign
the average to the next global iterate: ((++1) + *& -==1 K.
5

---

Published as a conference paper at ICLR 2023
3.1 DIFFICULTY OF ADAPTING THE SDE FRAMEWORK TO LOCAL SGD
A widely-adopted approach to understanding the dynamics of SGD is to approximate it from a con- tinuous perspective with the following SDE (3) which call the conventional SDE approximation. Below, we discuss why it cannot be directly adopted to characterize the behavior of Local SGD.
dx = V V
(3)
It is proved by Li et al. (2019a) that this SDE is first-order approximation to SGD, where each discrete step corresponds to a continuous time interval of 17. Several previous works adopt this SDE approximation and connect good generalization to having a large diffusion term v 1 > AWW in the SDE (Jastrzebski et al. 2017: Smith et al. 2020), because suitable amount of noise can be necessary for large- batch training to generalize well (see also Appendix
According to Finding 2.1 it is tempting to consider the limit - 0 and see if Local SGD can also be modeled via a variant of the conventional SDE. In this case the typical time length that guarantees a good SDE approximation error is O discrete steps (Li et al. 2019a 2021a). However, this time scaling is too short for the difference to appear between Local SGD and SGD. Indeed, Theorem 3.1 below shows that they closely track each other for steps.
steps.
Theorem 3.1. Assume that the loss function is Ca -s-smooth with bounded second and third order derivatives and that (e(0; E) is bounded. Let 'T' >> 0 be constant, O1*) be the s-th global iterate of Local SGD and wi be the t-th iterate of SGD with the same initialization wy ( and same 9. Bioc; K. Then any H S I and 8 = O(poly(n) it holds with probability at least 1 - 5 that
for all s S T q ||0(c] w.H l|2 O1 log )
We defer the proof to Appendix | See also Appendix D for Lin et al. (2020b 's attempt to model Local SGD with multiple conventional SDEs and discussions on why does not give much insight.
3.2 SDE APPROXIMATION NEAR THE MINIMIZER MANIFOLD
Inspired by a recent paper (Li et al., 2021b), our strategy to overcome the shortcomings of the conventional SDE is to design new SDE that can guarantee a good approximation for Onn discrete steps, much longer than the (17 discrete steps for the conventional SDE. Following their setting, we assume the existence of manifold I consisting only of local minimizers and track the global iterate ((*) around I' after it takes Ö(n steps to approach I. Though the expected gradient VL is near zero around I, the dynamics are still non-trivial because the noise can drive the iterate move significant distance in Onn-2) steps.
Assumption 3.1. The loss function C(-) and the matrix square root of the noise covariance Σ EA)(.) are -mo-smothh. Besides, we assume that [[ E)lla is bounded by constant for all O and &. Assumption 3.2. T is C -smooth, (d mi)ddinensi submanifold Rdd where any S E r a local minimizer of L. For all Ç E T, rank(V* = m. Additionally, there exists an open neighborhood of r, denoted as U, such that I arg mineet L()).
Assumption 3.3. I is compact manifold.
The smoothness assumption on L is generally satisfied when we use smooth activation functions, such Swish (Ramachandran et al. 2017), softplus and GeLU (Hendrycks & Gimpel 20166), which work equally well as ReLU in many cireumstances. The existence of a minimizer manifold with rank L(C)) m has also been made as key assumption in Fehrman et al. (022)); Li et al. (2021b); Lyu et al. (2022). where (annk((V(())))))))) = m ensures that the Hessian is maximally non- degenerate on the manifold and implies that the tangent space at Ç I' equals the null space of V²L((). The last assumption is made to prevent the analysis from being too technically involved.
Our SDE for Local SGD characterizes the training dynamics near I. For ease of presentation, we define the following projection operators Φb Pc for points and differential forms respectively.
Definition 3.1 (Gradient Flow Projection). Fix a point m & Г. For x E Rd, consider the gradient (t)(t) ((x(()) with =0)) = x. We denote the gradient flow projection of x as ((x).
flow
(xx) = lim xt) if the limit exists and belongs to I; otherwise, Φ(x) -
Definition 3.2. For any Ç E T and any differential form AdW, + bdt in Itô calculus, where A is matrix and b is a vector, we use Po(AWW, + bdl) as a shorthand for the differential form (x((S) AdW, (ab(C)b + ( dt.

---

Published as a conference paper at ICLR 2023
See Øksendal (2013) for an introduction to Itō calculus. Here Pe equals + by Itō calculus, which means that P3 projects an infinitesimal step from Ç, so that Ç after taking the projected step does not leave the manifold r. It can be shown by simple calculus that ($(() equals the projection matrix onto the tangent space of T at Ç. We decompose the noise covariance E(g for Ç E T into two parts: the noise in the tangent space (i) (())))((() and the noise in the rest 2x({) Σ({) [j(()) Now we are ready to state SDE for Local SGD.
(C)
Definition 3.3 (Slow SDE for Local SGD). Given 1), H > 0 and So r, define ((t) as the solution of the following SDE with initial condition ((0) = Co:
ac(I) Ps
1 VB
E 12 (C)dW,
EB V ((C))
2 V V (())
(4)
(a) diffasion
(b) drif--1
drif----11
Here Ed(C). 1 (C) E Edx are defined as
EO(C) = = EE
Arta; ( (E) 0.) D.ue.
(5)
(C) : = Winlt (Act2a9) (Eo 0.0, v.
(6)
Art
where {w.]e is set of eigenvectors of V (Q) that forms an orthonormal eigenbasis. and A Aa are the corresponding eigenvalues. Additionally, vbxx) = ==1+z for x + and
(0) =
The use of Pg keeps ((t) on the manifold I through projection. Σj (C) introduces diffusion
term the SDE in the tangent space. The two drift terms involve Eg (-) and (-). which can be intuitively understood as rescaling the entries of the noise covariance in the eigenbasis of Hessian. In the special case where V tiag(()," ... Aa) E Rdxd, we have Eo... = xutXy Eo,i.-
Vi U(nH(), Sa,i.. (xr) a monotonically increasing function, which goes from to as +x; x goes from 0 to infinity (see Figure
We name this SDE as the Slow SDE for Local SGD because we will show that each discrete step of Local SGD corresponds to a continuous time interval of 172 instead of an interval of r) in the conventional SDE. In this sense, track longer horizon). This Slow SDE is inspired by Li et al. (2021b). Under nearly the same set of assumptions, they proved that SGD can be tracked by an SDE that is essentially equivalent (4) with K = 1, namely, without the drift-Il term.
SDE is "slower" than the conventional SDE (and hence can
dc(t) = Pc
2/ (C)dW, VB
1
V C(C(()
2B
diffusion
(b) drift-1
We refer to (7) as the Slow SDE SGD. We remark that the drfit-ll term in (4) is novel and the key to separate the generalization behaviors of Local SGD and SGD in theory. We will discuss this point later in Section 3.3 Now we present our SDE approximation theorem for Local SGD.
3.7 to 3.3 hold. Let T > 0 be constant and ((t) be the solution to with the initial condition ((0) D (()) E r. If H is set to for some constant a > 0,
Theorem 3.2. Let Assumptions
then for any C* smooth function g( maxoses a Why
[lg(((0()) E[gg((s((SS
where Ō(-) hides log factors and constants that are independent of but can depend on g g(0). Theorem 3.3. For = (ppoly (1)),, with probability at least 1 - 8, it holds for all O log H) < Vo
S T. that ((())) E I and (0()) (((())) Oly am gt where O(-) hides constants an independent of n, and 8.
Theorem 3.2 suggests that the trajectories of the manifold projection and the solution the Slow SDE 4 close to each other in the weak approximation sense. That is, ( and (C(t) cannot be distinguished by evaluating test functions from a wide function class, including all poly- nomials. This measurement of closeness between the iterates of stochastic gradient algorithms and their SDE approximations is also adopted by Li et al. (2019a) 202 1a); Malladi et al. (2022). but their analyses are for conventional SDEs. Theorem 3.3 further states that the iterate 01s) keeps close to its manifold projection after the first few rounds.
7

---

Published as a conference paper at ICLR 2023
Remark 3.1. To connect to Finding 2.1
we remark that our theorems (1) do not require the model
to be pre-trained (as long as the gradient flow starting with 0(0) converges to I); (2) give better bounds for smaller ny (3) characterize a long training horizon n7. The need for tuning H will be discussed in Section 3.3.3
Technical Contribution. The proof technique for Theorem 3.2 is novel and significantly different from the Slow SDE analysis of SGD in Li et al. (2021a). Their analysis uses advanced stochastic calculus and invokes Katzenherger's theorem Katzenberger 199T) to show that SGD converges to the Slow SDE in distribution, but no quantitative error bounds are provided. Also, due to the local updates and multiple aggregation steps in Local SGD, it is unclear how to extend Katzenberger's theorem to our case. To overcome this difficulty, we develop a new approach to analyze the Slow SDEs, which is based on the method of moments (Li et al. 2019a) and can provide the quantitative error bound (170-25) in weak approximation. See Appendix for our proof outline. A by-product of our result is the first quantitative approximation bound for the Slow SDE approximation for SGD, which can be easily obtained by setting K = 1.
3.3 INTERPRETATION OF THE SLOW SDES
In this subsection, we compare the Slow SDEs for SGI) and Local SGD and provide important insight into why Local SGD generalizes better than SGD: Local SGD strengthens the drift term in the Slow SDE, which makes the implicit regularization of stochastic gradient noise more effective.
3.3.1 INTERPRETATION OF THE SLOW SDE FOR SGD.
The Slow SDE for SGD 7 consists of the diffusion and drift-I terms. The former injects noise into the dynamics in the tangent space; the latter one drives the dynamics to move along the negative gradient of 2B (V), Ec( projected onto the tangent space, but ignoring the dependency of (C) on Ç. This can be connected to the class of semi- arii-gradient methods which only computes a part of the gradient (Mnih al. 2013 Sutton & Barto 1998, Brandfonbrener & Bruna 2020). In this view, the long- nengttermmn behavior of SGD is similar to a stochastic semi comi-gradient method minimizing the implicit regularizer 28 (V))(C), (C)) on the minimizer manifold of the original loss C.
Though the cemi-gradient method may not perfectly optimize its objective, the above argument reveals that SGD has deterministic trend toward the region with smaller magnitude of Hessian, which is commonly believed to correlate with better generalization (Hochreiter & Schmidhuber, 1997| Keskar et al. 2017; Neyshabur et al., [017)] Jiang et al.] (020)) (see Appendix for more discussions). In contrast, the diffusion term can be regarded as a random perturbation to this trend, which can impede optimization when the drift-I term is not strong enough.
Based on this view, we conjecture that strengthening the drift term of the Slow SDE can help SGD to better regularize the model, yielding a better generalization performance. More specifically, we propose the following hypothesis, which compares the generalization performances of the following generalized Slow SDEs. Note that ( 2/3 ]SSow SDE corresponds to the Slow SDE for SGD Definition 3.4. For K1, K2 2 0, define (MI, x2) -Slow SDE to be the following:
dc(1)= Pc (V F (C)dW, - K2 V2 ((C)
(8)
Hypothesis 3.1. Starting at a minimizer Ço E I, run (K1, K2) -Soo SDE and (KI, (k1,k)) -SSow SDE respectively for the same amount of time T > and obtain ((T),,((T)) If K2 Kiz. then the expected test accuracy at C(T) is better than that at (T)).
Due to the No Free Lunch Theorem, we do not claim that our hypothesis is always true, but we do believe that the hypothesis holds when training usual neural networks (e.g., ResNets, VGGNets) on standard benchmarks (e.g.. CIFAR-10, ImageNet).
Example: Training with Label Noise Regularization. To exemplify the generalization benefit of having a larger drift term, we follow a line of thcoretical works (Li et al. 2021b] Blanc et al. 2020) Damian et al. 2021 to study the case of training over-parameterrize neural nets with label noise regularization. For a C-class classification task, the label noise regularization is as follows: every time we draw sample from the training set, make the true label as it is with probability 1 - p, and replace it with any other label with equal probability J-.. When we use cross-entropy loss, the Slow SDE for SGD turns out to be a simple deterministic gradient flow on I (instead of a semi- gradient method) for minimizing the trace of Hessian: dç(()= = "B (rtr(V?L(C))dd where Vr
8

---

Published as a conference paper at ICLR 2023
stands for the gradient of the function f projected to the tangent space of T. Checking the validity of our hypothesis reduces to the following question: Is minimizing the trace of Hessian beneficial to generalization? Many works prove positive results in concrete settings, including the line of works we just mentioned. We refer the readers to Appendix Gfor further discussion.
3.3.2 LOCAL SGD STRENGTHENS THE DRIFT TERM IN SLOW SDE.
Based on Hypothesis 3.1 we argue that Local SGD improves generalization by strengthening the drift term of the Slow SDE. First, it can be seen from 4 that the Slow SDE for Local SGD has an additional drfit-II II term. Similar to the drift-I term of the Slow SDE for SGD, this drift-ll term drives the dynamics to move along the negative semi semi-gradient of K-1 ( (C)) (with the 2
dependeney of L (C) on Ç ignored). Combining it with the implicit regularizer induced by the drift-I term, we can see that the long-term behavior of Local SGD is similar to stochastic semi-gu semi-gradient method minimizing the implicit regularizer 3/ 2 C(C), Ec(C) K-1 2 C(C). (C)) on I.
we can see that (C) is basically a rescaling of the entries of (C) in the eigenbasis of Hessian, where the rescaling factor (nnH (A: + x,)) for each entry is between 0 and (see Figure 9 for the plot of 4)). When n is small, the rescaling
Comparing the definitions of Σ. (C) (5) and (C) (o)
factors should be close to (0) = O, then (C) = 0, leading to almost no additional regularization. On the other hand, when nH is large. the rescaling factors should be close to (t(00) = 1, so
(C) R (C). We can then merge the two implicit regularizers as K (V)C(S), Ec()), and (4)
2B
becomes the ( 1 K) J-Slow SDE, which is restated below:
B 2B
d((t)= = Pc v II (C)dW, - 2 V (C)]dt)
E /"
From the above argument we know how the Slow SDE of Local SGD 4 changes as nHH transitions
from to +00. Initially, when 1/ 0, (4) is the same the ( 2B -Slow SDE for SGD. Then
increasing H strengthens the drift term of ( As nH - +00, transitions to the B 28 ) SDE, where the drift term becomes K times larger.
According to Hypothesis 3.1, the & K )-low SDE generalizes better than the ( 2B )-Slow SDE, so Local SGD with H - should generalize better than SGD. When hyH is chosen realistically as finite value, the generalization performance of Local SGD interpolates between these two cases, which results in worse generalization than HH +o0 but should still be better than SGD.
3.3.3 THEORETICAL INSIGHTS INTO TUNING THE NUMBER OF LOCAL STEPS
Based on Slow SDE approximations, now discuss how the number of local steps H affects the generalization of Local SGD. When is small but finite, tuning H offers a trade-off between regularization strength and SDE approximation quality. Larger a == H makes the regularization stronger in the SDE (as discussed in Section 3.3.2), but the SDE itself may lose track of Local SGD, which can be seen from the error bound O (ana an log(a/nó) in Theorem 3.3 Therefore, we expect the test accuracy to first increase and then decrease as we gradually increase H. Indeed, we observe in Figures 2(e) and 2(1) that the plot of test accuracy versus H is unimodal for each
It is thus necessary to tune H for the best generalization. When H is tuned together with other hyperparameters, such as learning rate 7). our Slow SDE approximation recommends setting H to be at least 11) so that a := 1/H does not vanish in the Slow SDE. Since larger a gives a stronger regularization effect, the optimal H should be set to the largest value so that the Slow SDE does not lose track of Local SGD. Indeed, we empirically observed that when is tuned optimally, a increases 77 decreases, suggesting that the optimal H grows faster than 2'n See Figure 5
CONCLUSIONS
In this paper, we analyze the long-term generalization behavior of Local SGD in the small learning rate regime by deriving the Slow SDE for Local SGD as a generalization of that for SGD (Li et al., 2021b). We attribute the generalization improvement over SGD to the larger drift term in the SDE for Local SGD. Our empirical validation shows that Local SGD indeed induces generalization benefits with small learning rate and long enough training time. The main limitation of our work is that our analysis does not imply any direct theoretical separation between SGD and Local SGD in test accuracy, which requires much deeper understanding of the loss landscape and the Slow SDEs and is left for future work. Another direction for future work is to design distributed training methods that provably generalize better than SGD based on the theoretical insights obtained from Slow SDEs.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENT AND DISCLOSURE OF FUNDING
The work of Xinran Gu and Longbo Huang is supported by the Technology and Innovation Major Project of the Ministry of Science and Technology of China under Grant 2020AAA0108400 and 2020AAA0108403. the Tsinghua University Initiative Scientific Research Program, and Tsinghua Precision Medicine Foundation 10001020109. The work of Kaifeng Lyu and Sanjeev Arora is supported by funding from NSF, ONR, Simons Foundation, DARPA and SRC.
REFERENCES
Kwangjun Ahn, Jingzhao Zhang, and Suvrit Sra. Understanding the unstable convergence of gra- dient descent. In Kamalika Chaudhuri, Stefanie Jegelka, Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings the 39th International Conference Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 247-257.PMLL, 17-23 Jul 2022.
Debraj Basu, Deepesh Data, Can Karakus, and Suhas Diggavi. Qsparse-local-SGD: Distributed SGD with quantization, sparsification and local computations. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Yoshua Bengio. Practical Recommendations for Gradient-Busee Training of Deep Architectures, pp. 437-478 Springer Berlin Heidelberg, Berlin, Heidelberg, 2012. ISBN 978-3-642-35289- doi: 10.1007/978- 3-1----44---289--- 26.
Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by ornstein-uhlenbeck like process. In Jacob Abernethy and Shivani Agarwal (eds.), Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pp. 483-513. PMLR, 09-12 Jul 2020.
David Brandfonbrener and Joan Bruna. Geometric insights into the convergence of nonlinear TD learning. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
Jianmin Chen, Xinghao Pan, Rajat Monga, Samy Bengio, and Rafal Jozefowicz. Revisiting dis- 2016.
tributed synchronous SGD. arXiv preprint 1604:
Kai Chen and Qiang Huo. Scalable training of deep learning machines by incremental block train- ing with intra-block parallel optimization and blockwise model-update filtering. In 2016 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 5880-50000 5884,
2016. doi: 10. 2A110091AASSPP2216677728800
Alex Damian, Tengyu Ma, and Jason D. Lee. Label noise SGD provably prefers flat global mini- mizers. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning. volume 70 of Proceedings of Machine Learning Research, pp- 1019-1028. PMLR, 06-11 Aug 2017.
Aijun Du and JinQiao Duan. Invariant manifold reduction for stochastic dynamical systems. Dy- namic Systems and Applications, 16:6816696, 2007.
KJ Falconer. Differentiation of the limit mapping in a dynamical system. Journal of the London Mathematical Society, 2(2):356-372, 1983.
Benjamin Fehrman, Benjamin Gess, and Arnulf Jentzen. Convergence rates for the stochastic gradi- ent descent method for non-convex objective functions. Journal of Machine Learning Research, 21:136, 2020.
Damir Filipović. Invariant manifolds for weak solutions to stochastic equations. Probability theory and related fields, 118(3):323-34 2000.
10

---

Published as a conference paper at ICLR 2023
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness- aware minimiza- tion for efficiently improving generalization. International Conference on Learning Represen- tations, 2021.
Margalit Glasgow, Honglin Yuan, and Tengyu Ma. Sharp bounds for federated averaging (Local SGD) and continuous perspective. In International Conference Arlificial Intelligence and Statistics, pp. 9050-9090. PMLR, 2022.
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, An- drew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch SGD: Training imagenet in hour. arXiv preprint N.kiv:/706.02677, 2017.
Farzin Haddadpour, Mohammad Mahdi Kamani, Mehrdad Mahdavi, and Viveck Cadambe. Local SGD with periodic averaging: Tighter analysis and adaptive synchronization. Advances in Neural Information Processing Systems, 32, 2019.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE international conference on computer vision, pp. 1026-1034, 2015.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint
wrXiv:1606.08415 2016.
Sepp Hochreiter and Jürren Schmidhuber. Flat minima. Neural computation, 9()):1-42, 1997.
Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: closing the general- ization gap in large batch training of neural networks. Advances neural information processing systems, 30, 2017.
Wenging Hu, Chris Junchi Li, Lei Li, and Jian-Guo Liu. On the diffuusion approximation of noncon- vex stochastic gradient descent. arXiv preprint arXiv: 1705.0070... 2017.
Hikaru Ibayashi and Masaaki Imaizumi. Exponential escape efficiency of SGD from sharp minima in non-stationary regime. arXiv preprint arXiv:21 2021.
Stanistaw Jastrzebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in SGD, arXiv preprint arXiv:1711.06233, 2017.
Xianyan Jia, Shutao Song, Wei He, Yangzihao Wang, Haidong Rong, Feihu Zhou, Liqiang Xie, Zhenyu Guo, Yuanzhou Yang, Liwei Yu, et al. Highly scalable deep learning training system with mixed-precision: Training imagenet in four minutes. Advances Neural Information Processing Systems, 2018.
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantas- tic generalization measures and where to find them. In International Conference on Learning Representations, 2020.
Peter Kairouz, Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Ad- vances and open problems in federated learning. Foundations and Trends@ in Machine Learning, 14(1-2):1-210, 2021.
Sai Praneeth Karimireddy, Satyen Kale. Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International Conference on Machine Learning, pp. 5132-5143. PMLR, 2020.
G. Katzenberger. Solutions of a stochastic differentiall equation forced onto a manifold by large drift. The Annals of Probability, 19(4):1587 1628, 1991.
11

---

Published as a conference paper at ICLR 2023
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On large- batch training for deep learning: Generalization gap and sharp minima. In International Conference Learning Representations, 2017.
Ahmed Khaled, Konstantin Mishchenko, and Peter Richtárik. Tighter theory for local SGD on iden- tical and heterogeneous data. In International Conference on Artificial Intelligence and Statistics, pp. 4519-4529. PMLR, 2020.
Bobby Kleinberg, Yuanzhi Li, and Yang Yuan. An alternative view: When does SGD escape lo- cal minima? In Jennifer Dy and Andreas Krause (eds.). Proceedings of the 35th International Conference Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp- 2698-2707. PMLR, 10-15 Jul 2018.
Alex Krizhevsky. One weird trick for parallelizing convolutional neural networks. arXiv preprini arXiv:1404.5997, 2014.
Alex Krizhevsky et al. Learning multiple layers of features from tiny images. 2009.
Guillaume Leclerc, Andrew Ilyas, Logan Engstrom, Sung Min Park, Hadi Salman, and Aleksander Madry. ffev. bttps::/gitthii
com wttp://w....///Littee//////E//
2022.
Yann A. LeCun, Léon Bottou, Genevieve B. Orr, and Klaus-Robert Müller. Efficient BackProp, pp- 9-48. Springer Berlin Heidelberg, Berlin, Heidelberg, 2012. ISBN 978-3-642-35289-8, doi:
Qianxiao Li. Cheng Tai, and Weinan Stochastic modified equations and dynamics of stochastic gradient algorithms i: Mathematical foundations, Journal of Machine Learning Research, 20(40); 1-47, 2019a
Xiang Li, Kaixuan Huang, Wenhao Yang. Shusen Wang, and Zhihua Zhang. On the convergence of fedavg on non-iid data. In International Conference on Learning Representations, 2019b.
Zhiyuan Li, Kaifeng Lyu, and Sanjeev Arora. Reconciling modern deep learning with traditional optimization analyses: The intrinsic learning rate. Advances Neural Information Processing Systems, 11114444-45555 2020.
Zhiyuan Li, Sadhika Malladi, and Sanjeev Arora. On the validity of modeling SGD with stochastic differential equations (sdes). Advances in Neural Information Processing Systems, 34:12712- 12725, 2021a.
Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after SGD reaches zero loss?-a mathematical framework. In International Conference on Learning Representations, 202
Zhiyuan Li, Tianhao Wang, and Dingli Yu. Fast mixing of stochastic gradient descent with normal- ization and weight decay. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022.
Tao Lin, Lingjing Kong, Sebastian Stich, and Martin Jaggi. Extrapolation for large-batch training in deep learning. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 6094-6104. PMLR, 13-18 Jul 2020a.
Tao Lin, Sebastian U. Stich, Kumar Kshitij Patel, and Martin Jaggi. Don't use large mini-batches, use Local SGD. In International Conference Learning Representations, 2020b.
Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normal- ization layers: Sharpness reduction, 2022.
Chao Ma and Lexing Ying. On linear stability of SGD and imput-smoothness of neural networks. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 16805-16817. Curran Associates, Inc., 2021.
12

---

Published as a conference paper at ICLR 2023
Sadhika Malladi, Kaifeng Lyu, Abhishek Panigrahi, and Sanjeev Arora. On the SDEs and scaling rules for adaptive gradient algorithms. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022.
Gideon Mann, Ryan T. McDonald, Mehryar Mohri, Nathan Silberman, and Dan Walker. Efficient large-scale distributed training of conditional maximum entropy models. In Advances in Neural Information Processing Systems 22, pp. 1231-1239, 2009.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-effiiien learning of deep networks from decentralized data. In Artificial intelli- gence and statistics, pp. 1273-1282. PMLR, 2017.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 518(7540):529-533 2015. 2015.
Behnam Neyshabur, Srinadh Bhojanapalli, David Mcallester, and Nati Srebro. Exploring gener- alization in deep learning. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Jose Javier Gonzalez Ortiz, Jonathan Frankle, Mike Rabbat, Ari Morcos, and Nicolas Ballas. Trade- offs of Local SGD scale: An empirical study, arXiv preprint arXiv:2110.08133, 2021.
Daniel Povey, Xiaohui Zhang, and Sanjeev Khudanpur. Parallel training of dnns with natural gradi- ent and parameter averaging. arXiv preprint arXiv:1410.7455, 2014.
Prajit Ramachandran, Barret Zoph, and Quoc V Le. Searching for activation functions. arXiv preprint arXiv: 2017.
Benjamin Recht, Christopher Ré, Stephen J. Wright, and Feng Niu. Hogwild: A lock-free approach to parallelizing stochastic gradient descent. In Advances in Neural Information Processing Sys- tems 24, pp. 693-701, 2011.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang. Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJCV), 1 115(3)211 25-22,,1, 2015. doi: 10.101 1263-015-0816-
Frank Seide, Hao Fu, Jasha Droppo, Gang Li, and Dong Yu. 1-bit stochastic gradient descent and its application to data-parallel distributed training of speech dnns. In Haizhou Li, Helen M. Meng, Bin Ma, Engsiong Chng. and Lei Xie (eds.), INTERSPEECH 2014, 15th Annual Conference of the International Speech Communication Association, Singapore, September 14-18, 2014, pp. 1058- 1062. ISCA, 2014. TRRhhtpp: / www. wwwp::iscasspppeee speech. wwwpeechh.rgg/.eooooe htsp//rcciiiveeiie http://ccvveeeeeespeeeeeeeeeeeee 2014/114 1058. 10m8.html
Christopher J. Shallue, Jaehoon Lee, Joseph Antognini, Jascha Sohl-Dickstein, Roy Frostig, and George E. Dahl. Measuring the effects of data parallelism neural network training. Journal of Machine Learning Research, 20( (1): -9, 2019.
K. Simonyan and A. Zisserman. Very deep convolutional networks for large-scale image recogni- tion. In International Conference on Learning Representations, May 2015.
Samuel Smith. Erich Elsen, and Soham De. On the generalization benefit of noise in stochastic gradient descent. In Hal Daumé and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 100990089900 90639006..110 PMLR, 13-18 Jul 2020.
Samuel L Smith, Benoit Dherin, David Barrett, and Soham De. On the origin of implicit regular- ization in stochastic gradient descent. In International Conference on Learning Representations, 2021.
Sebastian U Stich. Local SGD converges fast and communicates little. In International Conference Learning Representations, 2018.
13

---

Published as a conference paper at ICLR 2023
Nikko Strom. Scalable distributed DNN training using commodity GPU cloud computing. In IN- TERSPEECH 2015, 16th Anmual Conference of the International Speech Communication Associ- ation, Dresden, Germany, September 6-10, 2015, pp. 1488-1492. ISCA, 2015.
Hang Su and Haoyu Chen. Experiments parallel training of deep neural network using model averaging. arXiv preprint arXiv:1507.012239 2015.
Richard S. Sutton and Andrew G. Barto. Reinforcement learning an introduction. Adaptive com-
putation and machine learning. MIT Press, 1998. ISBN 778-0-2622193998-6
Jianyu Wang and Gauri Joshi. Adaptive communication strategies to achieve the best error-runtime trade-off in local-update SGD. Proceedings of Machine Learning and Systems, 1:212-229, 2019.
Jianyu Wang and Gauri Joshi. Cooperative SGD: A unified framework for the design and analysis of local-upe local-update SGD algorithms. Journal of Machine Learning Research, 22(213):1 2021.
Jianyu Wang, Rudrajit Das, Gauri Joshi, Satyen Kale, Zheng Xu, and Tong Zhang. On the unreasonable effectiveness of federated averaging with heterogeneous data. arXiv preprint arXiv:2206.044 047233 2022.
Blake Woodworth, Kumar Kshitij Patel, Sebastian Stich, Zhen Dai, Brian Bullins, Brendan Mcma- han, Ohad Shamir, and Nathan Srebro. Is local sgd hetter than minibatch sgd? In International Conference on Machine Learning, pp- 10334- 10343. PMLR, 2020a.
Blake E Woodworth, Kumar Kshitij Patel, and Nati Srebro. Minibatch vs Local SGD for heteroge- neous distributed learning. Advances in Neural Information Processing Systems, 33:6281-6299 2020b.
Lei Wu, Chao Ma, and Weinan E. How sgd selects the global minima in werrppaaameterized learning: A dynamical stability perspective. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
Zeke Xie, Issei Sato, and Masashi Sugiyama. A diffusion theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima. In International Conference on Learning Representations, 2021.
Yang You, Zhao Zhang, Cho-Jui Hsich, James Demmel, and Kurt Keutzer. Imagenet training in minutes. Proceedings of the 47th International Conference on Parallel Processing, pp. 1-10, 2018.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsich. Large batch optimization for deep learn- ing: Training BERT in 76 minutes. In International Conference on Learning Representations, 2020.
Hao Yu, Sen Yang, and Shenghuo Zhu. Parallel restarted SGD with faster convergence and less communication: Demystifying why model averaging works for deep learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 5693-5700, 2019.
Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems, 33: 15333-15393, 2020.
Xiaohui Zhang, Jan Trmal, Daniel Povey, and Sanjeev Khudanpur: Improving deep neural network acoustic models using generalized maxout networks. In 2014 IEEE International Conference on Acoustics. Speech and Signal Processing (ICASSP), pp. 215-219, 2014. 2014.6853589.
doi: 10.1109//C
Fan Zhou and Guojing Cong. On the convergence properties of k-step averaging stochastic gra- dient descent algorithm for nonconvex optimization. In Proceedings of the Twenty-Sevent Inter- national Joint Conference Artificial Intelligence. IJCAI-18, pp. 3219-3227. International Joint Conferences on Artificial Intelligence Organization, 7 2018. doi: 10.24963/jcci.200 2018/4477 URL https://doi wor. 24963/4933511 201ii/.011/4444
14

---

Published as a conference paper at ICLR 2023
Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects. arXiv preprint arXiv:1803.00195, 2018.
Martin Zinkevich, Markus Weimer, Lihong Li, and Alex Smola. Parallelized stochastic gradient descent. In J. Lafferty, C. Williams, J. Shylo, R. Zemel, and A. Culotta (eds.), Advances in Neural Information Processing Systems, volume 23. Curran Associates, Inc., 2010.
Bernt Øksendal. Stochastic differential equations: an introduction with applications. Springer Science & Business Media, 2013.
15

---

Published as a conference paper at ICLR 2023
CONTENTS
Introduction
2
When does Local SGD Generalize Better?
The Dehate on I Local SGD
2.2
Key Factors: Small Learning Rate and Sufficient Training Time
3
Theoretical Analysis of Local SGD: The Slow SDE
Difficulty of Adapting the SDE Framework to Local SGD
3.2
SDE Approximation near the Minimizer Manifold
6
3.3
Interpretation of the Slow SDEs
3.3.1
Interpretation of the Slow SDE for SGD.
8
3.3.2
Local SGD Strengthens the Drift Term in Slow SDE
3.3.3
Theoretical Insights into Tuning the Number of Local Steps
4
Conclusions
Additional Related Works
18
Additional Discussions
19
Implementation Details of Parallel SGD, Local SGD and Post-local SGD
20
D Modeling Local SGD with Multiple Conventional SDEs
23
Additional Interpretation of the Slow SDES
23
E.1 Understanding the Diffusion Term in the Slow SDE
23
The Effect of Global Batch Size on Generalization
24
Additional Experimental Results
25
Discussions on Local SGD with Label Noise Regularization
27
The Slow SDE for Local SGD with Label Noise Regularization
27
G.2 The Equivalence of Enlarging the Learning Rate and Adding Local Steps
28
H Deriving the Slow SDE after Applying the LSR
28
Proof of T'eoorerm...
30
Proof Outline of Main Theorems
33
Proof Details of Main Theorems
33
K. Additional Notations
34
16

---

Published as a conference paper at ICLR 2023
K.2
Computing the Derivatives of the Limiting Mapping
34
Preliminary Lemmas for GD and GF
35
K.4
Construction of working zones
38
Phase Iterate Approaching the Manifold
................................
39
K.5.J
Additional notations
39
K.5.2
Proof for Subphase
39
K.5.3
Proof for Subphase 2
43
K.6
Phase 2: Iterates Staying Close to Manifold
46
K.6.1
Additional notations
46
K.6.2
Proof for the High Probability Bounds
46
K.7
Summary of the dynamics and Proof of Theorems J.1 and J.2
51
K.8 Proof of Theorem 3.3
52
K.9
Computing the Moments for One "Giant Step"
53
K.10 Proof of Weak Approximation
66
K. 10.1 Preliminaries and additional notations
67
K.10.2 Proof of the approximation in context
68
Deriving the Slow SDE for Label Noise Regularization
72
M Experimental Details
74
M. Post-local SGD Experiments in Section
74
M.2 Experimental Details for Figures and
74
M.3 Details for Experiments in Figure
75
M.4 Details for Experiments on the Effect of the Diffusion Term
75
M.5 Details for Experiments on the Effect of Global Batch Size
76
M.6
Details for Experiments on Label Noise Regularization
76
17

---

Published as a conference paper at ICLR 2023
A ADDITIONAL RELATED WORKS
Optimization aspect of Local SGD. Local SGD a communication--fficieee variant of parallel SGD, where multiple workers perform SGD independently and average the model parameters peri- odically. Dating back Mann et al. (2009) and Zinkevich et al. (2010). this strategy has been widely adopted to reduce the communication cost and speed up training in both scenarios of data center dis- tributed training (Chen & Huo 2016 Zhang et al. 2014 Povey et al. 2014 Su & Chen 2015) and Federated Learning (McMahan al. 2017 Kairouz et al. 2021). To further accelerate training, Wang & Joshi (2019) and Haddadpour et al. (2019) proposed adaptive schemes for the averag- ing frequency, and Basu et al. (2019) combined Local SGD with gradient compression. Motivated to theoretically understand the empirical success of Local SGD, a lot of researchers analyzed the data and convex/non-convex objective functions. Among them, Yu et al. (019); [Stich] (2018); Khaled et al. (2020): |Wodwooth et al. (2020a) focus the homogeneous setting where data for each worker
convergence rate of Local SGD under various settings, e.g.. honnggenooushheeerggeneeee
independent and identically distributed (IID). Li et al. (2019)) Karimireddy et al. (2020); Glasgow al. (2022; Woodworth et al. (2020b); Wang al. (022) study the heterogencous setting, where workers have non-IlI data and local updates may induce "client drift" (Karimireddy et al., 2020) and hurt optimization. The error bound of Local SGD obtained by these works typically inferior to that of SGD with the same global batch size for fixed number of iterations/epochs and becomes worse as the number of local steps increases, revealing a trade-off between less communication and better optimization. In this paper, we interested in the generalization aspect of Local SGD in the homogencous setting, assuming the training loss can be optimized to a small value.
Gradient noise and generalization. The effect of stochastic gradient noise generalization has been studied from different aspects, e.g., changing the order of learning different patterns Li et al. (2019)). inducing an implicit regularizer in the second-order SDE approximation Smmihh et al. (2021)
Li et al. (2019)). Our work follows a line of works studying the effect of noise in the lens of sharpness, which is long believed be related to generalization Hochreiter & Schmidhuber 1997);
Neyshabur et al. (2017). Keskar et al. (2017) empirically observed that large-batch training leads to worse generalization and sharper minima than small-batch training. |Wu et al.] (2018); Hu et al. (2017): Ma & Ying (2021) showed that gradient noise destabilizes the training around sharp min- ima, and Kleinberg al. (0188); Zhu et al., (2018); Xie et al. (202)); Ibayashi & Imaizumi (2021)
quantitatively characterized how SGD escapes sharp minima. The most related papers are Blanc
al. (2020): [Damian et al.) (202T): Li et al. (2021b), which focus on the training dynamics near a manifold of minima and study the effect of noise on sharpness (see also Section 3.2). Though the
mathematical definition of sharpness may be vulnerable to the various symmetries in deep neural
nets (Dinh et al. 2017). sharpness still appears to be one of the most promising tools for predicting generalization Qlang et al. 2020 Foret et al. (02I).
Improving generalization in large-batch training. The generalization issue of the large-batch (or full-batch) training has been observed as early as (Bengio) 2012 LeCun et al. 2012). As men- tioned in Section []] the generalization issue of large-batch training could be due to the lack of sufficient amount of stochastic noise. To make up the noise in large-batch training, Krizhevsky 2014); Goyal et al. (0177) empirically discovered the Linear Scaling Rule for SGD, which sug- gests enlarging the learning rate proportionally to the batch size. Jastrzębski et al. (2017) adopted an SDE- analysis to justify that this scaling rule indeed retains the same amount of noise as small-batch training (see also Section However, the SDE approximation may fail if the learning rate is too large (Li et al. 2021a), especially in the early phase of training before the first learning rate decay (Smith et al.) (020). Shallue et al. (2019) demonstrated that generalization gap between small- and latch-batch training can also depend on many other training hyperparameters. Besides enlarging the learning rate, other approaches have also been proposed to reduce the gap, including training longer (Hoffer et al.) 2017). learning rate warmup (Goyal et al. 2017). LARS (You et al. 2018), LAMB (You et al. 22220. In this paper, focus on using Local SGD to improve general- ization, but adding local steps is a generic training trick that can also be combined with others, e.g.,
Local LARS (Lin et al. 2020b), Local Extrap-SGD (Lin et al. 2020a).
18

---

Published as a conference paper at ICLR 2023
B ADDITIONAL DISCUSSIONS
Connection to the conventional wisdom that the diffusion term matters more. As mentioned in Section 3.1 it is believed in the literature is that large diffusion term in the conventional SDE leads to good generalization. One may think that the diffusion term in the Slow SDE corresponds to that in the conventional SDE, and thus enlarging the diffusion term rather than the drift term should lead to better generalization. However, we note that both the diffusion and drift terms in the Slow SDEs result from the long-term effects of the diffusion term in the conventional SDE (Slow SDEs become stationary if Σ 0). This means our view characterizes the role of gradient noise in more detail, and therefore, goes one step further on the conventional wisdom.
Slow SDEs for neural nets with modern training techniques. In modern neural net training, it common to add normalization layers and weight decay (L (L--regularization) for better optimization and generalization. However, these techniques lead to violations of our assumptions, e.g., fixed point exists in the regularized loss (Li et al., 2020 Ahn et al. 2022). Still, a minimizer manifold can be expected to exist for the unregularized loss. Li al. (2022) noted that the drift and diffusion around the manifold proceeds faster in this case, and derived a Slow SDE for SGD that captures O log 5/ discrete steps instead of O (- We believe that our analysis also be extended to this case, and that adding local steps still results in the effect of strengthening the drift term.
19

---

Published as a conference paper at ICLR 2023
C IMPLEMENTATION DETAILS OF PARALLEL SGD, LOCAL SGD AND POST-LOCAL SGD
In this section, we present the formal procedures for Parallel SGD, Local SGD and Post-local SGD. Given a training dataset and a data augmentation function, Algorithms [ and 2 show the imple- mentations of distributed samplers for sampling local batches with and without replacement. Then Algorithms 3 to 5 show the implementations of parallel SGD, Local SGD and Post-local SGD that can run with either of the samplers.
Sampling with replacement. Our theory analyzes parallel SGD, Local SGD and Post-local SGD when local batches are sampled with replacement (Algorithm[]) T That is, local batches consist of IID samples from the same training distribution D where DD serves as an abstraction of the distribution of an augmented sample drawn from the training dataset. The mathematical formulations are given in Section
Sampling without replacement. Slightly different from our theory, we use the sampling without replacement (Algorithm 2 in our experiments unless otherwise stated. This sampling scheme standard in practice: is used by GyyGo et al. (2017) for parallel SGD and Lin et al. (2020b): ctt al. (2021) for Post-local/Local SGD. This sampling scheme works as follows. At the beginning of every epoch, the whole training dataset is shuffled and evenly partitioned into K shards. Each worker takes one shard and samples batches without replacement. When all workers pass their shard, the next epoch begins and the whole dataset is reshuffled. An alternative view is that the workers always share the same dataset. For each epoch, they perform local steps by sampling batches of data without replacement until the dataset contains too few data form batch. Then another epoch starts with the dataset reloaded to the initial state.
Discrepancy in Sampling Schemes. We argue that this discrepancy between theory and experi- ments on sample schemes is minor. Though sampling without replacement is standard in practice most previous works, e.g., WWang Joshi (2019); Li al. (2021a); la;; Zhang al. (2020), analyze
sampling with replacement for technical simplicity and yields meaningful results.
Moreover, even if change the sampling scheme with replacement, Local SGD can still improve the generalization of SGD (by merely adding local steps). See Appendix (erendixFJoocc for the experiments. We believe that the reasons for better generalization of Local SGD with either sampling scheme are similar and leave the analysis for sampling without replacement for future work.
20

---

Published as a conference paper at ICLR 2023
Algorithm 1: Distributed Sampler on K Workers (Sampling with Replacement)
Require: shared training dataset D, data augmentation function A(é) Hyperparameters: local batch size Bioe
Function Sample () worker k:
Draw Bioe IID samples EB,..., from D with replacement: & - A(Gx) for all 1 < b I Bioc ;
apply data augmentation
return (fi...., BBr
end
Algorithm 2: Distributed Sampler on K Workers (Sampling without Replacement)
Require shared training dataset D, data augmentation function A(E)
number of local batches per worker per epoch Hyperparameters: local batch size Boo
Constant: Nioc Bioe
Local Variables: (k)) + Blee for worker // number of samples drawn in this epoch
Function Sample() () worker k: if (kk) Nior then // start
epoch:
// synchronize Draw a random permutation P of 1.... [D] jointly with other workers so that the same
Wait until all the other workers reach this line;
permutation shared among all workers
// reshuffle the dataset partition the dataset
Q (8) t P(k Poo Bj (k) t 0:
for all 1 < S Nioc
10
11
end
12
for i l.... Bloc do
the (ll)) Q0 -th data point of D:
13
sample without replacement apply augmentation
AlEx);
15
end
(ik) (k) Bioc
16
17
return (E1,..., . B
18
end
21

---

Published as a conference paper at ICLR 2023
Algorithm 3: Parallel SGD on K Workers
Input: loss function ((0; E), initial parameter A Hyperparameters: total number of iterations T, learning rate 7,, local batch size Bloc
for t 0, O.... . - do for each worker k do in parallel
Ekk,t. +- Sample() ()
1/ sample local batch
Bie 7
ghht
Bloe
// computinng the local gradient all Reduce aggregation gradients // update the model
end
1 Σx
9:
K k= gb.t
0.+1 t θ: 71:9
10 end
Algorithm 4: Local SGD on K Workers
Input: loss function ((0: {), initial parameter (()
Hyperparameters: total number of rounds R, number of local steps H per round Hyperparameters: learning 7, local batch size Bioc
for 8 0,.... R 1 do
for each worker k do in parallel
[ - Ø0);
// mointain lccal copy of the global iterate
for t = O..... H 1 do
(s) (5)
(5.,111'... Sk.t. Bioe Sample ()
// sample local batca
9h.t Bisc
Σ Li=1
Bise Ve(e <(...):
1f computing the local gradient // update the local model
((*)
((*)
(6)
10
k,++1
11
end
12
end
(
K Ek=1
(5)
13
17 all Redace aggregation of local iterates
14 end
Algorithm 5: Post-local SGD on Workers
Input: loss function ((0: {), initial parameter O, Hyperparameters: total number of iterations T, learning rate local batch size Bioc Hyperparameters: switching time point fo, number of local steps H per round Ensure: T to is multiple of !t
Starting from Qg. run Parallel SGD for to iterations and obtain 0,,-
Starting from Og run Local SGD for 1 return the final global iterate of Local SGD;
H (T - (b) rounds with H local steps per round;
22

---

Published as a conference paper at ICLR 2023
D MODELING LOCAL SGD WITH MULTIPLE CONVENTIONAL SDES
Lin et al. 2020b) tried to informally explain the success of Local SGD by adopting the argument that larger diffusion term in the conventional SDE leads to better generalization (see Section [.1 and appendix A Basically, they attempted to write multiple SDEs, each of which describes the
H-step local training process of each worker in each round (from o to ( The key difference between each of these SDEs and the SDE for SGD (3) is that the former one has a larger diffusion term because the workers use batch size Bioc instead of B:
/ " (X)(W)
dx V V B
(10)
Lin et al. (2020b) then argue that the total amount of "noise" in the training dynamics of Local SGD larger than that of SGD. However, it is hard to see whether is indeed larger, since the model averaging step at the end of each round can reduce the variance in training and may cancel the effect of having larger diffusion terms.
More formally, complete modeling of Local SGD following this idea should view the sequence of global iterates {0(*} as a Markov process (*). , Let Px (z, B, t) the distribution of X(t) in (3)
with initial condition X (0) æ. Then the Markov transition should be X (*+1) 1/4 Σズ= X (5)
k.H"
where X ( from TTP..
X K are K" independent samples from Px (X Bioe; Hn). i.e., sampling
Consider one round of model averaging. It is true that (x Bloe: Hn) may have a larger vari- ance than the corresponding SGD baseline Px (X(+), (5) B, Hn) because the former one has a smaller batch size. However, it is unclear whether X (*+1) also has a larger variance than Px X (e) B, HT)- This is because X (++1) is the average of K samples, which means we have to compare K times the variance of Px [X (s), Broe; Hn) with the variance of Px (x(s), B, Hn). Then it is unclear which one is larger.
In the special case where Hn small, Px (x))), Blse, Hn) is approximately equal to the following Gaussian distribution:
H Bioc
N
x -
E(X(1))
(11)
Then averaging over K samples gives
H B -(x()))
NV ( - -)V(((())
(12)
which is exactly the same as the Gaussian approximation of the SGD baseline. This means there do exist certain cases where Lin et al. (2020b)* 's argument does not give a good separation between Local SGD and SGD.
Moreover, we do not gain any further insights from this modeling since it is hard to see how model averaging interacts with the SDEs.
E
ADDITIONAL INTERPRETATION OF THE SLOW SDEs
E.1
UNDERSTANDING THE DIFFUSION TERM THE SLOW SDE
So far, we have discussed why adding local steps enlarges the drift term in the Slow SDE and why enlarging the drift term can benefit generalization. Besides this, here we remark that another way to accelerate the corresponding semi- remi-gradient method for minimizing the implicit regularizer is to reduce the diffusion term, so that the trajectory more closely follows the drift term. More formally, we propose the following:
Hypothesis E.1. Starting at minimizer Ço T, run (K1, -SSoo SDE and (MI, -Sow SDE respectively for the same amount of time T O and obtain C(T), (T)) If Σn 0 and Ky < K then the expected test accuracy at C(T)) is better than that at (T)
23

---

Published as a conference paper at ICLR 2023
0.95
Top-1 1 Test Accuracy
0.75
Test Accuracy
0.94
0.70
K 250
K
0.93
0.65
0.92
61
100k
150k
200k
15000
20000
25000
30000 35000
Iterations
Iterations
(a) CIFAR-10, 1/ 600 for K K> 1.
(b) ImageNet, HI = 78 for K > 1.
Figure 3: Reducing the diffusion term of the Slow SDE for Local SGD leads to better generalization. Test accuracy improves we increase K with fixed 17 and H to reduce the diffusion term while keeping the drift term untouched. See Appendix M.4 for details.
Here we exclude the case of Σ = because in this case the diffusion term in the Slow SDE is always zero. To verify Hypothesis E. we set the product =- nyHl large, keep H, y fixed, increase the number of workers K, and compare the generalization performances after a fixed amount of training steps (but after different numbers of epochs). This case corresponds to the K 2 Slow SDE, so adding more workers should reduce the diffusion term, As shown in Figure 3 a higher test accuracy is indeed achieved for larger K.
Enlarging the learning rate is not equally effective as adding local steps. Given that Local SGD improves generalization by strengthening the drift term, it is natural to wonder enlarging the learning rate of SGD would also lead to similar improvements. While it is true that enlarging the learning rate effectively increases the drift term, it also increases the diffusion term simultaneously, which can hinder the implicit regularization by Hypothesis [.1] In contrast, adding local steps does not change the diffusion term. As shown in Figure(6(a) even when the learning rate of SGD is increased, SGD still underperforms Local SGD by about 2% in test accuracy.
Implication:
On the other hand, in the special case of where Σ₁ = 0. Hypothesis E. does not hold, and enlarging
the learning rate by VK results in the same Slow SDE as adding local steps (see Appendix G for derivation). Then these two actions should produce the same generalization improvement, unless the learning rate so large that Slow SDE loses track of the training dynamics. As an example of such a special case, an experiment with label noise regularization is presented in Figure g
E.2 THE EFFECT OF GLOBAL BATCH SIZE ON GENERALIZATION
In this section, we discuss the effect of global batch size on the generalization of Local SGD. Given that the computation power of single worker is limited, we consider the case where the local batch size Bio is fixed and the global batch size B = Bi is tuned by adding or removing the workers. This scenario is relevant the practice because possible to train the neural net without causing generalization degradation.
may want to know the maximum parallelism
For SGD, previous works have proposed the Linear Scaling Rule (LSR) (Krizhevsky) 2014 Goyal et al. 2017 Jastrzębski et al. 2017): scaling the learning rate 1) +> Knj linearly with the global batch size B RB yields the same conventional SDE under constant epoch budget, hence leading to almost the same generalization performance as long as the SDE approximation does not fail.
We show in Theorem H that the LSR does not change the Slow SDE of SGD either. Experiments in Figure 4 show that the LSR indeed holds nicely when we continue training with small learning rates from the same CIFAR-10 and ImageNet checkpoints in Figure Here we choose K 16 and K = 256 as the base settings for CIFAR- 10 and ImageNet, respectively, and then tune the learning rate maximize the test accuracy. As shown in Figures 4(a) and 4[b) the optimal learning rate turns out to be small enough that the LSR can be applied to scale the global batch size with only
a minor change in test accuracy.
Now, assuming the learning rate is scaled as LSR, we study how to tune the number of local steps H for Local SGD for better generalization. natural choice to tune in the base settings and keep unchanged via scaling H + H/K. Then the following SDE can be derived (see Theorem H.2
dc(1)= Pc
JB Σ / (S)dW, 2B V ((C)
= V ((S) [(S)
(13)
2B
(a) diffuusio (unchanged)
(b) drift-I (unchanged)
(4) drin-I1 (rescaled)
24

---

Published as a conference paper at ICLR 2023
0.93
0.760
Test Accuracy
Text Accuracy
0.92
optimal
0.758
0.756
0.91
optimal
constand x
0.754
32
64
16
32
64
128
Number workers K' (log scale)
Number workers K (log scale)
(a) CIFAR-10, start from #250.
(b) ImageNet, start from #100.
2560
256
optimal constant
optimal constans
a (log scale)
640
a (log scale)
64-
160
16-
40-
10
16
128
512
Number workers K' (log scale)
Number workers K (log scale)
(c) CIFAR-I0, 10, start from #250.
(d) ImageNet, start from #100.
Figure 4: For training from CIFAR-IO and ImageNet checkpoints, Local SGD consistently outper- forms SGD (H = 1) across different batch sizes B (fixing Bio and varying K), where the learning rate is scaled by the LSR v x B. Two possible ways of tuning the number of local steps H are con- sidered: (1). Tune H for the best test accuracy for K 16 and K 256 respectively CIFAR-10 and ImageNet, then scale H as H ix 1/B so that a == H is constant; (2). Tune H specifically for each K. See Appendix M.5 for training details.
Compared with 4 the drift-Il term here is rescaled by positive factor. Again, when is large, can follow the argument in Section 3.3.2 approximate (C) (C) and obtain the following ( &K )-Slo SDE:
dc(() = Pc
2 VB 11 (S)
Rk V E(C))E 2B
(14)
The drift term of the above SDE is always stronger than SGD 7 long there exists more than one worker after the scaling (i.e., RK > 1). As expected from Hypothesis 3.1] we observed in the experiments that the generalization performance of Local SGD is always better than or least comparable to SGD across different batch sizes (see Figures 4(a) and 4(b)).
Taking closer look into the drift term in the Slow SDE (4). we can find that it scales linearly with k. According to Hypothesis 3.I the SDE is expected to generalize better when adding more workers (x 1) and to generalize worse when removing some workers (N < 1). For the latter case, we indeed observed that the test accuracy of Local SGD drops when removing workers. For the case of adding workers, however, we also need to take into account that the LSR specifies larger learning rate and causes a larger SDE approximation error for the same a, which may cancel the generalization improvement brought by strengthening the drift term. In the experiments, we observed that the test accuracy does not rise when adding more workers to the base settings.
would be beneficial to decrease to better trade-off between regularization strength and approximation quality. Figures 4(c) and 4(d). we plot the optimal value of a for each batch size, and we indeed observed that the optimal drops
Since a also controls the regularization strength (Section 3.3.3 for large batch size
scale up K. Conversely, smaller batch size (and hence smaller learning rate) allows for using a larger a to enhance regularization while still keeping a low approximation error (Theorem 3.3). The test accuracy curves in Figures 4(a) and 4(b) indeed show that setting a larger a can compensate for the accuracy drop when reducing the batch size.
F ADDITIONAL EXPERIMENTAL RESULTS
In this section, we present additional experimental results further verify finding.
25

---

Published as a conference paper at ICLR 2023
Supplementary Plot: Training time should be long enough. Figures ((a) and 5(b) show
larged views for Figures(20 and [(c) respectively, showing that Local SGD can generalize worse than SGD in the first few epochs.
shows that reducing the learn- ing rate from 0.32 to 0.064 does not lead to test accuracy drop for Local SGD CIFAR-10, if the training time is allowed to longer and the number of local steps is set properly. Figure (
Supplementary Plot: Learning rate should be small. Figure ((C)
presents the case where, with large learning rate, the generalization improvement of Local SGD disappears even starting from a pre-trained model.
Supplementary Plot: Reconciling
main finding with Ortiz al. (002)) In Figure (
the generalization benefit of Local SGD with H = 24 becomes less significant after the learning rate decay at epoch 226, which consistent with the observation by Ortiz et al. (22TT) that the generalization benefit of Local SGD usually disappears after the learning rate decay. But we can preserve the improvement by increasing to 900. Here, we use Local SGD with momentum.
Supplementary Plot: Optimal gets larger for smaller 1). In Figure 5(t) we summarize the optimal a = 7H that enables the highest test accuracy for each learning rate in Figure 2(f) We can see that the optimal increases as we decrease the learning rate. The reason is that the approxima- tion error bound O( an log f in Theorem 3.3 decreases with v), allowing for a larger value of better regularize the model.
Test Accusccy
0.8
Top-1 Text Accurccy
0.70
Text Aucuracy
0.95
0.8
180
0.83
0.73
0.27
N
0.06 H 3000
0.2g
300
I03
0.02
250
0c0C0
60003
Epochs
Epochs
Epochs
(a) CIFAR 10, start from random.
(b) ImageNet. start from #250.
(c) CIFAR-10, start from #100.
Opnimal a (log scale)
Tet Acaroor
0.72
Tex Accuracy
0.70
0.93
1.5
580
00.88
0.92
200
300
600
0.8
0.054
0.30
Number Local Staps # (lng seab))
Epochs
Leamnng Rate (Ing scale)
(d) ImageNet, start from #100.
(e) CIFAR-10, start $150.
(f) ImageNet, optimal a v.s. 77.
Figure 5: Additional experimental results about the effect of the learning rate, training time and the number local steps. See Appendix M.2 for details.
0.91
0.99
Tex Accuracy
Test Acourtty
Text Accuracy
0.00
0.322
1. Fare 2253
0.91
0.30
E 680, 0.6 Fir 4053
0.85
0.90-
0.39
0.64
1.78
2.56
16
0.37
0.64
1.28
40000
6000
Rata stale)
Losening Rate y stale)
(a) SGD with various 77.
(b) SGD with larger batch sizes.
(c) Post-local SGD, sampling with replacement.
Figure 6: Additional experimental results on CIFAR-10. 10. See Appendix M.3 for details.
SGD generalizes worse with extensively tuned learning rates. In Figure 6(a) run SGD from both random initialization and the pre-trained model for another 3. 0000 epochs with various learning rates and report the test accuracy. We can that none of the SGD runs beat Local SGD with the fixed learning rate = 0.32. Therefore, the inferior performance SGD in Figures 2(a)
and 2(b) is not due the improper learning rate and Local SGD indeed generalizes better.
26

---

Published as a conference paper at ICLR 2023
enlarge the batch size of SGD and report the test accuracy for various learning rates. We can see that SGD with larger batch sizes performs
SGD with larger batch sizes performs better. In Figure 6(b)
better and none of the SGD runs outperform Local SGD with the fixed learning rate yj 0.32. This result is unsurprising since it is well established in the literature (Jastrzebski al. 2017: Smith et al. 2020 Keskar al., 2017) that larger batch size typically leads worse generalization. See Appendix A for a survey of empirical and theoretical works on understanding and resolving this phenomenon.
Sampling with or without replacement does not matter. Note that there is a slight discrep- ancy sampling schemes between our theoretical and experimental setup: the update rules (I) and ( assume that data are sampled with replacement while most experiments use sampling without replacement (Appendix Cp. To eliminate the effect of this discrepancy, we conduct additional exper- iments on Post-local SGD using sampling with replacement (see Figure [6()) and Post-local SGD significantly outperforms SGD.
G DISCUSSIONS ON LOCAL SGD WITH LABEL NOISE REGULARIZATION
G.1 THE SLOW SDE FOR LoCaL SGD WITH LABEL. NOISE REGULARIZATION
In this subsection, present the Slow SDE for Local SGD in the case of label noise regularization and show that Local SGD indeed induces a stronger regularization term, which presumably leads to better generalization.
Theorem G.1 (Slow SDE for Local SGD with label noise regularization), For C-class classifi- cation task with cross-entropy loss, the slow SDE of Local SGD with label noise has the following form:
+ 1 Vr ( 11) ((F(2HNV*C((M) dt, ac(t) = 4B
(15)
2Hn
where F=x)== So (y)dd and interpreted matrix function, Additionally, Vr stands for the gradient of fuuccii I projected to the tangent space Tf
Proof. See Appendix L
Note that the magnitude of the RHS in (15) becomes larger H increases. By letting H go infinity, we further have the following theorem.
Theorem G.2. As the number of local steps goes to infinity, the slow SDE of Local SGD with label noise can be simplified as:
dc(t) = 4B ((()(V
(16)
Proof. We obtain the corollary by simply taking limit. By L L'Hospital's rule,
F(ar)
dF(az)
lim
lim
lim av(ar)
#000
+400
dx
++00
Therefore,
Ir (F((2H)))((()) - -tr((
lim
(17)
+00000
2Hn
Substituting (17) into 15) yields
16
As introduced in Section 3.3 the Slow SDE for SGD with label noise regularization has the follow-
ing form:
dc(t)
4B C())
(18)
which is deterministic flow that keeps reducing the trace of Hessian.
27

---

Published as a conference paper at ICLR 2023
Accuracy
0,85
Text Accuracy
05,004
0.80
0.75
/ www /
6,344
0.74
0.75
0.006
1:.2
0.4
11.
0.006
0.2
0.4
11.
0.8
Million
Million Iterations
(a) ResNet-56 GroupNorm.
(b) VGG-16 w/o normalization.
Figure 7: Local SGD with label noise regularization CIFAR-IO without data augmentation using K = 32 .Bo = 128. A larger number of local steps indeed enables higher test accuracy. For both architectures, replace ReLU with Swish. See Appendix MM6 for training details.
As the trace of Hessian can be seen as a measure for the sharpness the local loss landscape, (18) indicates that SGD with label noise regularization has implicit bias toward flatter minima,
which presumably promotes generalization (Hochreiter & Schmidhuber 1997 Keskar et al. 2017]
Neyshabur al. 2017). More concretely, [lance et al. (2020) and Li et al. (2021b) connect mini- mizing the trace of Hessian to finding sparse or low-rank solutions for training two-layer linear nets.
Damian et al. 202T empirically showed that good generalization correlates with smaller trace
of Hessian in training ResNets with label noise. Besides, Ma & Ying (202I) connect the trace of Hessian the smoothness of the function represented by deep neural net.
From Theorems G.I and G.2 we can conclude that Local SGD accelerates the process of sharp- ness reduction, thereby leading to better generalization. Furthermore, the regularization effect gets stronger for larger H and is approximately K times that of SGD. We also conduct experiments on non-au non-augmented CIFAR-10 with label noise regularization verify our conclusion. As shown in Figure increasing the number of local steps indeed gives better generalization performance.
G.2 THE EQUIVALENCE OF ENLARGING THE LEARNING RATE AND ADDING LOCAL STEPS
In this subsection, we explain in detail why training with label noise regularization is special case where enlarging the learning rate of SGD can bring the same generalization benefit as adding local steps. TWhen we scale up the learning rate SGD 7j + кл (while keeping other hyperparameters unchanged), the corresponding Slow SDE is (8) with time horizon x²T instead of where SGD tracks a continuous interval of 1=²1? per step instead of 17². After rescaling the time horizon to so that SGD tracks continuous interval of 1/ per step, obtain
A² 4B Vr
dc(1)
(19)
Let VK in 19 and we obtain the same Slow SDE as (6), which is for Local SGD with large number local steps. In Figure we conduct experiments to verify that SGD indeed achieves comparable test accuracy that of Local SGD with a large H if its learning rate scaled up by / VK that of Local SGD.
H DERIVING THE SLOW SDE AFTER APPLYING THE LSR
In this section, we derive the Slow SDEs for SGD and Local SGD after applying the LSR in Ap- pendix E.2 The results are formally summarized the following theorems.
Theorem H.1 (Slow SDE for SGD after applying the LSR). Let Assumptions 3.1 to 3.3 hold.
Assume that run SGD with learning rate
and the number of workers K' KK for some constant > O. Let T > 0 be a constant and ((t) be the solution to with the initial condition ((0) 1 (0,) I. Then for any -mmoohh function g()),
[Elg( - /x)]] (1700.25), where O(.) hides log factors and
maxos.5 =n 7 stants that independent of ní but can depend on g()).
28

---

Published as a conference paper at ICLR 2023
Test Accuracy
0.7400
0.7375
0.7350
0:0006
0.5
1.5
Million Iterations
Figure 8: Local SGD with label noise regularization CIFAR-IO without data augmentation using
K 4, Bise = 128. SGD (H = 1) indeed achieves comparable test accuracy as Local SGD with
large H when we scale up its learning rate to v K times that of Local SGD. See Appendix M.6 for training details.
Proof. Replacing B with xB in the original Slow SDE for Local SGD 7 gives the following Slow SDE:
dc(t) Ps
VKB D
(C)dW:
V ((S)E
(20)
(a)
(b) drif--1
Note that the continuous time horizon for (20) is KI' instead of since after applying the LSR, SGD tracks continuous interval of x²21²² per step instead of 17² while the total number of steps is scaled down by K. We can then rescale the time scaling to obtain 7 that holds for T.
Theorem H.2 (Slow SDE for Local SGD after applying the LSR). Let Assumptions 3.7 to 3.3 hold.
Assume that the number of local steps H' a for some constants o,, K > O. Let > O be a constant and
Local SGD) with learning rate 1/' ka, the number workers K" RKK and
xn ((t) be the solution to (1) with the initial condition {(0) = ( E Then for any C3 -mooohhh function g(0), maxuses niūz [ - (lg(((s /x>]] = (1100.25), where Ó(-) hides log factors and constants that are independent of 1/ but can depend on g(0).
dc(t) = P
1 ()) /B
2B V* (S))
=k ((II>(((S)) (C)
(21)
2B
(a) diffission (unchanged)
(b) drift-I (unchunged)
(c) drift-ll (resculed)
Proof. Replacing B with BB in the original Slow SDE for Local SGD gives the following Slow
SDE:
dc(t)
Pc
E 1 VRB
(C)dW, I V ((C)
KK-1 v (
(22)
2
2xB
(a) diffusion
(b) drift-1
(c) drift-1l
Note that the continuous time horizon for (22 is KT instead of T since after applying the LSR,
Local SGD tracks continuous interval of nf per step instead of 1/2 while the total number of
steps is scaled down by k. We can then rescale the time scaling to obtain (2I) that holds for T. O
29

---

Published as a conference paper at ICLR 2023
1 PROOF OF THEOREM 3.1
This section presents the proof for Theorem 3.1 First, we introduce some notations that will be used
throughout this section. For the sequence of Local SGD iterates {0 k [K], = 1 H, 2 0), we introduce an auxiliary sequence (ü, her, which consists of GD iterates from O(0)
üu 0 U+t1 < dt V[(()
For convenience, let ú, = ú.H+t and H = (s) We will use ú,"' and ú. 2g, and + interchangeably. Recall that have assumed that is C3-smooth with bounded second and third order derivatives. Let M2 supoER [V and Va supoEr* [ Since Ve(0; () is bounded, the gradient noise = is also bounded. We denote by max an upper bound such that | S max holds for all s, k,
k
To prove Theorem 3.1 will show that both Local SGD iterates ((*) and SGD iterates w.H track GD iterates ü,H closely with high probability. For each client k, define the following sequence {2n, :t 0}, which will be used in the proof for bounding the overall effect of noise.
ZR,-= = Σ II (I
(I-)V<((())
Zk,7,
Zk.0 0,
vk [K].
W1++1
The following lemma shows that Zu is concentrated around the origin.
Lemma I.1 (Concentration property of of( (Z With probability at least - 8, the following holds
simultaneously for all k. E [K], 0 < 1 < 1
2T
2TK
Zk, Cj a max
log
T]
81j
where Õ == ex2(
Proof. For each ZR.t, construct sequence { Zu..
[-1
Σ
II
(I -I- ((-))(((()))
(5)
7(0) = 0.
k,000
T=0
11T+1
Since [ ll2 S U/ for all 2 0, the following holds for all ST <t and < 1:1:
1- II (I (
S1 +1+ p211)' S exp(TTv) = ET.
So (āx,, Mi is a martingale with 12 - 2 S C Since Zr,t Z by Azuma-Hon nouma-Hoeffding inequality,
22
P Zm,ele 26) < 2xxp
21
Ö
1 max
Taking union bound on all & E [K] and S S
[ we can conclude that with probability least
1 5,
2TK log
T
12x.el22 5 Ci C1axx
vO < / <
k E [K]
8n
77
30

---

Published as a conference paper at ICLR 2023
The following lemma states that, with high probability, Local SGD iterates 0 and 0(s) closely track the gradient descent iterates ü. for H rounds.
Lemma 1.2. For d = O(poly(i)), the following inequalities hold with probability at least
I
T
100 + S Csy/ n fog n76"
vk 5K]
8<<
.5sts., H,
Hn
and
00 u.H1l2 Ezy/n nogn log np""
T
Hn
where C3 is constant independent of and H.
Proof. Let b,2 and GD iterates. According to the update rule for
dt and A (s) = ((s)
in be the differences between the Local SGD lett and a!!)
k,t
vVC(())) - -1 ú d út -)(
(23)
k,++1 R,t
(24)
Subtracting 23 by 24 gives
= kst (VL((())) - vL(a - 17
k
= -I -V-
A (e)
k.t n + t/0k.11
(25)
where 05 is a remainder term with norm ||
2 À (1) 12. For the s-th round of Local SGD,
S
can apply 25 & times obtain the following:
サ-1 II - (I---(Z(a())))
( k.t
(
(I (
K.0
1
2k.T
(26)
Σ (1- (-
7=01=7+
Here, T can expressed the following form:
#+
T
Zz. sHt
II 1=sH
(I (I-
Zr.aH.
Substituting in / H and taking the average, we derive the following recursion:
A (++1)
Σ A
(s)
K
KE]K]
II -I- (V---((((()))) ( -==
(++1)H-1
K Σ ZE((+1)11
11
Σ
II
(I (I- (I-nV=((ai)
Zeo0
K
kE[K]
kE(K]
I=eH
H-1
(1 I
(
(27)
KE[K]
1=7+1
31

---

Published as a conference paper at ICLR 2023
Applying (27
times yields
d-1 8-18-1 MΣ
sH
(8)
7
Σ
ZR.SHl
II (I- 11-
B (28)
kE[K]
r=u T=U kE[K]
Jmri+t-1
Substitute 28
into 26
and we have
++H+t11
(8)
7
ZWn - 12 +1
II 1=sH
(I (I-nv>
Zr,CH
2
K
[EE
H-1
II
(I (I- (I-NV2
K
P=0
7=0
[E[[]]
U=rH+T+1
+
II
(I-n
+=sH+++11
By Cauwhy-Schwartz inequality and triangle inequality, we have
A
7 <
Σ
ll
+ +
Ze,
+ -nOil Zk
K
[E
(29)
nCivs 1/3
H-1
nCivs
r-1
ΣΣ Σ r=0
(r)
+
2K
2
R
r=
ME(K]
r=0
where Õ = exp(v2T))
Below prove by induction that for 6 = Oppoly(n)), if
2T
2TK
1 ll2 S C1 01 max
log
VO 20<6<
-
k E [K],
(30)
np
then there exists a constant O2 such that for all E [K], S 8 <
n
and S L15 H.
A (0) la S E2 1oo
2TK
(31)
n70
First, for all & E [K], | (0)
k 0 and hence 31
holds. Assuming that 31 holds for all Å k(.r
where k' E [K] 0 S
s, 0 S H and
s. 0.
then by 29 for all & [K], the
following holds:
2TK
2TK
À k.lll | < 3 o max 2T log
+ é CTT log
178
Let O₂ 2 60fommmx 2T. Then for sufficiently small n, probability at least 8. Furthermore, notice that (()) ūsH have the lemma.
31
holds. By Lemma
1.1
30
holds with
K ELe[K]
Hence
R.1f
The iterates of standard SGD can be viewed as the local iterates single client with the number of local steps 1 Therefore, we can directly apply Lemma [.2 and obtain the following lemma about the SGD iterates wt
Corollary I.1. For 8 (poly(n)), the following holds with probability at least 8:
I v T
|w.. - úH llz - C' nog no¹
Hn
where C3 is the same constant as in Lemma 1.2
Applying Lemma 1.2 and Corollary L1 and taking the union bound, we have Theorem 3.1
32

---

Published as a conference paper at ICLR 2023
J PROOF OUTLINE OF MAIN THEOREMS
We adopt the general framework proposed by Li et al. (2019a) to bound the closeness of discrete algorithms and SDE solutions via the method of moments. However, their framework is not di- rectly applicable to our case since they provide approximation guarantees for discrete algorithms with learning rate 17 for Onn=²) steps while want to capture Local SGD for O(n ?) steps. To overcome this difficulty, we treat Rgm == Laype rounds as a "giant step" of Local SGD with an "effective" learning rate ny'_ô, where 5 is a constant in (0, 1), and derive the recursive formulas to compute the moments for the change in every step, every round, and every Rgrp rounds. The formulation of the recursions requires detailed analysis of the limiting dynamics of the iterate and careful control of approximation errors.
The dynamics of the iterate can divided into two phases: the approaching phase (Phase 1) and the drift phase (Phase 2). The approaching phase only lasts for O 5) rounds, during which the iterate is quickly driven to the minimizer manifold by the negative gradient and ends up within only ((VT) from I' (see Appendix K.5). After that, the iterate enters the drifting phase and moves in the tangent space of I' while staying close to I' (see Appendix K. The closeness of the iterates (local and global) and r is summarized in the following theorem.
Theorem J.1 (Closeness of the iterates and I)., For ó = O(poly(n)), with probability at least 1 1- 6, for all Olog (t) S (T/(H))))
1 ((0()) = -0
(((() )r,
r78
Also. for all Oloo ;) = (T | 1. & [K] and 0 S st H.
1 10 k.t - . ((())) =0 (Vnlos
=0
n1
Here, (-) hides constants independent of and 5.
To control the approximation errors, we also provide high probability bound for the change of the manifold projection within Reep rounds.
Theorem J.2 (High probability bound for the change of manifold projection). For 8 = Oppoy((n))
with probability at least 1 - 5, for all 0 S s S [H)?) Rgep and 0 S S Rgrp
| - कه (0()))) = (106-0.58 log n5
((0()), ( (ol eT,
where O()) hides constants independent of n and 6.
The proof of Theorems and J.2 is based on the analysis of the dynamics of the iterate and presented in Appendix K.7
Utilizing Theorems J.I and J.2 we move on to estimate the first and second moments of the change of the manifold projection every Rgep rounds. However, the randomness during training might drive the iterate far from the manifold (with a low probability, though), making the dynamics intractable.
To tackle this issue, we construct a well-behaved auxiliary sequence (0(.), (*) J. which is constrained
to the neighborhood of I and equals the original sequence (O(.)) with high probability (see Defini- tion K.5). Then we can formulate recursions for the change of manifold projection of the auxiliary sequence using the nice properties near T. The estimate of moments is summarized in Theorem K.2
Finally, based on the moment estimates, we apply the framework in Li et al. (2019a) to show that the manifold projection and the SDE solution are weak approximations of each other in Appendix K.10
K
PROOF DETAILS OF MAIN THEOREMS
The detailed proof is organized follows. In Appendix K. we introduce the notations that will be used throughout the proof. To establish preliminary knowledge, Appendix (.22dixxXXIpIroiiess explicit ex- pression for the projection operator ((-), and Appendix K. presents lemmas about gradient descent
33

---

Published as a conference paper at ICLR 2023
(GD) and gradient flow (GF). Based on the preliminary knowledge, construct nested working zone to characterize the closeness of the iterate and I' in Appendix K.4, Appendices K.5 to K.10 I..
make up the main body of the proof. Specifically, Appendices K.5 and K.6 analyze the dynamics of
Local SGD iterates for phases and 2, respectively. Utilizing these analyses, we provide the proof of Theorems J.1 and J.2 in Appendix [K7 and the proof of Theorem [33 in Appendix [K_]] Then
derive the estimation for the first and second moments of one "giant step (0(0)*+Rgm) Rgv) ((0()))
in Appendix K.9 Finally, we prove the approximation theorem 3.2 in Appendix K.10
K.1 ADDITIONAL NOTATIONS
Let Biot T be the total number of rounds. Denote by (b() the manifold projection of the global iterate at the beginning of round s. Let (=(*)) k, := 0 k,t (+)) be the difference between the local iterate and the manifold projection the global iterate. Also define =/4² == ExE(K] (6) [] R. and (x) = K [kE[K] k,0 which is the average of a among K workers at step 0 and H. Then for all &s [K]. 2,0
(x)
(a) x ((+) (b(), Finally, Since Vε(O:ç) is bounded, the gradient noise
(4) Ri is also bounded and we denote by max the upper bound such that ||= F Vs, k.
We first introduce the notion of μ-PL. We will later show that there exists neighborhood of the minimizer manifold I' where satisfies μ--LL.
Definition K.1 (Polyak-Lojasieviee Condition). For μ > O, we say function L()) satisfies p- Polyak-Lojasiewicc condition (abbreviated as p-PL) on set U if
|VCO ((C(0) inf L()")).
We then introduce the definitions of the e-ball at a point and the e-neighborhood of a set. For 0 E Rd and > O. (0) = (0' - } the open e-ball centered O. For C Rd 2e = Uøez (0) is the e-neighborhood of Z.
K.2 COMPUTING THE DERIVATIVES OF THE LIMITING MAPPING
In subsection, we present lemmas that relate the derivatives the limiting mapping ((-) to the derivatives of the loss function L(-). We first introduce the operator VH.
Definition K.2. For semi-definite symmetric matrix H E Rdxdd let Xj, Vú be the j-th eigenvalue and eigenvector and vo 's form an orthonormal basis of Rdd. Then, define the operator Vh R Rd as
VH(M) =
Σ
du +1 (M, (M, v.0, VM Rddd
1
LOVA
Intuitively, this operator projects M the base matrix faae and sums up the projections with weights X
Additionally, for 0 E I, denote by To and Tel
the tangent and normal space of at respectively. (202lb). We include them to make the paper self-contained.
Lemmas K. to K.4 are from Li et al.
Lemma K.1 (Lemma C.1 of Li et al. V2L(0)v
(0021 For any O E I and any v (T)),, it holds that
Lemma K.2 (Lemma 4.3 of Li et al.
2021b ) For any 0 r, aD(0) Bdxd is the projection
matrix onto the langeni space (T)))
Lemma K.3 (Lemma C.4 of Li et al. that
2021b ). For any 0 E u E Rd and TT(T), holds
i u] = -0
v²C(0) uu] v2 C() v (100__000000
Lemma K.4 (Lemma C.6 of Li (021b)). For any 0 e and Σ span u E T (T)),
(0) =0000000
34

---

Published as a conference paper at ICLR 2023
Lemma K.5. For all 0 r, u (T),, it holds that
(
O.
(32)
Proof. This proof is inspired by Lemma C.4 of Li et al. (202lb) For any 0 E T, consider a parameterized smooth curve v(t,, 2 0 on I such that v(O) 0 and v' v'( Let P (((
0 ()()( Pi (t) I (v((()) and (t) [()((()) By Lemma C.I and 4.3 in Li et al.
2021b).
Htt)= = P (t)
Take the derivative with respect to t on both sides,
(1) PL (L( + (1()H()) (1 H (1) (L) = (1(())
At t = 0, we have
P1(0) (0) (()
(33)
WLOG let (())
diag( A1
Ad): E Raxd where = for all m 1 S d. Therefore
Im
0
0
(1(
P (0)
Decompose P(()), (()) and H' (0)as follows.
0
0
0
Ia
d-m
PP'
I (0)
P
(0)
()) 0
H
H
H (0)
(0)
(0))
(0)
H (0)
P
1121
(0)
[0)
0
H (0)
H
Substituting the decomposition into
33 we have
PP 11 (0) Hn (0) P (0) Hn (0) 0]
0
H (0)
H
Therefore, (0) = 0 and
0
(()) (P((O))
H (0)
Any u E (T) can be decomposed as u - [0.u2] where u2 E Bd-m. With this decomposition, we have P (0)H (0) O. Also, note that H'( [0) Hence,
[1(0)V=((0)(or"] 0.
K.3 PRELIMINARY LEMMAS FOR GD AND GF
In this subsection, we introduce a few useful preliminary lemmas about gradient descent and gradient flow. Before presenting the lemmas, we introduce some notations and assumptions that will be used in this subsection.
Assume that the loss function L(0) is p-smooth and p-PL in an open, convex neighborhood U of a local minimizer 0- Denote by L := L((*) the minimum value for simplicity. Let e' be the radius of the open b'llll centered at 0* such that Bed (0*) C U. We also define a potential function (0) (CO)) L*.
Consider gradient descent iterates {ūr)EN following the update rule u = û; L(ū.). We first introduce the descent lemma for gradient descent.
Lemma K.6 (Descent lemma for GD). If ū, E U and 11 < 1, then
[VC(t)) < L(ū.) -(ūr+1);
and
((0141) - L S (1 (C)) L²).
35

---

Published as a conference paper at ICLR 2023
Proof. By p-smoothness,
prj
C(ū++1) 1 [(ū,) + (VL(ū,) ū,- ú,)
2 | -4.12
-((ss) 111 mM
2 V [[
$ C(ū.) ! 2 [V
By the definition of p-PL, we have
[(ü++1) -- s (1 (nn)))) -(*).
Then we prove the Lipschitzness of L (0).
Lemma K.7 (Lipschitzness of ())))). (0) is 12p-Lipschitz for 0 U. That for any O,, 02 E U,
(b(0)) - (02) S /2pl - 02112-
Proof. Fix 01 and 02. Denote by (t)) == (1 101--)0 t)0 + t02 the convex combination of 0 and 0₂ where / E [0, 1]. Further define f(t) ()())). Below we consider two cases.
Case 1. If E (0, 1).1). f(t) > 0, then f(t) is differentiable on (0, 1).
(02) (0,) [f(1) f(o)
/ r(()d
1// (vū(e(t)). 02 8, dt
(VC(O())), 02 - 01) db
VD(O(1))
(VC((())] dt.
02 01l12 -C((())-
By pmmootthssss of L, for all U,
|VC(0)) < 2p (C(0) - L²).
Since VL(0(t) L for all E (0, -2(0
IVC(1())
V2p. Therefore,
(02) (0)) S v 2p2 02 01112-
Case 2. 3t' (0, such that f(t') then
((2) (0,)) = I5(1) - S(o)
= = = t²) f(1) ((1) - f(t') f(t),
+t' (((t) - f(0)
1-
8
[(1) S(0)
max
1-1²¹
Since O(t') minimizes L in an open set, VL(((t')) = 0. By p-smoothness of C, for all O E U.
(()) < 0 - f 0('')
((() - / - ((t))l22.
36

---

Published as a conference paper at ICLR 2023
Therefore,
f(1) - V 2 00 0())l22 (1 - (1-t)) / 10 - 01ll2
P
f(0) - ( 01 0(t)1l2 = t' / 01112
Then we have
6 L1062 01112. (02) (0:)) - V 2
Combining case and case 2, we conclude the proof.
Below introduce lemma that relates the movement of step gradient descent the change of the potential function. Lemma K.8 (Lemma G.I in Lyu et al. 2022 ). If E U and nj S 1/p2 then
((())) (ú(+1)) 2 /2p n [ 122
Proof.
C(ū,) ((ū++1)
(ū:) (((++1))
v(ū,) + v (00+11)
((ū,+1) C(ue)
2
77 P271/ /2) /2/
(
where the two inequalities uses Lemma K.6
By p-PL,, ( - APp
Therefore, we
have (ū;)) (ū,+1))
V 1la N Lيب V((o)ll
Based on Lemma K.8 we have the following lemma that bounds the movement of GD over multiple steps.
Lemma K.9 (Bounding the movement of GD). úg initialized such that ( 0-112 S 1 V S.c,
then for all t 2 O, iz Bdd (0*) and
8 | uo Sool2 S V ((o)).
Proof. We prove the proposition by induction. When = 0, it trivially holds. Assume that the proposition holds for ú,, 0 < T < t. For step t, since ú, B (8*). we apply Lemma K.8 and obtain
| - uol12 v
1
VL(
8 S
/ (w(ao) - ((())) s V
((())-
Further by p-smoothness of C(.).
88 ((ū)) s 2 [io - 00 - / 11 uall2 S V P
1 2
Therefore, || 0" S A u + و
0 < c', which concludes the proof.
Finally, we introduce lemma adapted from Thm. D.4 of which bounds the movement of GF. Lyu al. (002).
37

---

Published as a conference paper at ICLR 2023
Lemma K.10. Assume that ||00 lg converges to a point U and
1- < / Vge. The gradient flow O(1)
d2(0(0) dt
starting at
0u lim O(t)
2 S VV C(0b) - L 1 V 0o
121
+00
Proof. Let T == it: : 0 $ U). Then for all < T,
dt (C(6) c2) = (C(O) L2)-112. (VC(O)
d
d0
dt
1 [C( V
de 1
2
dt
J-PL, (VC(0)l /2u(C(0) L* Hence,
dt (C(0) - c2) 1/2 5 <- /2 d6
2 dt 122.
Integrating both sides, we have
(0(T))
2 - < /f 00 < é
/
dr 22p
where the second inequality uses p-smoothness of Therefore, T to0 and (t)) converges to some point in U.
K.4 CONSTRUCTION OF WORKING ZONES
We construct four nested working zones (T (2,,, in the neighborhood of Later we will show that the local iterates i,z E t and the global iterates (()) E T with high probability after Olog 4t rounds. The following lemma illustrates the properties the working zones should satisfy.
Lemma K.11 (Working zone lemma). There exists constants 6o < 61 < E2
< 63 such that
119,19 1 satisfy the following properties:
1. C satisfies μ-PL in I for some p >
2. Any gradient flow starting in m converges to some point in r. Then, by Falconer 1983 (()) is C in I
3. Any E T has n-neighhorhood B (0) such that B (0) C
4. Any gradient descent starting in Eo with sufficiently small learning rate will stay in T
Proof. Let ( be initialized such that (()()) E r. Let Z be the set of all points on the gradient flow trajectory starting from g(0) and Z" be the e-neighborhood of 2, where is positive constant. Since the gradient flow converges to (0), Z and z are bounded.
We construct four nested working zones. By Lemma H.3 in Lyu et 2022). there exists an 63- neighborhood of I, 1s,, such that L satisfies μ-PL for some μ > O. Let M be the convex hull of T U 2 and M be the n--neigboorhood of MM where e is positive constant. Then M is bounded.
Define P2 suPde.s [(2((0))l2 and Ps sup (())()) By Lemma K.10 we can
construct ez-gehhhboroood of r where 62 < V Peg such that all GF starting in T verges to By Falconer (1983) +(-) C² in 1 Define V1 = supaerss [0)(0))2 and 1/2
supoer [002()))2 We also construct an éj neighborhood Ir where ej S J 1 E c3 such that all 0 T has 61 neighborhood where t is well defined. Finally, by Lemma K.9
there
exists an eo-neighborhood of I where eg < 1 H c1 such that all gradient descent iterates starting will stay in le.
in Ee with P2
38

---

corresponds to the loss value Published as a conference paper at ICLR 2023
Note that the notions of 2", M P2. Ps. M1, and V2 defined in the proof will be useful in the remaining part this section. When analyzing the limiting dynamics of Local SGD, we will show that all 0 stays in 1 Ui re, ( E 1 with high probability after Oloo + rounds.
(a)
K.5 PHASE ITERATE APPROACHING THE MANIFOLD
The approaching phase can be further divided into two subphases. In the first subphase, (()) is initialized such that () E r. We will show that after a constant number of rounds 8o. ((00) goes to the inner part of T such that ||@(=o) < 1|2 I Ceu with high probability, where 0 < < 1 and the constants will be specified later (see Appendix (.5.2). In the second subphase, we show that the iterate can reach within (v))) distance from after (log ",) rounds with high probability (see Appendix K.5.3).
K.5.1 ADDITIONAL NOTATIONS
(+) Consider auxiliary sequence {ū, where ūn
((*) and u a.!) 17 VL sts
(s) H 1. Define A == ket dh to be the difference between the local iterate and the gradient
descent iterate. Notice that A (+) 0, for all and
Consider gradient flow {u(t))e2a with the initial condition u(0) = (0) and converges to (()) Т. For simplicity, let u, u(sar (sa + fn) be the gradient flow after rounds plus steps. Let 8p be the smallest number such that uo of 1).
(40) (0) 1| S 1 co Note that is constant independent
In this subsection, the minimum value of the loss in Appendix K.3
r, i... _ L()), VΦ e r.
We also define the following sequence (7 ())
St=0 that will be used in the proof. Define
[-] I (--V
Σ
(e)
[(e) = 0.
2,.7
T=0 1-++1
K.5.2 PROOF FOR SUBPHASE
First, we have the following lemma about the concentration of [
Lemma K.12 (Concentration property of 2 )t Given ((s) such that ú E I U Z for all S t < HI then with probability least 1 5,
1 S 0
2H log 8
2HK
AA SISH
where Õ₁ exp(ap2).
Proof. For each (()) k.t construct a sequence ( (7(1) k.1,1' 1o=o
2(1 (+)
k.t.t' ==
II (I - v ((((()))))
=k.t
(0)
2 =
7=0
1=++1
Since ( E U we have (( |2 - P2 for all Stg H. Then, for all and
II (I - (-)((((((()))
< p2n) S exp((ap2) či
1=++1
39

---

Published as a conference paper at ICLR 2023
Notice that for all 0 < 1 S H, k,,, Hi= is martingale with 12 By Azuma-Hoefdding's inequality,
2(5) -,,,t'-1112 -1l22 S Cio
max-
E²
P(
+ c') S 2xxp
I 2 exx
21
((immux
2
2H
(Ci
(Ciomax
Taking union bound on all l E [K] and 0 St S H, we conclude that with probability least
6,
2HK
1 ll2 I C max V 22 log
AA stss H, E
The following lemma states that the gradient descent iterates will closely track the gradient flow with the same initial point.
Lemma K.13. Denote : sup120 VL(u(t))l as the upper bound the gradient on the gradi-
flow trajectory. If [|u)) ui 2 O (v))), then for all 0 sts H, the closeness of up and
(0)
(s)
u4 is bounded by
u, ( 112 S Õ2 1ú) 1/2 + +GnG.
where Č₁ exp(app))
Proof. We prove by induction that
| (x)
u
(5) 2 S (1 + p2n) (1 - u 2 227² Σ (1 + p21)
(34)
7=0
When = 0, 34 holds trivially. Assume that (34) holds for 0 S S then
sat((+))
(0)
(()) - -1/V [VV((( u1
VL(u(v))de)
i+
U1+1
satty
((0)) 4 -
(VL()(()) - vl(u()))
sat((+1)
(VL(u((") v((u(v))) do.
satta
By smoothness of L.
[V((u))) vL(u(e))]l S p2llu! u(v)ll2
S P2
[V(((w))lladw
$ P2n
Since G + p27)* I 1/ pen) S exp(ap2)nG, then [| u O(vT),
which implies that a(!²) E M Hence, [V - C(u!")) )ll & p211 u,
(*)
122-
By triangle inequality, |جايا 222i1l2 < (1 + p27) ū. - u.,²) ||2 + P2
s (1 + p27) ++1 |(( (00 ||ا + p2n2 (1+
T=0
which concludes the induction step. Appling
p2n exp(p2n), have the lemma.
40

---

Published as a conference paper at ICLR 2023
Utilizing the concentration probability of { (2 can obtain the following lemma which implies that the Local SGD iterates will closely track the gradient descent iterates with high probability.
Lemma K.14. Given Ø(*) such that u(.)) r U 2' for all 0 sts H, then for 6 Oppoly(n)),
with probability at least 5, there exists a constant C33 such that
|10 (8)
ul("|12 Coy nlog no
I S [K),
and
|0( u s Csyy n
log no
Proof. Since u(/)) r u for all s H, we have (((((u)) )2 S P2. According to the
update rule for ((0) k,t
and u!!)
(s)
1,4+ ( vVC (5)
(35)
(a) (4) -VVC(u
t+1 u,
(36)
Subtracting (36 from (35) gives
(s)
A (6) (VC(O()) vL(T, (3))))) (+)
b,++1
(I C( A
(3) ) k,t n (5)
+ 1 k,t''
(37)
Here, (3) = (1 s
ket + 86 (0))))))
" where g E (0, 1) depends on k.t and (0) Therefore,
| ll 2 (4)
if o0 E M Applying
37 t times, we have
(s) A
II [I v ((a(t)))
A (s) R,0
Σ II (1 -
h.t
===
+=01=7+
I-
+ "
By Cauchy-Schwartz inequality, triangle inequality and the definition of kg if Sr
and & [K], of
MM then we have
|
(+) k.
a
n/|
TP3 Σ CH
(38)
2
Applying Lemma K.12 and substituting in the value of H, we have that with probability at least
1 6,
2a 2nK
7 1 S Čia max
V
log
Vk K H.
(39)
1 nō
Now show by induction that for O(poly(n)), when 39 holds, there exists constant
BaK O2 > 20 max v /2a0 such that || Á (illl S /2/ 17 log m
When 0, A (e)
0 0. Assume that | (0) C2
2aK for all k E 1 qē
Then for all 0 < < t 1, o E M Therefore,
apply 38) and obtain
A b.t
la n/| kit 1|2 + 2
Σ C 2 T=0
C
2aK
0, Co log 2aK 17
Omamax
200 log
+ )
41

---

Published as a conference paper at ICLR 2023
Given that C2 2 20 max v 2a0 and 8 O(poly(n)) when yj is sufficiently small, | S O2 Vn 2a
2 for all To sum up, for 5 = (poly(nn) with probability at least 1 - 6, | S C25 nlo
no
k E [K],0 S H. By triangle inequality,
[10(ot1) ü 2
Σ kE[K]
A R.Hlla S Eeyinlog
2a n7
The combination of Lemma K.13 and Lemma K.14 leads to the following lemma, which states that the Local SGD iterate will enter s1 after so rounds with high probability.
Lemma K.15. Given g such that ((0(0)) r, then for 8 O(poly(n)), there exists a positive constant C4 such that with probability at least 5, 5,
[(((s)) 13 -
1
F Eo + Cay/noon no P2
4
Proof.
First, we prove by induction that for ā = Oppoly(1)), when
2/
kgt 212 S Ei o max V 2H log
v0 < 11 H,
80.
(40)
the closeness (()) and uo is bounded by
||0(,) uo 1
2 ct 11 Coyy n
80 log ny
5 80-
(41)
When 0, ((0)
un (0)
Assume that
holds for round
Then by Lemma K.13 for all
0 S t S H,
[((()] u 2 s E (
up 2 Enng
(8)
= = - up
(o)
la + ÖnnG
80 nG + Cnyinlog
+ +LngG
no
Therefore, for sufficiently small 17, (0) Lemma K. 14 we have
2 v0 S 1 s H. Combing the above inequality with
[10
up
(++1)
2| (((++] u1 l|2
< ((++))
(2) su a4t n/ (n Csy/ log
u 1|2 +
11
uH
ll2
which concludes the induction.
Therefore, when (40) holds, there exists a positive constant Co such that
|10(00) (((0))) 1|2 s Eay n log n1
I
By definition of (0))
1 F
[1(()))
1100 <
Ezg. nlog no"
/V P2
Finally, according to Lemma K.12
440) holds with probability at least 8.
42

---

Published as a conference paper at ICLR 2023
K.5.3 PROOF FOR SUBPHASE
In subphase 2, show that the iterate can reach within Ō(VT) distance from I' after O 1) rounds with high probability. The following lemma manifests how the potential function (((=)) evolves after one round.
Lemma K.16. Given 0(0) for 6 O(poly(n)), with probability least 6,
( h,t E E
& ((())) - Csy 1og n7 ( k
T
4k E [K]. S1SH
and
Øl+t1) E r ((0(+t1)) & exp(-oo1/22 (xpp(-op/)(ū(00)) + Ossnnlog nō
where Cs is positive constant.
Proof. Since ((*) to then for all st H, ai!") pe by the definition of the working zone.
By Lemma K.6 for
P2
[(ū))) (() L S - mi) (C(0(0)) -2-) s ((((()) - L-, VOLIS H.
Specially, for & == H,
((ū())) L < (1 (1- (1-) (C(Ol)))- -c*) (ex((-a))((((0)) L*).
Therefore,
(u exp(-ap1/ (0!.)). (g1o))
According to the proof of Lemma K.14 for å = Oppoly( (1)), when
2a
2
lll 2 & E1 Cia
log
Vk [K],sst H,
(42)
V
17
n
there exists constant Ö3 such that
[|0(3] k,t ai! ui!") 1|2 S Csp
log noo
H, 0A H, [K],
and
[0(++]] - s Osy.1 1 no
(+) Since u,
E Te, 40 < <t H, 0 (+1) T and 0 re H,
By Lemma K. 1 (-) is Lipp-LLps in M Therefore, when 42 holds, there exists a constant
Cs = C2 such that
w(O(!)) s (a( v 2 a(("1) ( + E' nlo n1
1
and
( S ((((())
/ 2p2 0(++1) 2
- ll2
< exp( (ap(-au/2)]0 2 + +G3 no no
Finally, by Lemma
K.12
42 holds with probability at least 8.
43

---

Published as a conference paper at ICLR 2023
We thus led to the following lemma which characterizes the evolution of the potential ((((s)) and v ( over multiple rounds. 2,1
Lemma K.17 Given ||0 2b - V P2
H co for ā = (ppoly(1)) and any integer
S R S
Reor with probability at least Ó,
R
0 1 re, (0(+)) I exp( 2us
1) log
44
3 S R.
exp(-aap/2)
17
(43)
Furthermore,
R 0 E r ( ((()))) + Esp n yogg no¹
V0SISHOS><<RAEM (44)
Proof. We prove induction that for d - (1poly(1)), when
12 S Ci max
2a
log
2Roc
4k:
H,
(45)
"
1
then for all 0 < 8 5 R. 43 and
44)
hold.
When 0. (()) n and trivially holds. By Lemma
K. 16
holds. Assume that
43
and hold for round 1. Then for round s, by Lemma
K. 16
0
T and
R
(((())) S <xxp(=ap/2))(00-1 + Ct n log
n
< expp(aaus/2) + - exp( exp(-au/22 Cs
7 log
n
where the second inequality comes from the induction hypothesis. By Lemma K.10
2
||0(.)
c5( 12
(0()))
22
S 2
2
R noo
(0)0)))
Ös
n
22
22p(1 exp(- aap/2))
n7
$ co + 22p(1 - (app-aa//2)) Cap 1) log
2
R
2
no
Here, the last inequality uses (0())) < 22 ((() 11
1| 5 / E0n Hence, when v is suffi-
ciently small, ((*) E 1 Still by Lemma K.16 ((0) E 2 and
R n5
( S T (((e)) C' /
r log
Finally, according Lemma K.12 (45 holds with probability at least 8.
The following corollary is direct consequence of Lemma K.17 and Lemma K.10
Corollary K.1. Let 81 := ap log 1. Given ||@(0) 10)0 pa probability at least - 8,
Eo, for 8 Oppoly(n)), with
(ola)) s Coy 1o log n
| 0101)112 s Cog
Bay/loog 11 115"
(46)
where ČG is constant.
44

---

Published as a conference paper at ICLR 2023
Proof. Substituting in = 81 to Lemma K.17 and applying [|0(.)) - 2 - v2 for
(())) E oo we have the lemma.
L
Finally, we provide high probability bound for the change of the projection the manifold after 81 rounds [ ( 22-
Lemma K.18. Let $1 =
20 ou
For 8 (ppol((n)), with Eo log 1. Given |0 - -000022 V 02
probability at least - 8,
1 [ - 22 I Ös log
17 V 17 log 77
Proof. From Lemma K.17 for ô = Oppol((n)), when
2a log
2
|
Cromaxx
Ak (KIK
(47)
7
(1)) then ((+) E 100, for all 0 S S 81. By the definition of Th ū E r for all 0 S H,O I 81. By triangle inequality, [|<b())) - 1 can be decomposed as follows.
(
(0) 1
1 (--1)) - -
( - (0()
)|12 +
((()+)))) 2(1(61))22
(48)
8=0
By Lemma K.14 when 47 hold, then for all 8 < 81
-
||0(++]] u 112 S Es log
81
This implies that 0l++1) E Ba (ū!!)) Since for all A E r |[0((0)]22 I then F(-) is v1-
Lipschitz Bei (())) This gives
[00()()))) - ((((( ll2 S v1110(0t1)
üH ll2
S vi Vv
/ 17
|n
81 log np
(49)
Then we analyze |0(++1) / 1l2. By Lemma K.9 and the definition of I and r there exists & E such that úi E BB) Vo S < H. Therefore, we can expand ((ū(4)) as follows:
((ū(f1)= = =(ū!!" (5) -
v )))
1/2 g2 (V(((( v((u("))
= no ū (0) [ (0))
- 12 ga (. (. +(1 c. (0) (1)) (V(() v<(al")) }]
(s)
where (00) E 10. Then we have
H
( ((z("))) 112 -
172
(1) (1 ( ( [((()) 12
2
1=0
1
H
<
1/2
[V((2.)))
)1.
2
45

---

Published as a conference paper at ICLR 2023
By Lemma K.6
1 V ) - ((ūb))) ((ū(-1) Therefore,
[ (5)
4(ūb )||2 S ((( -((ū)))
- u1 (0())))
Čhn
s U2l 2 2exp(-asp) (0 + (1 exp( oou)))
log
$1 17
(50)
where the last inequality Cauchy-Schwand inequality and Lemma K.17 we obtain
Summing up
50
$1 53
( )l2b2))2 s v2n
20(((())
Σ
exp(-auu) + (1 exp(-exppppp au/ 2))
log
176
==0
STn log log
(51)
no
where O7 is constant. Substituting 449 and
51
into
48
for sufficiently small y,, we have
[1bb(r)) - 10 S vn viC2s1
11 log 1
31
+ Crn log log ") n
1
s Es log
n log
1)
n7ō
where Ös constant. Finally, according to Lemma
K.12
(47
holds with probability at least
1- 8.
K.6 PHASE 2: ITERATES STAYING CLOSE TO MANIFOLD
In this subsection, we show that ||a = ( and [[0(etr) - -()))) = 0(16.500.50) 4A
r < Rgrp with high probability.
K.6.1 ADDITIONAL NOTATIONS
Before presenting the lemmas, we define the following martingale (mm()) ( that will be useful in the proof:
mk.0 O.
R
We also define P: gd + [ddxd as an extension of 04:
P(0) (0). e
0,
otherwise.
Finally, we define a martingale (Z
0,0 0
(x)
1 K
Σ
(()
Σ
Σ
- = 0.
ket
+
KE[K
r=0
T=0
KE[K] 7=0
K.6.2 PROOF THE HIGH PROBABILITY BOUNDS
direct application of Aoema-Hoe''dngg's inequality yields the following lemma.
Lemma K.19 (Concentration property of k.t).
(*)
With probubility at least 1 1-ô, the following holds:
m 12 & Cor/ log n
1
YO SLSHH, [K].0 S < Rgrp
where Čo is constant.
46

---

Published as a conference paper at ICLR 2023
(4) Proof. Notice that 1,1+1
mi,ele Omax- Then by Azuma-Hoeffiings inequality,
e2
P( ( d) S exp
2to max
Taking union bound on K clients, H local steps and Agm rounds, we obtain that the following inequality holds with probability least 1 ò:
2K 2KH
m 2 S max 2H log
v0 [K],0 Rgrp
Substituting in and Rgr
Lams
yields the lemma.
Again applying Azuma- Jzema-Hoeffding's inequality, we have the following lemma about the concentra- tion property of z(()
Lemma K.20 (Concentration property of 2 With probability at least 6, the following in- equality holds:
I
2 C12n
0.5 0.5
log 118"
40 S s < Rgrp.
Proof. Notice that | 1 120 40 30555 and 12(++1) By Azuma-Hoeffding' inequality,
-z/1 S 1/20 max-
P( 00 222)
2(sH + 1)2302
max
Taking union bound on Rgep rounds, we obtain that the following inequality holds with probability at least 1 - 6:
2R 1 11 S Tmax22 2H Rgrp log
4O S < Rgrp.
8
Substituting in H and Rgep
Larp yields the lemma.
We proceed to present a direct corollary of Lemma K.17 which provides a bound for the potential function over Rgep rounds.
Lemma K.21. Given |000 (0)) 12 S Coy Coyn nlog where Co is constant, then for ó
(ppoly(n)), with probability at least 1 - 8.
g(o) E r
(((()))) < Ciy poog no"
< 5 < Rgeps
(52)
and
( k.t
E
v(0 S Ci / 1lo 170
VO < RETps 0 0.50 H1 k E [K],
(53)
where C1 is constant that can depend on Co-
Furthermore,
(0((:))) Croy nlog 1
n9
where Co is constant independent of Co-
47

---

Published as a conference paper at ICLR 2023
Proof. By P2-smoothness of C. ( < Cov nea log 1 Substituting Rgrp = Langz and 2
( 0) S
log into Lemma K.17, for d = (poly(nn)) with probability at least 1 - 6, 52 and (53) where C₁ a constant that can depend on Co.
Coy
nes
2
Furthermore, for round 0 ((
Rgw T ( S exp( (xp((-((T-9) + 1 - exp(- aap/) Coynlog nd I /unynlos n7
where Cg is constant independent of Co-
Lemma K.22. Given ||0(0] (0(0) l2 s Cov n log where Co is constant, then for ô
Oppoly(n)), with probability at least 1 - 8. for all 0 S so < Rgrp, 0 S t H. k [K].
1
1
|a & C2p/ 1 log 176"
|0 12 S C2// 17 log
n/
- - 1 S C2
log
(0(+t)) - 010) 2|2 S C2 C2. nog log 175"
where C2 is a constant that can depend Co. Furthermore,
I 1 ( (B 1la S Cupn 170*
where Cn1 is a constant independent of Co-
Proof. Decomposing titt by triangle inequality, we have
|
la
[0
060) 2/2 - |10(,) sb(*) 1|2-
We first bound [|0()) ( 2/ By Lemma K.21 for = O(poly(n)) with probability at least
2
2 log
(( s ci
v0
8 < RRD
(54)
no
w())!)) s Ci
11
log
2
V0 S s < Rgep 0 5 H,
(55)
n70"
and
2 ( s &Coo n log no
(56)
where C2 is constant that may depend on Co and Čι0 is a constant independent of Co. When (4 and (56) hold, by Lemma K.10
110(.) - 1 1|2 S
V 1
(0())) sci/
2
2n log
2
(57)
n
(Rrp)))) (Rarr) 1|2 <
3 /
(0() < C10y/
2n
log
2
(58)
118"
Then we bound [|0[[]] - ||2 By the update rule, have
(( ( n [ n > 20 x
VL
(5)
(9)
7
7=0
48

---

Published as a conference paper at ICLR 2023
Still by triangle inequality, have
(+) k,t 0 S
|C(O + ml|mm +
Due to p2-smoothness of L, when 55) holds,
2 (222 s 2 n7
(V(0) 11 I
(59)
By Lemma K.19 with probability at least
2
|l Imfllll S Co Co. log no
2
0A S [K].05 Rge.
(60)
Combining 55 and (60). when 55 and can depend Co such that
(56) hold simultaneously, there exists a constant C3 which
10 112 C3p log no
AA [K],0 IK). H.
(61)
By triangle inequality,
[0(++1) - -0(z))12 S n 1 17
Combining 57)
58 and (61). we complete the proof.
Then we provide high probability bounds for the movement of (bi*) within Rgep rounds.
Lemma K.23. Given ||0(0] (0) la s Co nnog 1 where Co is a constant, then for å (p)oy(nn)), with probability at least
(1 11 S Cay 0.5
log no
v1 S Rgrp
where C4 is a constant that can depend on Co-
Proof. By the update rule of Local SGD,
H-1 (L(
H-1
0 = ((5)
- 1 Σ 20
17
Averaging among K clients gives
H 17
( 1) = ((+)
Σ
VL
Σ Σ
(s) k,t
K
K
L=
kE[K]
I=0
kE|K]
By Lemma K.22 for 8 O(poly(1)). the following holds with probability least 1 1 -5/3.
[10(]] g(+) 112 I C2y nlog
no o
3 (o(()) v0 Rgpp 0sts H, [K].
(62)
[0+++1] 0(0)112 S C2 nog
3 0(s), 0(++1) (ol))))) VO 5 Rgpp
(63)
710*
49

---

Published as a conference paper at ICLR 2023
When 62 and 663
hold, we can expand & ((((++1) follows:
= =0(0) (6(+1))
(
)()))
1
& (0()) (00))))(06++ -(()) 0(+1+) o
H-1
#-1
=5(0)
"1
T. (
a (T
Σ
K
K
[K[K]
1 kE|K]
7
Ts²
1 8
(a(50)))))
(1
0l+)) g(0), g(o+1) 0(0}}.
where (()) E (0, 1). Telescoping from round 0 to 8 - 1, we have
s-1
[ 66 = ΣT
Σ
Ts
r=0
From 63
3 log can bound - ll2 by 11T/ 122 S C3n
We proceed to bound [|T[0) 112-
np
When 62 and (63) hold, we have
a (
00 VC -
((a), VC(( k.t)
Po1 0(.) (1 10000000011111 0 vL(( k.)
where (5) k.t
E (0, By Lemma K.17 with probability least 1 8/3, the following holds:
2 (0 s Ciyy C
2og no
Yk E [K] < t H,O s < Rgp- (64)
When 62 and 664 hold simultaneously, we have for all 0 5 < Rgip
H
|| <
TJU2
10
o
1211
K
t=0
aV2 Ou2 C2
3 n7
n log
K
Finally, we bound|
|| Σi= Tr(1)11 By Lemma K.20
the following inequality holds with probability
at least 1 - 8/3:
(5) H C1277 0.5-
2 S
log no
4A < s < R Rgrp
(65)
When 62
63
and
65) hold simultaneously, we have
12 17|
/ 1l S C127 0.50
3 log no
YO 1056 < Rer
==0
Combining the bounds for |T 2|2 || I 1 and |TT and taking union bound, we obtain that for 8 Oppol((n)) the following inequality holds with probability at least 1 - 8:
1|
I 1
| o (0) S Can 0.5-0.9
log
v1 S s S Rep.
where Ca is constant that can depend on Co.
50

---

Published as a conference paper at ICLR 2023
K.7 SUMMARY OF THE DYNAMICS AND PROOF OF THEOREMS AND J.2
Based on the results in Appendix K.5 and Appendix K.6, we summarize the dynamics of Local SGD iterates and then present the proof of Theorems J.1 and J.2 in this subsection. For convenience, we first introduce the definition of global step and ò-good step.
Definition K.3 (Global step). Define I as the index set {(s.t) 2 0,0 S f H) with lexico-
graphical order, which means ($1, t1) 3 ($2, 12) if and only if < 82 or (81 = $2 and t S t2). A
global step is indexed by (s. t) corresponding to the t-th local step at round
Definition K.4 (ô-good step). the training process of Local SGD, we say the global step (s, t) 3
(Reot, 0) is 8-good if the following inequalities hold:
(r) k,r
le S explaxp2))msse /
2H log
6H RHot
Vk E [K], ( (8.1),
6K H Root
|m k,7 12 I mmax| y
2H log
4k E [K], (r, 7) 3 (s,t),
8
2
(r) H 12 S 0 RH grp log
v4 S <
6
Applying the concentration properties of - 2, and / (Lemmma K.20 K.19 and K. yields the following theorem.
(r)
Theorem K.1. For 5 (ppol((n)). with probability at least 1 - 6, all global steps (s, = (Reat, 0) are 8-good.
In the remainder of this subsection, we use (-) notation to hide constants independent of d and n
Below we present a summary of the dynamics of Local SGD when g(0) is initialized such that o (((()) E I and all global steps are 5-good. Phase I lasts for so + 81 = =Oooo (/) rounds. At the end of phase the iterate reaches within O(y n log + from i..., [[((Hots))) (6600t1)))
O1 log ) The change of the projection on manifold over so + $1 rounds, (0( - (000)22
is bounded by Ollog 1 1) log 175
After so $1 rounds, the dynamic enters phase 2 when the iterates stay close to I with ((s) re, Vso +81 SS S Reot and of:)) r Vk E [K], (so + 81, 0) 3 (s, () N (Brot: 0). Furthermore, [1a) ll2 and || (+) llz satisfy the following equations:
|a=()]lee O V log +/)
Ak [],0 H, 80 + s1 3158 Roots
|a< |l2 O(v 1 nlo ) Moreover, for + < can be bounded as follows:
Vsp + S < Root-
< Reot Rgrp the change of the manifold projection within Rgrp rounds
[(((tt)- - (x) 12 O( 0.5 0.5-0.58
0.5
log
A1
S Rgrp-
no
After combing through the dynamics of Local SGD iterates during the approaching and drift phase,
we are ready to present the proof of Theorems [.1 and J.2 which are direct consequences of the lemmas in Appendix K.5 and K.6
Proof of Theorem J.7 By Lemmas K.15
K. K.22 and Corollary K. for & = O(poly(n)), when all
global steps are d-good, 0(*) E 1s Vso + 81 S Riot and
ket E r Vk [K], (50 + s1,0) <
(s. > (Reot, 0) and ||a 1l2.
||=/ ll2 satisfy the following equations:
(c)
| ll2 O(v log ) | ||2 O(V llo #
Vk E [K],0 H, + 31 S < Biots
Aso + 81 S 8 < Boot-
51

---

Published as a conference paper at ICLR 2023
Hence
a Reoi) 1l2 = (O(()))))) = 0
( Rio (12) = =1( मेي by smoothness of C and
Lemma K.10 According Theorem K.1 with probabiiyyy least 1 6, all global steps are d-good, thus completing the proof.
Proof of Theorem J.2
By Lemma K.23 for 5 = O(poly(1)), when all global steps are 5-good, then
vss + S Reot
Rgp.
[ - -b 2|2 0110.500.55 0.58). 4D
S
Rugip
Also, by Lemma K.18 when all global steps are d-good, the change of projection on manifold over 8o 81 rounds (i.e., Phase 1), (0((0++=)))++++++))))))))))))( x
(0) ||2 is bounded by ((vn)- According to Theorem K.1
with probability at least - 6, all global steps d-good, thus completing the proof.
K.8 PROOF OF THEOREM 3.3
In this subsection, we explicitly derive the dependency of the upproximation error on a. The proofs are quite similar to those in Appendix [_. and hence we only state the key proof idea for brevity. With the same method as the proofs in Appendix K.5.2 we can show that with high probability, ||0(*)
2 12 $ 2 VV pa after si = (()) rounds. Below we focus on the dynamics Local SGD
thereafter. We first remind the readers of the definition of (Zg,):
(I (-- (I-nv><(a(())
(+)
[ = 0.
R,r
2,,7
r=0
M=++1
We have the following lemma that controls the norm of the matrix product I (I v ((u)
Lemma K.24. Given ((0) E 1 then there exists positive constant C independent of such that for all 0 E T < H,
II (I-nV ( SC'. |=++1
Proof. Since ((e) 1o0 then at E re: for all 0 S < We first bound the minimum
eigenvalue of 4 V²((ū,))
(o)) Due to the PL condition, by Lemma K.6 for n I p2'
C(ū( L s (1 -n) (C(E)))) L*) s (xp((put)(C(B"> L*), VO H.
Therefore,
(ūn!" S (xxp(-pbn/22)((001)))
22, mh By Weyl's inequality. Let Cj = p3/
Amin (V2L(ū(())
= |Amin V2 (2 Amiin (v
S psl 2= ((ul))) (²(("("("))) ))
S pall (((6 -v(ū(())
2 expl ssp(-utm/200 (2)
s psy P
s C expl- putttt//// /
52

---

Published as a conference paper at ICLR 2023
where the last two inequalities use Lemmas K.10 and K.7 respectively. Therefore, for all 0 st H
and I 7 I 1, 1,
II (1 (1- [=7+1
II
(1 m/mi {²L(ū(()))
1=++1
+ 7/ mmnu
V2 (²((((()))
1=0
S exp
(xp(-pln/2))
(66)
1=0
For sufficiently small 7,, there exists a constant C3 such that
1
C2
exp(-nln/2) = 1 cxp(=pm/2 1
(67)
1=0
Substituting 667 into 66). we obtain the lemma.
Based on Lemma K.24 we obtain the following lemma about the concentration property of (0) which can be derived in the same way Lemma K.12 Lemma K.25. Given (()) r then with probability at least 1 - 5,
2,2
2x
2aK
1l S C'grmax V R,t
log
VO SOSISH [
17
nō
where C/2 is defined in Lemma K.24
The following lemma can be derived analogously to Lemma K.14 but the error bound tighter in terms of its dependency a.
Lemma K.26. Given Ø(+) T then for 5 Oppol((1)) with probability least 5, there exists a constant C₁₁ independent of a such that
[10()]
x 122 s Cl an no"
Ciyannoo
4A v0 H [K]
and
[0++1]) 2// S Cipy on log a no"
Then, similar to Lemma 1.17 we can show that for = O(poly(1)) (1)) and simultaneously % + s where si O log 1), it holds with probability least that [[(())
8 >
Φ (5) ||
O oor log g.) Note that to eliminate the dependency of the second term's denominator on a in
44.. we can discuss the cases of a > Co and a < co respectively where cp can be an arbitrary
positive constant independent of a. For the of co group [ rounds together and repeat the arguments in this subsection to analyze the closeness between Local SGD and GD iterates as well as the evolution of loss.
K.9 COMPUTING THE MOMENTS FOR ONE "GIANT STEP"
In this subsection, we compute the first and second moments for the change of manifold projection every Reep rounds of Local SGD. Since the randomness training might drive the iterate out of the
working zone, making the dynamic intractable, analyze more well-behaved sequence (0(t) t) 3 (Roo. 0), le [K]} which is equal to {o(*)) with high probability. Specifically, 0%:)) equal o(0) if the global step (s, is 1) 1200-good and is point dmall E I otherwise. The formal definition is as follows.
53

---

Published as a conference paper at ICLR 2023
(((x)
10
30
Figure 9: A plot of (xx).
Definition
Denote {global step (s, 1) is 17000 -good). Define a well-behaved sequence which satisfies the following update rule:
K.S
(Well-behaved
sequence).
by E( the event
0/0 k,t = 0 1 e(0) + amull 1 g(r)
E.1t1 [
=0(²) c + Pmull 11
(68)
0 v(( k,1) - 1
(0) -1
(O
n)(((()))) -)kk)))
Ex Pnull
(69)
=
By Theorem K. with probability at least 1 - 1,100 0
0 k.
Ak E [K], (8, (Reor, 0).
Similar to we define the following variables with respect to (0(.)
0 ((-1)) avg
1
(5)
d(())
(((()))
K
H.
KE[K]
(e) a k
= ( bit
(10),
((s) avg
B60),
a (+)
Σ
(1()) R.H
avg.0
svg.
K
kE[K]
Notice that (0) (
k00 avg.o for all & [K]. Finally, we introduce the following mapping (0)
- Rdxd which is closely related to V defined in Theorem 3.2
Definition K.6. For we define the mapping Le (0) Radd,
(0)
Σ (nHH(A + X₂)) (E))), v n (JE(d]
where X;, U2 are the i-th eigenvalue and eigenvector of V and v, 's form an orthonormal basis Rfd Additionally, (xx) e-"-1+z and (0) 0: see Figure foo a plot.
Remark K.1. Intuitively, & (0) rescales the entries of E()) in the eigenbasis of 42 C(O) When
(() = Biag((A,.... ....)) E Rdxdd where X; = 0 for all m < xis & d, V (Eois, = (nH + Xz)) Note that & (0) also be written as
(ec(4((0) ((H C(O) V²L(())) vec vec)))())
where @ denotes the Kronecker sum A A B = A I Ia @ B, vecl-) is the vectorization operator of matrix and v(-) interpreted as a matrix function.
Now we are ready to present the result about the moments of S(stRgn) - ((aa)
Theorem K.2, For 80 + 81 S $ S Reat Ryep and 0 < 3 < 0.5, the first and second moments of Rlo) (b()) are as follows:
Bl Rood) (5()) ()))),(0)) c(5) 1=
2B (
+ +(K - 1)
(70)
Ö 1.5 2-22 + O(n))
El(bl (++ Rusv) (l())) (o ((at)) - ((0) c(o)) = 1/'² Σ (z)))) + 15-266 + Ö(n).
B
(71)
where ((.) hides log terms and constants independent of n.
54

---

Published as a conference paper at ICLR 2023
Remark K.2. By Theorem K. and the definition of 0/. (70) and (71 still hold when we replace ei with ② and replace of Rere) with ob(++ (
We shall have Theorem K.2 if we prove the following theorem, which directly gives Theorem K.2
with a simple shift of index. For brevity, denote by Ad(s)) == Φ(6) (0)) Eo = I()("))
Z. ())))
Theorem K.3. Given | 10012 = O Olynlo 1,,, for 0 < < 0.5, the first and second
moments of (or) are follows:
11' 0 (3) + (K - 1) + 0(1)15-23) + Ö(n),
E BlA
2B
1 71
E E R (Rxo))
AΦ RemD
Σa.1 061 Ö(n).
K. in the remainder of this subsection. For convenience, we introduce more notations that will be used throughout the proof. Let Ho (())). By Assumption 3.2 rank Ho) = m. WLOG, assume H, = (A1, ... Ad) E Raxd where X; = 0 for all m
We will prove Theorem
i S d and \ 2 x.... ... 2 Am. By Lemma K.2 ( is the projection matrix onto the tangent
space T T2:0 (T) (i.e. the null space of 4 V² (() and therefore, (D()(0) =
0 0
Let
0
Id-m
P1 = 00( and P1 == La Pl.
Let A
E E[a
avg. avg. q. = B[add
(T)T
k. and B = Eia | The latter two
(s)
notations are independent of l since 0 i
0 are identically distributed. The following lemma
computes the first and second moments of the change of manifold projection every round.
Lemma K.27. Given | 800 = ol log b),, for 0 S s < Rgrp- the first and second
moments of g(s+1) ((0)) are as follows:
- - = Pi Pi
(6) + 2 [B]
+ 2 ( (+) avg. + (11
1
(72)
2
- [[((((+1))
( (())) P A
(x)
(x) avg
PI + Ō 1551-500.00 0.5
(73)
Proof. By Taylor expansion, we have (l(+1)= = 1 (3(0) + 2
((00 0 ( )
avg +
1 Ma) #
(
WH. O11 (()
(10) ( 10)
A ag. 1 )
ا a agg.u
(
+ 0(1 m (9)
(3)
= (0) P ن (c)
2 1 (+) A
+
8² 4 (
avg.Ht
2
+ 0 A
avg. n + | ||2
avg.
(s)
112 + [| 11
Rearrange the terms and we obtain:
((o+1) - (0)( Pie
(a) W + 0 ( 10 T
(+)
1
2 ( (all 1 Log
2
(74)
+ + (s) (2
svg. lle M
(+)
|2||@ avg. H |
(o)
Moreover,
- - (d(et) - -
ty H avg. H
a (S)
+ 11
20 12). (75)
55

---

Published as a conference paper at ICLR 2023
Noticing that ( A ( are identically distributed for all &: E [K], we have El0 Ag(+)T]]
arg.
X [xE() B[c (a) AΦ (+) B Then taking expectation of both sides of 74 gives
B(glo+1) -661] PI + 0 (3 1
2 8 2
(0))
O A .1)
12 avg. H [1a] B la (
avg. (821 + El avg. 11)))
Again taking expectation of both sides of 75 yields
- A) P A Pl + ((EII
(())))))+
|||
|2 avg. (1)).
By Lemmas
K.22 and K.23
the following holds simultancously with probability at least - 17 100
A (+) ||2 O(10.-0050) Lv ll2 (67"
Furthermore, since for all l E [K] and (s, 3 (Reae, 0), ((:) stays in nea which is bounded set, | lla and || ll2 are also bounded. Therefore, we have
E
|||| D (5)
(e) avR, H [|2] = 0(1)15--9)
(76)
E|
A (+)
2a avg, /1 l Ö(1) 1.5 -550..9)
(+)
(77)
E (5)
avg, H 11 Ö(n)1.6,
(78)
which concludes the proof.
We compute A C. and B,d) by solving set recursinss, which formulated the following (+)
lemma. Additionally, define A = E[a (c)T [ and M³ == [l k.1 (k: + 1).
Lemma K.28. Given 11000 200122 O 17 log ) for 0 < Rap and 0 S < H, we have the following recursions.
1 = q. () Ho - 7 [B, 1 v 14
(79)
2
1+1 = A H (3) - n Ho + n/2
(
Eo + (1122500.55)
(80)
Bidc
( M (+) M (8) - Ho + 01 2.5-.53).
(+)
(81)
t-1 (I HHHH B (0) + 21522-58)
(82)
Moreover,
A 1 AG)
+ (1- M (s)
(83)
avg K
K
H
M (0 1) (0+1) = P, Á()) P. + (1) 55 0.58),
(84)
d (0-1) P11H (6) - ()) )
1
(0) avgl + 00)-5557)
(0)
(85)
( P B(6) + P A P + 1-550
(86)
Proof. We first derive the recursion for qt Recall the update rule for 0(:)
0(:) - - - n + (la)
k,t+1
56

---

Published as a conference paper at ICLR 2023
Subtracting ((() from both sides gives
(d) (d) 2,,+1
+
k
( 11
(V + [>
0(1+))11)
(9) 0(1e())112
= (0))
(2< + V L( + a/!
+ (0(000))+00000))))11111
[
(1) 1 (s) + + +2
2
a nHo v 0( -
" v
( k
+ O k.r
|1
(87)
where the second and third equality perform Taylor expansion. Taking expectation on both sides gives
q = (I- H) - VV'
" 2(
+
[1 + y - kit | + E
2. + E[ l22
By Theorem K.1 with probability at least 17 1000
vk [K], (s, 3 (Rgrp 0). Also
notice that both 0 (()) and Guull belong to the bounded set r Therefore, | bounded and we
have E| E| k.t [|] = (v)200). Combining this with 76 to 78 yields 79
Secondly, we derive the recursion for ((a). Multiplying both sides of by A and taking expectation, have
B (I B.
Or A
11
+ |
Still by Theorem K. and
76)
78
have (82)
Thirdly, we derive the recursion for A By
87
we have
1 A H A
(5)
Ho
Eo ( + [[]()]la)
Bioc
On
(10 18
|
2 (0)
11 ||2]
n12
- ( ((
Eo + 2.3-0.5)
Bioc
which establishes (80).
Fourthly, derive the recursion for M( Multiplying both sides of expectation, # k. obtain
87 by a(00 and taking
M (s)
M TH MM (x) Ho + 0nE ا l|2))
+ On | ll2 + [12)
By similar argument to the proof of Lemma K.27
have
lla] = 06))
E|| ||||| || A (0) [|2] = Ö 1.50.50)
which yields 81
57

---

Published as a conference paper at ICLR 2023
Now proceed prove (83) to (86). By definition of Aavg.
A
avg
(+)
K
K
(8)
)
Σ
kk
1
kE|K]
KEIK]
B
T
Σ
Bla T 1
K2
K2
kE|K]
=
+ (1
M
K
K
which demonstrates 83
(s+1) Then we derive 84). By definition of o avgg.0>
(++1) avg.0
(3() + a (a)
avg. ((b(.)) a I1)
((0)) + (5)
(3(0) + (T(()( ).)
avg. 0(1 (+)
(s)
avg. 111)
avg.
a00)
P1 0 (11 Lag. + 011 12)
avg,
Pi ((()) aug, H
011
(+) avg. 2/2 + avg. 11 ||2
(88)
Hence,
M(+t1
(s+11 E [(
ava, avn.o)
A P1 (
avg H
(a)
11 ( (l2)).
avg
By 76 and 78 we obtain
By
(++1)
avg. H + Ovl. | +
|| ||a. 112).
(89)
Combining 88
and 89 gives
E 2 o -3(0)) 4
(x)
P Ö(7) 15500.88)
avg
Therefore,
B (++1) = B|a avg.0
(++1)
( E| (A + -
P B (s)
P1
P (()
P + Ö(1 1.5
avg
Finally, we apply Lemma K.27 to derive 85
4o
(s+1)
E[
(1+1)
E[ (5)
Φ (++1) - [((0)
avg.
4H
(5)
P 1H
(8)
(1)
8 (0)) (s) 1-5
P
Qu
(+)
- 8" x & (0)
[B
2
1
(0) 11
avg.
(3)
+ 0
which concludes the proof.
With the assumption that the hessian at 20)) is diagonal, we have the following corollary that for- mulates the recursions for each matrix element.
58

---

Published as a conference paper at ICLR 2023
Corollary K.2. Given 1 - Of f) the following elementwise recursions.
for S s < Rerp and 0 S < H, we have
A (1 (Ar + (1) A + T² Eo + 011225-0.50)
(90)
Binc
M 141,J (1 (A+ + X1) m/ +0(1).5-0.50)
(91)
B(A) 141 (1 B, Ö((72.5-0)
(92)
A
1 [A M
Avgik.
MHH.j) + 11,.33'
(93)
K
-0(11.5-0.55
1sis m, j S m, otherwise.
(94)
((1)25-0055)
B(b)
HiJ + Ale)
Auggn in
+ (
1 i m, m < S d,
0,3
((++1)
B(G) + ((11
1 i m, 1 S m,
S
(95)
((1)15-)),
m<isd. d.
Having formulated the recursions, we are ready to solve out the explicit expressions. We will split each matrix into four parts and them one by on. Specifically, a matrix M can be split into P. M P in the tangent space of I' at 23(), P1MP1 in the normal space, along with P and P across both spaces.
We first compute the elements of PLA(") P1 and P1 A PL.
Lemma K.29 (General formula for P.A P1 and P. (5) P.... P1). Let Rg = =[ 10 log 11 Then
for 1 S m. S j S m and Ro S < R Rgrp
A Avg. (^: + \3) B T + [r 2350059
1
A (1- 1
(Dz (A+ ))
Mo
1 (Ai + X1) Bioe
Eosis + 01p
K
(Ai Binc
For < Ro, Ale), = Ön) and Aw. = O())
Proof. For 1 sis m, < S m. A > 0, > 0. By (0).
+ Σ (1 (2++ +x)m)" n²
1-1
A0l = (1 - (Ar + 1y)n)
Eo,is
D,i.j
Bioc
7=0
1-1
+ 0 Σ (1 (Ax + 23)m) ,,
0=0
= (1 - (A: + Ay )m) 164)/ + 1- (1 (Ai + 21)m)²
Eo + (616.500.58))
(), + B
where the second inequality uses Σt=b (1 -
+ \ymm)==
(A By 9I.
(Art2yla
M2, (1 - (A. + +A)))) Ō (1 - (1- (A, + 725))
-=0
=11- - (A, 11)m)( 1600
+ (-1..000.)
59

---

Published as a conference paper at ICLR 2023
where the second equality uses M M6 (0+1)
A ( By (93) and (94).
A A = (1 - (A, + 29)1) n.... + + (A: + +y)nn) Abhis + +01)1
(A: + 1y) Bioe
Abst11 Auvg..s + (1122.50053)
1 (1 - (Ai + {)))) H
nT (1 -2; + 21)n) H A + Ö61)
(A: A
Then we obtain
+ 1 1- (A, +
Ab% (1- (de + +y)))))) 100
1y Σ (1 (1+ +2)))))))
(De + B10
s-1
+ +011115-0.555 1.58
Σ (1: (A, +
+ )
r=
Notice that |1 - (): + >y)m1 < 1 and
(1 - (2x + )1)m)) - exx(-((A + = Aj)a).
(96)
Therefore,
E1- (A: (2)))
1 (1 +
+ Xz)m)
1
1- (1 - (A, 2y)n)
1 1- exxeppp(pp + Az)a)
Then we have
1 (1 (h: + (1)))
A/40 = (1 (A++ +
+x)))))))at My)m)aHHAH
A +
Mo + Ö11
(As X2)
Finally,
we demonstrate that for 2 Ro. A and Alt) A is approximately equal to
(Art2,) Eo,i.t- By (96). when s 2 Ro. (1 (As Ay)n) O() which gives
(()) A (>1 \j) Bloe
1 To + Ö(1 )
(1- (A: (Ai+
1 1 K
(A: + Bloc noo.i
7 (Ai + X, Bioc Co,,, + -33..300
For Ro, since {!)) (1) 2 (6) = have A is O()) and Al. Ö(n).
Secondly, we compute P A P1 and P. P1.
Lemma K.30 (General formula for P1
(5) (() P1 and P Alo Pt). For 1sis m,m m < <jsd
1(1) = 1 - (1 - Xmn)
Eo + 0(111100050
ABBroc
-1- Co + 011
Avggi.
Proof. Note that for 1 5 i m,m <j S d and X > 0, 0. By (90) and 94
((0) = (1 >m)) ((6) 1 (1 >m1)"
E),., + Ö(1 15.0033
A,Bloc
- (1 - >zn))
Eo.es 16
Bio
By 91 and 94).
M (x)
Then,
(,6.
1- 1- (1 A H
(s)
Eo + O(1 1.0003
avg.i.j
A Bloc
60

---

Published as a conference paper at ICLR 2023
Similar to Lemma K.30 have the following lemma for the general formula of P A Pi and Pi AIA (x) avg Pt.
Lemma K.31 (General formula for Pl P/A.)) P1 and P A P1). For <is ddddaa and 1 $ m,
A = (1 (1- X
2, nEa.s + O(1) (.66.60).))
- (1- )11)
( avgii.
Doss + 011) MyKBtoc
Finally, we derive the general formula for Pl PIA, Pl and P A PI.
Lemma K.32 (General formula for PIAy PlA, P and P PlA A PT). For < i d and < s d.
(+))
H1i² Conss + 1(111.500650
avg.i.
KB000
tnp Ea,is + 01g
Bioc
Proof. Note that for m <i d, m <j S d and A = X = 0. 00090) is then simplified as
A41.5 = A.1 + "1²
Cass + (1722500.58)
Bioc
Therefore,
A(10 Ab. +
Eosi. + 0611.5-0.50)
(97)
Bioc
According to (11). and ( yields
Ö11 for <is d and < j S d. Combining (11.
94
Hu? KBooc
A (e)
Eo.i,j + 10115-0050
avg.iy =
Now, we move on to compute the general formula for B
Lemma K.33 (The general formula for P1 B())) Pt). Note that for 1 S i S m and m < S d, when Ro = 10
log 1 < < Rgrp.
Amo
(1- den)
AKo Tposs + Ö(11.5-8).
Biad
For Ro, Bl,)) Ö(n).
Proof. Note that for 1 1 m, By 92
[() Bt,,i (1 B() + 2-52.-8)).
Hence,
B(00) (1 > [3(e) 16511.5-8)
According to 995
B(o+1)
-B(L), + A(0) ks + (1725-88)
U,i.j
Avg.
=(1 An)) H ((0)) + A(5) + 16)
61

---

Published as a conference paper at ICLR 2023
Then we have
BOll, = (1 - -1-Xom)a Bu,is + A(d)
Avg.ú.j
Σ - Xm)) + 0 06
(1
p=0
r=0
+ 1- (1 (1 A
(1 - 2pp)) B
avg.... (1)2.3-8)
1 (1 \,1))
Xe1) te
A4)
avg.isJ + (1111550).
1 (1 X,17)
where the second equality uses (66 and the last inequality uses Bi
Ae a = 0. For
(0)
2 Ro. A (w)) avgiins
1-(1-A:n) " A. Rioe
Mo 15-0555 which gives
B
1
Enjis + O11
Bie
Therefore,
(1 (1m)
B
Boo n/ + (6p1.5-))
As For Ro, A(a) 4, (n) and therefore. B Ō(1).
Lemma K.34 (General formula for the elements of P1 PI ). For 1 sis m and 1 S S m,
B(C)) (-)
Proof. Note that for 1 sis m. > By (92).
Biflis (1 B + (1)2.--0).
Hence,
B = (1- (1--)))) B + -61).5-1)
By (55)
((a)
HJ. + (112255-))
(1 Am)) B(o)) + Ō1
B + +0 6-1
(1 Amn) He 10)
(1- H
(1 Am))
4H ((0)
Boss 151
= (
where the last inequality uses ( 0.
Lemma K.35 (General formula for P B(a)), For m <is d, B, = 06g15-0)
Proof. Note that X; = 0 for m <is d. By 92 and 95
B = BB(a) + Ō( 2-5 B((0) (-1
Therefore,
B 10(11 Bbx²) 161115-0)
Having obtained the expressions for ( A and A now provide explicit expressions for the first and second moments of the change of manifold projection every round the following two lemmas.
62

---

Published as a conference paper at ICLR 2023
Lemma K.36. The expectation of the change of manifold projection every round is
E(pl(t1)- - -66
28 (()),
Hy 0 |Ea ((((0)) + + (1)
Ro < < Rgrp
(98)
Ro
log 7 where Ro = a 1)
Proof. We first compute Fl(3(++1) - ( By 72
we only need to compute P ( by relating
to these matrices. Multiplying both sides of 799 by P gives
piá. (e) -pP
B (
"
211255-))
(99)
Similarly, according
85 have
P P
P 0² 2
1
pi
(( + +017)--A)
(100)
Combining 99 and 100 yields
#-1
p (0)
P
8 2 (A)
"1 -
L( ~
Σ 1
2
avg
(101)
H-1
v (0)
C(
B B(a}| P ( | 11))
t=
By Lemmas K.29
K.32 and K.30
for
1 Ro = | log J A Ö(n), A = On)) and
{( ()) O()). Therefore, E E[
{(0) Ö(r). For fors> Ro. A A +
Substituting (101) into (72) gives
El - ((0) P1
P
(A)
P 02 x (0) (1)))) [B(]] (+)
2
avg
T
I-1
| CP
PP
1
1
A
Σ
Ta
Below compute Ti and To for Ro respectively. By Lemma
K.3
P.P4((:0) P Á avg P Pi 8 o
|P|
A avg
Pi] - 0,
())
(+) P = ( ( P
By Lemma K.4
P ( (+) P = 0.
Therefore, for 8 > Ro.
Hip ( 0615.550.50)
P avg
2K
where we apply Lemma K.32 Similarly, for 8 > Ro.
P 0 (
0 ( B Pl] :6n115-3)
where we apply Lemma K.35
Hence,
!n12
Ti
8 (0) | + Ō
(102)
2B
63

---

Published as a conference paper at ICLR 2023
We move to show that
Hn²² T2 2B
( E0.1 + (K (k))+(1)) (}
(103)
Similar to the way compute A matrices and then substituting them into the linear operator -1 V [( one by one. First, we show that
A and B, compute To by splitting Ta into four
Hn² 0 1 (3)
EE (K - 1)b (),
2B -11
(104)
where (+() is interpreted as an elementwise matrix function here. By Lemmas K.29 and K.34 for 1 S & m, 1 S j I m and Ro.
1 = (1 1 (1 K
(1 + ),)m)* no + (A; + Bto (A++Ay)
"
Eois + 017
B(,)) Ö(1 1.5
Therefore,
H-1
1l = - (1 (1- K
(1- - (Ar + 21)m)
Hn
Eass + (); + Bto
Eo.is + Ö11
()i + X,) Bioc
Hn
Ea.i..
K(A; + A
(1 K
Hn
(1- (A, + Ay)m))
Do.is + 050 8).
(Ai )z) Bine
H1i(A: + X3)
B
Ö6n -1
I=
Then we simplify T4. Notice that
(1 - (Ax + Ar)m) = =exp(-Hii HI; + x1) +
=exp(- HH( Ay)n) 0()).
Therefore,
T4 (((); + (1))))++))))))
Substituting Ta back into the expression for EH=o1 Al,,) gives
#-1
Hy
Eo.i.. + (1-1) 1
Hmpp Hn) Eos. + 033
Σ 1=0
1 = K(A Bjoc
(1-K) K
(A, + X3) Broc
Combining the elementwise results, obtain the following matrix form expression:
Hi²
2B P (oo,1 +(K - 1}{((E,,))
+ 155
By Lemma K.4 we have 1004
Secondly, we show that for $ > Ro.
C { + P1P
117 8 8 (0)
(105)
B EEo + (K - 1) (Eo.,.))) + (6125-5)
64

---

Published as a conference paper at ICLR 2023
where (--) V ())()
interpreted elementwise matrix function here. By symmetry of A and
H
n- C9((111)
Σ
Σ P
P1
Σ
PLA
Pl
Therefore, we only have to evaluate
H-1 Σ
H PL (A + B PI Σ
((B())
P
P1
1=0
==0
To compute the elements " (A + +B(")) P1. we combine Lemmas K.30 and K.33 to obtain that for 1 S m and < d.
H
H-1 (1 \mn)²
(()
Ie.is + 0(1)05--))
t=0
t=0
Hn EoSis
1 (1- >11)
Eoiid + (170.-
A,Bio
B/00
Hn
(1 1 (1 Amn) H
Eass + 0-588
A;Hn
H1 (X,HH))
((10.--8)
B1o
and
H-1
H-1
(1 AmnI''
Σ
Σ
Moo... + --)).
X K Bioc
1=0
1=0
1 (1 >in)
Eoii + 0(10,--1))
H1
Hn
(1 1-- 1- Amn)
(1-Am)
Eo,is
Eo,is 0630.3
X
29111
Hn
Hn
Ar 2o.i.
(A,lH1))oo.1 + 0500.
B|oc
Therefore, the matrix form of Et= P1 (A(
H-1
((()) Pl is
#-1
H H (Σo.L.I + (K - (Eo.,.11)))) + 0-)05-))
Σ PL( + B PI
1=0
where vi(-) is interpreted an elementwise matrix function here. Furthermore, by Lemma K.35 -H-1
B(a) a) 055 Applying Lemma
K.3 we have (105). Finally, directly applying
Lemma K.S we have
p V L
[P T3P)
(106)
Notice that do Eo,1) 0 where (b-)) operates on each element. Combining 1044)
(105) and 106
we obtain (103) By (122) and (103). we have (88)
Lemma K.37 The second moment of the change of manifold projection every round is
Hy Eo,1 + 05-055)) Ro Reg
B((((0+1) - -
O()).
Ro
where Ro = 10 Xma
log 1 1
Proof. Directly apply Lemma K.32 and Lemma K.27 and have the lemma.
65

---

Published as a conference paper at ICLR 2023
With Lemmas K.36 and K.37
we ready prove Theorem K.3
Proof of Theorem K.3 We first derive E | c (B Recall that Rep
Lamp
Hyt+3 o(1)
where 0 < 0.5. By Lemma K.36
Ra
B[ - =0001= = Σ
E|
(((B))11) 66101 +
B( -3601]
==U
s Ra
+1
11 '1-
En + (3 + 5-) O()). 0² (0))
2
Then we compute P AΦl T.
Rgp
E
( -
Σ (5( (blot)) (((0)) (
Rgm
- -b>()) _ (B(6t1) - ((+)))|+ Σ B((
- T] -
17' 1-8
B Co,1 + + Ö111551155))
where the last inequality uses B[(g((+!)) -( Ō(12).
K.10 PROOF WEAK APPROXIMATION
We are now in position utilize the estimate of moments obtained in previous subsections to prove the closeness the sequence ( 122 [T)(Hn)) and the SDE solution {{ t E [o, T]} in the sense of
weak approximation. Recall the SDE that we expect the manifold projection { (0()))) (T/(H+))) track:
( to
dc(t) Pc JB E ( (C)dW, 28 V LCC))E K- C(C)
(107)
(+) diffusion
(b) drift-1
(c) drift-ll
According to Lemma
and Lemma
the drift term in total can be written the following
form:
(b) + (c)
2B 0² (C)(E) + (K 1)))
Then by definition of Pc. [107 equivalent to the following SDE:
+ _ (C) [E(() + (K (K))
dc(t)= VB
aN
(108)
2B
Therefore, we only have to show that (6() closely tracks {C(t)] satisfying Equation 108 By Lemma K.11 there exists E3 neighborhood of I", 1 where ((.) is C -mooth... Due com- pactness of T. I is bounded and the mappings (-)()), aD(-). E.)((.) E(-) and (t) are all Lipschitz in I By Kirszbraun theorem, both the drift and diffusion term of [008] can extended Lipschitz functions on Rd Therefore, the solution to the extended SDE exists and is unique. We further show that the solution, initialized as a point on r, always stays on the manifold almost surely.
As preparation, we first show that I' has no boundary.
Lemma K.38. Under Assumptions 3. to
has no boundary.
Proof. We prove by contradiction. If has boundary ar, WLOG, for point p Ør. let the Hessian at be diagonal with the form 2 (2p(p) = diag(A,,""> Ad) where x > for < i S m and X; = 0 for < i d.
66

---

Published as a conference paper at ICLR 2023
Denote by Z.j (111" (1 - j) the ++ 1) dimensional vector formed by the i-th to j-th coordinates of æ. Since -(V((p) diag(A), Am) is invertible, by the implicit function theorem, there exists an open neighborhood V of Pm+1:d such that VL(v) = O. ve E V. Then, [(v) [(p) minget (()) and hence V r, which contradicts with p or.
Therefore, I' is a closed manifold (i.e., compact and without boundary). Then we have the following lemma stating that I is invariant for (108).
Lemma K.39. Let ((t) be the solution to 108 with ((0) E I', then ((t) E f all t 0. In other
words, is invariant for 108).
Proof. According to Filipović 2000) and Du & Duan (2007), for closed manifold M to be viable for the SDE dx (t) F(X (t) (t) where F I + R² and B IRd + R² are locally Lipschitz, only have verify the following Nagumo type consistency condition:
((x) F(x) [(B((())) (MM,, B,(x) TM(M),
where D[:] is the Jacobian operator and (x) denotes the j-th column of B(x).
In context, since for r, ( is a projection matrix onto TT(T), each column of ( ())/(((())) belongs T'')), (I'), verifying the second condition. Denote by Pi(Ø) == La (Ø)) the projection onto the normal space of I' at Φ. To verify the first condition, it suffices to show that as follows.
P1 (() We evaluate E, P1 (
Σ (L = B Σ
1 PL (0)
8 (00
(4),))
B
V"[())+>><(>> (0)),
(109)
B
where the last inequality uses Lemma K.3
Again applying Lemma
we have
P (o) = 1 VV (o))
1
(110)
2B
Combining (109) and 110p. we can verify the first condition.
In order to establish Theorem 3.2 it suffices to prove the following theorem, which captures the closeness of (6() and ((t) every Repp rounds.
Theorem K.4. If (1)(0)
((0) ll2
O( Oly b) and ((0) (0)) then for Rgep
75 every test function E C3,
Eg(ol(BRao)) - Eg((((nn (.75)))))))))))) S Cgm 0.2 (log 2))
max n==0,-. T775-7
where Cg > 0 is a constant independent of ' but can depend on g(.) and b > 0 is constant independent of n and g(-).
1.10.1 PRELIMINARIES AND ADDITIONAL NOTATIONS
We first introduce general formulation for stochastic gradient algorithms (SGAs) and then specify the components of this formulation in our context. Consider the following SGA:
En+1 = 2n + Te (xx, En):
where In E Rd is the parameter, Te is the learning rate, he(,..) is the update which depends 2n and a random vector &,. sampled from some distribution E(x,). Also, consider the following Stochastic Differential Equation (SDE).
dx((t) b(X (t) + (X
67

---

Published as a conference paper at ICLR 2023
where b(-) Rd -> Rd is the drift function and or)) Rdxd Raxd is the diffusion matrix.
Denote by Px [x, s, t) the distribution of X (t) with the initial condition X (s) = x: e.Define
A (2, = X (n+1)n a.
where X (n+ ~ Px (2, noees (n + 1)me):
which characterizes the update in one step.
In our context, we view the change of manifold projection over Rgep := Lاn E (.,)
rounds as one "giant step". Hence the (n( corresponds to the discrete time random variable In corresponds to and (tt) corresponds to the continuous time random variable X,, According to Theorem K.2 we set
1
To 1 b(()= = 2B
(C) [E(<) + (K - 1)) o(S) = VB
2 (4(
Due to compactness of T, b(.) and (-) are Lipschitz on r.
As for the update in one step, Á (c, is defined in our context as:
(Ø n) == S(n+ line Φ For convenience, we further define
where S(n+1)n ~ Pc(Ø, niee (n + 1)7)) and Φ E P.
A (n) ((xx
À (m) Å ( n)) (n) = = (()R))
6 (x) b(Φ (m RecD)).
We Cg,, denote constants that can depend the test function y and independent of ne- The following lemma relates the moments of Á ((Φ, n) b(Φ) and (Φ)).
Lemma K.40. There exists positive constant Co independent of ne and g such that for all & I,
[ n)] mabi(()) s Cone.
A1 S i d,
(6, (), Te
Σ (..(o)o (6) S Coné.
AA vi Sd
I=1
6
E
II A.. (0, (b.
Cone:
A1 i ,66 $
The lemma below states that the expectation of the test function is smooth with respect to the initial value.
Proof. Noticing that (i) the solution to (108) always stays on I' almost surely if its initial value ((0) belongs I (ii) b(.) and (-) are C and (iii) I is compact,
directly apply Lemma B.3 in Malladi et al. (2022) and Lemma 26 in et al. (2019a) to obtain the above lemma.
The following lemma states that the expectation ((((t)) for g E Ca smooth with respect the initial value of the SDE solution.
Lemma K.41. Let E [0, T) & T and g E C3, For t E [s, Ti, define
u(d, 8, 1) = ES (G
Then u(,, t) C3 uniformly in s,t.
Proof. A slight modification of Lemma B.4 in Malladi et al. 2022 will give the above lemma.
K.10.2 PROOF OF THE APPROXIMATION IN OUR CONTEXT
For 9 E (0, 0.5), define 71
1-529 72 == g and then 1 71 < 1.5, < 2.
We introduce the following lemma which serves key step control the approximation error. Specifically, this lemma bounds the difference in one step change between the discrete process and the continuous as well the product of higher orders.
68

---

Published as a conference paper at ICLR 2023
Lemma K.42. If |1000 1/ = 0(( n log }),, then there exist positive constants C₁ and independent of The and g such that for all S n 1T/me].
1.
E (1)
(1)
E g(n
1 Cinz (log ) + C ne n
sigisd,
|B| (m)
(m)
E e(n
- Cn (log )0 + Cin y,, vI sisjsd
2.
6
IF
A (n)
el Repp)
S SPn 2 (log no )
A1 i, , S d,
6
E
A
(m)
E c(nBarr
I CiTe 201 (log 2
$ 41. isdd.
Proof. According Appendix
K.7
have
E
II
A (m)
(u
(11)-38).
Since 71 < 1.5 and 72 < 2, we can utilize Theorem
K. and conclude that there exist positive
constants C2 and b independent of 17 and such that
T nobg E cfmBBood) S C2n (log , + C2n (log 1 v1 sisd,
(111)
d BIA, (1) Ann) 1) Σ
(n)
&
efnBawd)
S C2n (log ) + C2n (log 1)²², VI < i,j < d,
1=1
(112)
E
(n)
Ed
S c2n 2 (log ne 2 v1 < 11, ... , v d. (113)
Combining
111
113 with Lemma K.40 gives the above lemma.
Lemma K.43. For test g E C3, let un, (Φ) == u(Φ, MM,e mnee) B.C-Pe (c ,m. (C())
If - 1101122 = Oly p0ly/oo log 1,), then for all 0 < 1 S n and S n (T/ ///elll
E[ (( + A (1) ( A (++1)) ((unl) S Cy + n log( (
where Cg.1 is a positive constant independent of 1 and (( but can depend on
(b) E C3 for all / and n. That is, there exists K(-) E G such that for all I, n. ui, (() and its partial derivatives up the third order are bounded by K (()).
Proof. By Lemma K.41
By the law of total expectation and triangle inequality.
E Elu[+1, (1 (( + (1) 11 ( 3 (R + A [1)] (1Baa))
Bluu+1,0 ((rp)) + ()) u1+1, ((Rr))
(()) 3 (())) ((R
+ 100 (6 (())
(1))
3 ( a(lRar)
Aa
100 B I, ( (
A (1)) (1
a(lRwv)
1,n
As
69

---

Published as a conference paper at ICLR 2023
We first bound A₂ and Ag. Since ((())) Г, both (()))))) + A and @ Á belong to I. Due to compactness of I and smoothness of u441, on r, there exist a positive constant Cg22 such that A2 + As I Cg
We proceed to bound A1. Expanding (1+1,n()) at ((ll and by triangle inequality,
A
d
E
111+1.n (bll)ao))
ellB(o)) - (( en
Do:
Bi
1
8 uutiIn (B())))))
d((RET)) Eo c((Ror)) A -
E
2
05,dō;
B2
+ R + [R],
where the remainders R and R are
8 Puiia
R
E
(
A (1)
AA A ( Eg cftRoro)]
do.06,dbp
Ishiipd
R
1
Σ
E
(1()) 05
8 Putia
03 A A. 1 A (1 Ey e((Rere),
IsijpEd
for some 0,0 (0, 1). Since ( belongs to I which is compact, there exists a constant C9.3 such that for all 1 S i.j S d s < S [T/mell
Du+i.
8 Put+in ( (1 S Cg.3.
(( Cg.3,
0Φ:
ds,00;
By Lemma K.42
B₂ S d C C1.0 + t )log 1 1
Bi g. n n (log 2
Now we bound the remainders. By Cauchy-Schwarrd inequality,
Juutin 0o,
(6 (R +0A ( AA A A 16 (((R Eg c(lBxao))
B
(8 ( + 01 0 (1) 2 ,
8 Juutin
((lka))) e ))
1/2
<
E
x
22) ((1B (f e(nRau)) 1/2
Since ( and (r))) + +A (1) both belong to I' which is compact, there exists constant C/,4
such that for all 1 S i.j.P < d, 0 I and S n S [T/nel.
Parii ((R AA 3
$ c22 c
00.00,00,
Combining the above inequality with Lemma K.42
we have
JuILIn (11R + +0A4 11) A
A
(1)
(1xx)) el c(llBar))
< Cg..Cime log( 22)
D0p Hence, for Snn s
&
R
CynCin log( (
Similarly, we can show that there exists constant Cy.s such that for all 1 s n [T/ne], 555
J d Cn ())
Combining the bounds Ai As. have the lemma.
70

---

Published as a conference paper at ICLR 2023
Finally, prove Theorem K.4
Proof. For 0 I v n, define the random variable Śi. which follows the distribution (o( n, conditioned on @(1B))) Therefore, P(Gm,, = (((Rar) = 1 and So.n San-- Denote by ul (1 1) == Pe ())(()) and Ti+1, u1t1, (0 (6 + + (1) (1+1)% m)))
2411n R R(o))
A (1) (1 + Me, me).
3 B n Bnn
11 - Elg( (nae)))
Blg((nm)) s(So,n) Rem)) + O(n 200)
BlolC(+1,n) -(((.)) Eo cfnBood O00)
Blu( 1) (1 1) nTe)) u( m))) & ) O17000)
B[T++1, efmRoodd + O111100)
Noticing that [T |
Ef ) = [4411 ((1 ( cllRen)) | efnRBer) we can apply
Lemma K.43 and obtain that for all 0 S [T) [T/Te).
E Elg( Big(S(nne))]
s nn + (log 1 TC g. (n=-1 + n -----11 (log 1===11 n. 2.)0
n/
Notice that n + = = 0.5-8 + r and T, C,21 are both constants that are independent of TJe. Let 3 0.25 and we have Theorem K.4
Having established Theorem K.4 we are thus led to prove Theorem 3.2
Proof of Theurem 3.2 Denote by 8els = sg + 81 O(log #) which is the time the global iterate
g(s) will reach within Ö()) from I with high probability. Define E(t) be the solution to the limiting SDE 108) conditioned on E( e(*at) and (()) (f(e))) By Theorem K.4 we have
E Eig( (nlletnlonts 0(((n10077) | (((a))) gf g(ta))
S 0. (log 7) 2),
max n=0,-. 1
where Rgep
Larp.r.. Noticing that (i) g E C3 (ii) b, E C and (ii) <(t). E er, e (0,
almost surely, we can conclude that given E c(ma)))
|I((1) - E(1](2 Ö(V),, E T.
Then there exists positive constant 1/ independent of and g. and C% which is independent of but can depend on such that
Eig(o g(((nn) 0.75 + sels (otsH1 s 1, 12=00" 0.77
0.25 (log 1
max
We can view the random variable pairs (( +sa.) Snmp -5 ssnaan) : 0.... 1T757 as reference points and then approximate the value of (( ()) and g(((sHnn2) with the value the and
nearest reference points. By Lemmas K.18 and K.2 for 0 Sr I Ryrp and S S Reat r,
( [|2] Ö1 0.375)
Since the values of (0)) and { are restricted to bounded set, g(+) is Lipschitz on that set. Therefore, have the theorem.
71

---

Published as a conference paper at ICLR 2023
L DERIVING THE SLOW SDE FOR LABEL NOISE REGULARIZATION
In this section, we formulate how label noise regularization works and provide a detailed derivation of the theoretical results in Appendix G
Consider training a model for C-class classification on dataset D = { y:) where 2; denotes the input and yi E [C] denotes the label. Denote by 44 the (C - oopen simplex. Let f(0; a:) E A be the model output on input with parameter O, whose j-th coordinate fz æ)) stands for the probability æ belonging to class j. Let E(0; ær,y) y) be the cross entropy loss given input æ and label y, i.e. E(0; æ, y) log ()
Adding label noise means replacing the true label y with a fresh noisy label ý every time we access the sample. Specifically, ÿ is set the true label y with probability - p and any other label with probability -=1- where p is the fixed corruption probability. The training loss is defined as (()) N Σ (f m, 0:)] where the expectation is taken over the stochasticity of jt- Notice that given sample (x, y),
E[((O; æ. 0)] -(1 p) log (0(0
P
Σ log (0 3Ø
(114)
C-1
By the property of cross-entropy loss, 114) attains its global minimum if and only if f, نه for all j [C].j # y and fy = p. p. Due the large expressiveness of modern deep learning models, there typically exists a set S == (0 f((0) = B[ús), Vi E [N]] such that all elements of 8 minimizes C(0). Then, the manifold is a subset of S The following lemma relates the noise covariance Σ(0) = EIENN B[(Ve(() x,, 0:) (()) (VE(O; x,, 0c) - V2(0)) to the hessian V²)( for all 0 E S*
Lemma L.1. 1f 9.) is C² smooth on Rd given any i E [N] Os E [C] and S + 2, then for all 0 S" Σ(0) = L()).
Proof. Since L(-) is Cy-smooth, VL 0 for all 0 E S* To prove the above lemma, it suffices to show that Vi [N], E(VE(0 x, 9.) VI(0: o,, T1 = C() W.L.O.G, let y and therefore for all 0 E S*
[1 a)) = 1 p=:d1
z) = C-1 02 b L.jE
Additionally, let hxx)= == - log(2)) x E IRt. The stochastic gradient vf(0; 9)) follows the distri- bution:
04 Vi(0; x 1) = (a(a1)
w.p. 1-p.
00
(a(2))
5h
w.p. E > 1.
E
Then the covariance of the gradient noise is:
20 0) = (1 - (1(((()))) (0)
(0h(0')
00
00
((2)()2)) 2
(f,
(0')
C-1 Σ
+
00~
00
1>1
And the hessian is:
8h
phz(az)
fz
VCC(0) (1 (1-p)h((a))
+
002
C C-1 1
002
j>1
T
0f
(0f)
T
ph"
0f
())((0)
+(1-p) (a.)))
+
00
00
C
80
00
j>1
72

---

Published as a conference paper at ICLR 2023
Since ΣjE(o] fi =
811
0211
Σ
(115)
002
002
s>1
Also, notice that h'(x) = = Therefore,
ph
(1 (1-p) (a1)
(116)
C-1
Substituting (115) and (16 into the expression of T gives T = 0, which simplifies V²((0) as the
following form:
0f
(f)
ph" (22)
0f,
(ff(0)
V2L(0) = (1-r))"(aat
+
00 00
C
00
00
Again notice that (x) = (x)) for all x E RR Therefore, =
With the property Σ ((), we are ready to prove Theorem G.1
Proof of Theorem
G.1
Recall the general form of the slow SDE:
(t(t)
(() +
d (c) [E(() + (K dc)
(117)
VB
2B
where & is defined in Definition K.6 Since for Ç E E. = v²C(C), then
a (c) 0.
(118)
Now we show that
0 (< = -((r((
(119)
Since v2C(C) E((), VoPciC) Σ I. By Lemma K.4
8 (()(E)(( - -(0 v
Finally, we show that
{ = Vr (F((2HNV*C(())
(120)
2Hn
Define ((x) = rι((x) -11 By definition of (C(),, when Σ(c) V2C(C),
(C() ((2n))V²((()) where (()) is interpreted matrix function. Since ((2)HV=L(() span(uu' uu u E Té (T)). by Lemma K.4
8 (C) -
CLC)
By the chain rule, we have 120 Combining
118
119 and 120 gives the theorem.
73

---

Published as a conference paper at ICLR 2023
M EXPERIMENTAL DETAILS
In this section, we specify the experimental details that are omitted in the main text. Our experiments are conducted on CIFAR-I0 10 (Krizhevsky et al. 2009) and ImageNet Russakovsky et al. (2015). Our code is available at https: / //ggttttuuuu thu..b...... com hww.r 128 Loca SGD Our implementation of ResNet-56 (He et al. 2016) and VGG-16 (Simonyan an & Zisserman 2015) is based on the high- high-starred repository by Wei Yang and we use the implementation of ResNet-50 50 from torchvision 0.3.1. We
run all 10 experiments with Bloe = 128 on 8 NVIDIA Tesla P100 GPUs while ImageNet experiments are on 8 NVIDIA A5000 GPUS with Binc 32. All ImageNet experiments trained with ResNet-50.
We generally adopt the following training strategies. We do not add any momentum unless otherwise stated. We follow the suggestions by Jia et al. (2018) and do not add weight decay to the bias and learnable parameters in the normalization layers. For all models with BatchNorm layers, we go through 100 batches of data with batch size Bio to estimate the running mean and variance before evaluation. Experiments on both datasets follow the standard data augmentation pipeline in al. (201)) except the label noise experiments. Additionally, we use FFCV (Leclerc et al. 2022) to accelerate data loading for ImageNet training.
Slightly different from the update rule of Local SGD in Section I we use sampling without replace- ment unless otherwise stated. See Appendix [ for implementation details and discussion.
M.1 POST-LOCAL SGD EXPERIMENTS IN SECTION II
CIFAR-10 experiments. We simulate 32 clients with B = 4096. We follow the linear scaling rule and linear learning rate warmup strategy suggested by/ooyl et al. (2017). We first run 250 epochs of SGD with the learning rate gradually ramping up from 0.1 to 3.2 for the first 50 epochs. Resuming from the model obtained at epoch 250, we run Local SGD with 1) = 0.32. Note that we conduct grid search for the initial learning rate among (0.005, 0.01, 0.05, 0.1, 0.15, 0.2] and choose the learning rate with which parallel SGD (H = 1) achieves the best test accuracy. We also make sure that the optimal learning rate resides in the middle of the set. The weight decay À is set as 5 x 10-4, As for the initialization scheme, we follow Lin et al. (2020b) and Goyal et al. (2017). Specifically, we use Kaiming Normal (He et al.) 2013) for the weights of convolutional layers and initialize the weights of fully-connected layers by a Gaussian distribution with mean zero and standard deviation 0.01. The weights for normalization layers are initialized one. All bias parameters are initialized as zero. We report the mean and standard deviation over 5 runs.
ImageNet experiments. We simulate 256 workers with B = 8192. We follow the linear scaling rule and linear learning rate warmup strategy suggested by Goyal et al. (217)). We first run 100 epochs of SGD where the learning rate linearly ramps up from 0.5 to 16 for the first 5 epochs and then decays by a factor of 0.1 epoch 50. Resuming from epoch 100, we run Local SGD with n 0.16. Note that we conduct grid search for the initial learning rate among {0.05, 0.1, 0.5, 1} and choose the learning rate with which parallel SGD (Hf = 1) achieves the best test accuracy. We also make sure that the optimal learning rate resides in the middle of the set. The weight decay is set as 1 x 10-4 and we do not add any momentum. The initialization scheme follows the implementation of torchvision 0.3.1. We report the mean and standard deviation over runs.
M.2 EXPERIMENTAL DETAILS FOR FIGURES 2 AND 5
CIFAR-10 experiments. We use ResNet-56 for all CIFAR-10 experiments in the two figures. We
simulate 32 workers with B = 4096 and set the weight decay 5 x 10 For Figures 2nd)and ((b) we set 7 = 0.32, which is the same as the learning rate after decay in Figure [(a)
and
For Figure 2(a) we
adopt the same initialization scheme introduced in the corresponding paragraph in Appendix M.
For Figures 2(b) [ and [(c)), we use the model at epoch 250 in Figure ((a) as the pre-traines
model. Additionally, we use training budget of 250 epochs for Figure 2(e) In Figure we use Local SGD with momentum 0.9, where the momentum buffer is kept locally and never averaged.
We run SGD with momentum 0.9 for 150 epochs to obtain the pre-trained model, where the learning
74

---

Published as a conference paper at ICLR 2023
0.76
Test Accuracy
0.92
Accuracy
0.90
0.74
K St.
0.73
K 250. W -000
0.58
250
300
350
400
450
500
100
110
120
130
140
150
Epochs
Epochs
(a) CIFAR-10, start from #250.
(b) ImageNet, start from #100.
0.76
Accuracy
0.92
Test Accuracy
0.75
R 16. H., 24000
0.90
K B.g
20.00
0.74
Hay
Km Har 30.0
Mage
K 61. Bay 480
0.88
Hig.
0.73
R 250, Mi.
Hige
250
300
350
400
450
500
100
110
120
130
140
150
Epochs
Epochs
(c) CIFAR-10, start from #250, optimal I.
(d) ImageNet, start from #100, optimal H.
Figure 10: The learning curves for experiments in Figure
rate ramps up from 0.05 to 1.6 linearly in the first 150 epochs. Note that we conduct grid search for initial learning rate among {0.01, 0.05, 0.1, 0.15, 0.2] and choose the learning rate with which parallel SGD (H = 1) achieves the highest test accuracy. We also make sure that the optimal learning rate resides the middle of the set. Resuming from epoch 150, we run Local SGD H (i.e.. SGD) and 24 with 7 0.16 and decay 7) by 0.1 at epoch 226. For Local SGD H 900, resume from the model at epoch 226 of H 24 with 17 = 0.016. We report the mean and standard deviation over runs for Figures (a) 2 and [(c) and over 5 runs for Figure
ImageNet experiments. We simulate 256 clients with B = 8192 and set the weight decay as 1 10-1 In Figure 2(d) both Local SGD and SGD start from the same random initialization.
We warm the learning rate from 0.1 to 3.2 in the first 5 epochs and decay the learning rate by
factor of 0.1 epochs 50 and 100. For Figures [ 2 and ((d) we use the model at epoch 100
in Figure [(b]aa the pre-trained model. In Figure 2(c) we set the learning rate as 0.16, which is the
same as the learning rate after epoch 100 in Figure ((b) Finally, in Figures ((c) 2( 5(b) and (d)
we report the mean and average over runs.
M.3 DETAILS FOR EXPERIMENTS IN FIGURE 6
For all experiments in Figure 6 we train a ResNet-56 model on CIFAR-1O. We report mean test accuracy over three runs and the shaded area reflects the standard deviation. For Figure 6(a] use the same setup as Figures ((a) and for training from random initialization and from a apre-triieed model respectively except the learning rate. For Figure 6(b) we resume from the model obtained at epoch 250 in Figure (aa and train for another 250 epochs. For Figure 6(c) we follow the same
procedure Figure ((a) except that
sampling with replacement. We also ensure that the
total numbers of iterations in Figures ((a) and [(c) are the same.
M.4 DETAILS FOR EXPERIMENTS ON THE EFFECT OF THE DIFFUSION TERM
CIFAR-10 experiments. The model we use is ResNet-56. For Figure 3(a) we first run SGD with batch size 128 and learning rate 7 0.5 for 250 epochs to obtain the pre-trained model. The initialization scheme is the same as the corresponding paragraph Appendix M. Resuming from epoch 250 with 7 0.05, we run Local SGD with K 16 until epoch 6000 and run all other setups for the same number of iterations. We report the mean and standard deviation over 3 runs.
ImageNet experiments. For Figures 3(b) and 4(b) 100 in Figure [[b] In Figure ((b) we run Local SGD with K = 256 for another 150 epochs with n 0.032. We run all other setups for the number of iterations with the same learning rate.
start from the model obtained epoch
75

---

Published as a conference paper at ICLR 2023
M.5 DETAILS FOR EXPERIMENTS ON THE EFFECT OF GLOBAL BATCH SIZE
CIFAR-10 experiments. The model we use is ResNet-56. We resume from the model obtained in Figure [[a] epoch 250 and train for another 250 epochs. The local batch size for all runs is Biac 128. We first make grid search of v for SGD with K = 16 among (0.04, 0.08, 0.16, 0.32, 0.64 and find that the final test accuracy varies little across different learning rates (within 0.1%). Then we choose 77 0.32. For the green curve in Figure 4(a) we search for the optimal H for K = 16 and keep a fixed when scaling ?) with K. For the red curve Figure 4(a) we search for the optimal H for each K among {6, 12, 60, 120, 300, 750, 1500, 3000, 6000, 12000. 24000} and also make sure that H does not exceed the total number of iterations for 250 epochs. The learning curves for constant and optimal are visualized in Figures I0(a) and 10(c) respectively. We report the mean and standard deviation over three runs.
ImageNet experiments. We start from the model obtained at epoch 100 in ((b)ree[(bjaag train for another 50 epochs. The local batch size for all runs is Bise 32. We first make grid search among {0.032, 0.064, 0.16, 0.32] for H 1 to achieve the best test accuracy and choose H 0.064. For the orange curve in Figure 4(b) we search H among {2, 4, 6, 13, 26, 52. 78, 156} for K = 256 to achieve the optimal test accuracy and the keep a constant as we scale 17 with K. To obtain the optimal H for each K, we search among {6240. 7800, 10400, 12480, 15600, 20800, 24960, 31200] for K = 16, {1600, 3120, 4160, 5200, 6240, 7800, 10400) for K = 32, (312, 480, 520, 624, 800, 975, 1040, 1248, 1560, 1950} for K = 64, and (1, 2, 6. 13) for K" = 512. The learning curves for constant and optimal a are visualized in Figures [0(b) and [O(d) respectively. We report the mean and standard deviation over three runs.
M.6 DETAILS FOR EXPERIMENTS ON LABEL NOISE REGULARIZATION
For all label noise experiments, we do not use data augmentation, use sampling with replacement, and set the corruption probability as 0.1. 1. We simulate 32 workers with B 4096 in Figure and
4 workers with B = 512 in Figure X We use ResNet-56 with GroupNorm with the number of groups 8 for Figure 7aa and VGG- 16 without normalization for Figures 7(b) and Below we list the training details for ResNet-56 and VGG-16 respectively.
ResNet-56. As for the model architecture, we replace the batch normalization layer in Yang's implementation with group normalization such that the training loss is independent of the sampling order. We also use Swish activation (Ramachandran et al.] (201)) in place of ReLU to ensure the smoothness of the loss function. We generate the pre-trained model by running label noise SGD with corruption probability = 0.1 for 500 epochs (6, 000 iterations). We initialize the model by the same strategy introduced in the first paragraph of Appendix M. Applying the linear warmup scheme proposed by Gooal et al. (2017). we gradually ramp up the learning rate from 0.1 to 3.2 for the first 20 epochs and multiply the learning rate by 0.1 at epoch 250. All subsequent experiments in Figure 7(a) (a) use learning rate 0.1. The weight decay X is set as 5 x 10-1 Note that adding weight decay the presence of normalization accelerates the limiting dynamies and will not affect
the implicit regularization on the original loss function (Li al. 2022).
VGG-16. We follow Yang's 's implementation of the model architecture except that we replace max- imum pooling with average pooling and use Swish activation (Ramachandran et al. (017)) make the training loss smooth. We initialize all weight parameters by Kaiming Normal and all bias param- eters as zero. The pre-trained model is obtained by running label noise SGD with total batch size 4096 and corruption probability
0.1 for 6000 iterations. We use a linear learning rate warmup from 0.1 to 0.5 in the first 500 iterations. All runs in Figures 7 and g resume from the model
obtained by SGD with label noise. In Figure 7b) we use learning rate 7] 0.1. In Figure R we set 7 0.005 for H 97, 000 and y = 0.01 for SGD (H- 1). The weight decay À set as zero.
76


Reviews:
[{"review_id": "iGQyqZrxky9", "reviewer": "Reviewer_KHJA", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Local SGD (LSGD) is a communication-efficient variant of SGD for large-scale training, where multiple GPUs perform SGD independently and average the model parameters periodically. It has been shown previously that LSGD generalizes better when with a small enough learning rate and sufficient training time. To analyze this observation theoretically,  this paper introduces a new SDE that approximates LSGD in continuous time and compares it with the SDE of SGD. This comparison shows that LSGD has a stronger drift term which results in a stronger effect of regularizer and that leads to a faster reduction of sharpness. ", "summary_of_the_review": "comments \n1- Although remark 3.1 says that the analysis holds when SGDL starts out of zero loss manifold, however in the thm 3.2 it is needed that \\zeta(0) should be in the manifold. Where does this discrepancy come from? \n\n2- Since thm 3.2 needs that \\zeta(0) should be in the manifold, isn’t true that the presented analysis is more related to the Post LSGD? \n\n3- It would be more justified if you add some empirical results to the main body of the paper about your hypothesis 3.1. \n\nMinor comments: \nFor Fig 2, what is \\eta_1 and shouldn’t that be 0.32? Also, plots (e) and (f) don’t have proper legend i.e. not clear what are the values corresponding to each color. \n", "strength_and_weaknesses": "They provide novel and simpler ways to analyze the proposed SDE compared to existing theoretical results by adding a mild assumption. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "e5n7iuuR-lO", "reviewer": "Reviewer_QgEL", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "**Note**: Throughout the review, I will make a distinction between dense mini-batch SGD, i.e., mini-batch SGD with batch size $B$ and large mini-batch SGD with batch size $B\\cdot H$, where $H$ is the number of local steps, for the local-SGD algorithm with the exact computation and communication cost, and $B$ is the baseline batch-size used in this paper. Thus, dense mini-batch SGD is the same as parallel SGD in this paper's nomenclature.\n\n**Summary**: The paper studies the generalization behavior of Local SGD and its variant Post-local SGD ([Lin et al.](https://arxiv.org/pdf/1808.07217.pdf)). Specifically, it tries to understand and justify the better generalization performance of post-local SGD over dense mini-batch SGD. This problem is interesting, and while previous work explored this question, there is no satisfying, empirical, or theoretical explanation. \n\nAuthors first reproduce some experiments from previous works ([Lin et al.](https://arxiv.org/pdf/1808.07217.pdf), [Ortiz et al.](https://arxiv.org/abs/2110.08133?context=cs.CV)) showing that Post-local SGD, indeed has a generalization benefit over dense mini-batch SGD on CIFAR-10 and ImageNet. Then it compares local-SGD runs with different local steps but the same constant learning rate and shows that both local SGD and post-local SGD have a generalization benefit over dense mini-batch SGD if the learning rate is small enough. Further, the authors show that the optimal learning rate for local SGD depends on the number of local steps $H$. This observation is not surprising, given the optimization literature on local SGD in both convex and non-convex settings. Finally, the paper presents some theoretical results using the stochastic differential approximation of local SGD to justify the empirical behavior. \n", "summary_of_the_review": "The paper tries to answer an interesting question using new empirical and theoretical insights. However, I believe the comparisons between different optimizers are unfair due to incorrect hyperparameter tuning. As a result, the paper doesn't study what it claims to study. Some important baselines are also missing from the experiments. The theoretical result is not a standalone contribution but is used to explain the empirical findings. Since the finding seems less valuable, it is also unclear if the developed theory helps understand local SGD any better. More importantly, there is a mismatch between the settings of the experiments and the theory. The theoretical techniques might be of independent interest, but they don't constitute the paper's main contribution, the way it is written currently. Thus, I don't recommend accepting the work in its current form.  ", "strength_and_weaknesses": "#### **Unfair hyper-parameter tuning**\nThe paper almost treats the step size as a part of the problem and not a part of the algorithm/optimizer. This is why the authors make unfair comparisons and come up with questionable conclusions. For instance, why should one consider the same step size while changing the number of local steps in figure 1? It is well understood that at least in the convex setting, the optimal step size for local SGD inversely depends on the number of local steps $H$ (see the [optimal rate for local SGD]((https://arxiv.org/abs/2111.03741)) in the homogeneous setting in [this paper](https://arxiv.org/abs/2002.07839)). \n\nThe relationship might be different in the non-convex setting (as hinted by figure 2(e)-(f)). Thus the correct thing to do is to tune the step size separately for each $H$. Most of the experiments in [Lin et al.](https://arxiv.org/pdf/1808.07217.pdf) use the step-size schedule of dense/large mini-batch SGD for local SGD, thus giving mini-batch SGD the benefit of the doubt. The authors seem to be doing the opposite here, and I conjecture that dense/large mini-batch SGD will be better than the local SGD baselines in figure 2(a) if it uses the correct hyper-parameters. I also expect that $H=1$ might be optimal for larger step sizes in figures 2(e)-(f). \n\nIdeally, the authors must first fix a training budget and tune the step size to obtain the best validation loss/accuracy for the \"final model\" of each optimizer **separately**. This would make the optimal step size a function of $H$ and the training budget. Finally, it is important to clarify the stopping criterion for optimization; I couldn't find it in the appendix. \n\n\n#### **Missing comparison to Large Mini-batch SGD**\nThere is a lot of work trying to understand the optimization/generalization properties of local SGD in the convex setting. But the theory has been disappointing in showing the benefit of local SGD over large mini-batch SGD and single-machine SGD ([Woodworth et al.](https://arxiv.org/abs/2102.01583)). In fact, in the heterogeneous convex setting, which is the setting of the experiments, large mini-batch SGD is almost always better than local SGD [Woodworth et al.](https://arxiv.org/abs/2006.04735). Thus it is important to include large mini-batch SGD in all the experiments and use the correct step size for it (usual [scaling rules often fail](https://arxiv.org/abs/1811.03600)). The relevant setting, then, is the one where post-local SGD is better than both variants of mini-batch SGD. I expect that for small $H$, large-mini-batch SGD would be the correct algorithm to use, but for very large $H$ post-local SGD would be better. It is also important to compare against the SDE approximation of large mini-batch SGD in section 3. However, it might be challenging because the larger optimal step size might hurt the approximation error and make the comparison with the slow-SDE inconclusive. In particular, note that when $\\eta$ is very small local SGD is expected to behave as large mini-batch SGD. Can the authors comment on this theoretically from the SDE perspective?  \n\n\n#### **Inconsistent optimization setup**\nThe optimization setup introduced in section 1 and used throughout section 3 is that of stochastic optimization. However, all the experiments are performed using SGD with multiple passes, i.e., sampling without replacement, along with hyper-parameter tuning on the validation data set. Furthermore, the data is split between the machines for parallel training. While this is fine for the mini-batch algorithms, it introduces heterogeneity in client updates for local SGD; check [this](https://arxiv.org/abs/2006.04735) and [this](https://arxiv.org/abs/2206.04723). For a large number of local steps, this could cause a [client drift](http://proceedings.mlr.press/v130/charles21a.html) and have a regularization effect for local SGD. Several conflating factors make it hard to reconcile the experiments with each other and the theoretical guarantees. To make things worse, the writing doesn't acknowledge these nuances in optimization.    \n\n> Given a sufficiently small learning rate and a sufficiently long training time, Local SGD exhibits better generalization than SGD if the number of local steps H per round is tuned correctly according to the learning rate. This holds for both training from random initialization and from pre-trained models.\n\nIn light of the above comments, the above finding, which seems to be the paper's main contribution, is a bit vague. It is **not interesting** to identify which hyper-parameters make local SGD look good. Ideally, different optimizers should be compared after \"convergence,\" as the main goal is generalization here. Even if the training budget is fixed apriori and the authors want to consider the effect of different training budgets, they must change the step sizes accordingly. The interesting question is why post-local SGD generalizes better than mini-batch SGD, even if it potentially uses a non-optimal step size ([Lin et al.](https://arxiv.org/pdf/1808.07217.pdf)). The paper does not answer this question.\n\n#### **Highlight connections to the optimization literature**\nAs mentioned above, several connections to the optimization literature about local SGD must be mentioned in the main paper. This includes observations in figure 2(e)-(f), the client drift in the heterogeneous setting, min-max optimal algorithms in the homogeneous setting, etc. Note that many of these papers are generalization guarantees as they directly optimize on distribution $\\tilde{\\mathcal{D}}$, using stochastic first-order oracles and thus offer similar insights as in section 3. \n\n#### **Theoretical results**\nThe theoretical results and some of the developed techniques are novel. But the authors should improve the writing in section three, as the main result of the section doesn't come off clearly.\n\nI have already mentioned the missing comparison to large mini-batch SGD. I expect that for very small $\\eta = O(1/H)$, local-SGD and large mini-batch SGD should behave similarly, as the local gradients are all computed at very similar points (for context see [this paper](https://arxiv.org/pdf/1910.06378.pdf)). Thus, the interesting regime is where local SGD (or its approximation) can be shown to improve over both variants of mini-batch SGD. Can the current theoretical results even highlight such a regime because $\\eta = \\theta(1/H)$ in theorem 3.2? Also, what can the authors conclude about large mini-batch SGD, assuming hypothesis 3.1? \n\nThere are two levels of approximation in section 3.2, one more approximation leading up to (9), and these might not hold very well for the optimal step size for local SGD, which can be larger than $1/H$ (at least in the convex setting). Even If the optimal step size is smaller than $1/H$, then large mini-batch SGD might have comparable performance. On top of this, the generalization benefit would only hold in the regime where hypothesis 3.1 is correct. Thus it is a bit hard to make sense of the theoretical result here. The writing will benefit from a discussion of these different regimes. \n\nFinally, I'm not too fond of using the section to explain finding 2.1 (due to the above reasons). The theoretical results would be more interesting as a standalone contribution after more polishing.\n\n\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "TrBuHfgGW90", "reviewer": "Reviewer_ELvh", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper uses the derivation of an SDE to explain why and under what condition local SGD generalizes better than SGD. Empirical results are also provided to show that small learning rate + long enough training time + carefully tuned $H$ results in generalization improvement over SGD.", "summary_of_the_review": "This paper uses the derivation of an SDE to explain why and under what condition local SGD generalizes better than SGD. Empirical results are also provided to show that small learning rate + long enough training time + carefully tuned $H$ results in generalization improvement over SGD. I wonder if the batch size also plays an important role in the SDE analysis.", "strength_and_weaknesses": "Strength:\n1. This paper uses the derivation of an SDE to explain why and under what condition local SGD generalizes better than SGD. \n2. Empirical results are also provided to show that small learning rate + long enough training time + carefully tuned $H$ results in generalization improvement over SGD.\n3. The theoretical and empirical results provides some general guidelines and intuition of using post-local SGD, which is easy to follow.\n\nWeakness (concerns and questions):\n1. I recommend to show the detailed algorithms of SGD, local SGD, and post-local SGD in the paper, appendix is also fine, so that it will be easier for the readers to refer to these algorithms, especially for those who are not familiar with local SGD and post-local SGD. I believe post-local SGD is much less well-known to the general readers.\n2. The theoretical and empirical analysis focuses on learning rate, total number of steps, and $H$. I wonder if the SDE analysis could also show how the global batch sizes affects the generalization of local SGD.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster