Why (and When) does Local SGD Generalize Better than SGD?
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 [39], though the training regimes for this to happen are still in debate [47].
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 with a small learning rate, after approaching the manifold of 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 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.
1. Introduction
Recent advances suggest that the ultimate performance of deep learning on test sets can be drastically improved by scaling up the dataset and increasing the model size, but this requires more computation.
In response, recent works [15, 24, 65] seek to speed up standard training methods by exploiting data parallelism in a distributed computing setting and most works focus on improving
Stochastic Gradient Descent (SGD).
SGD tries to solve problems of the form min θ ∈ R d E ξ ∼ ˜ D
[ ℓ ( θ ; ξ )] ,
where θ ∈ R
d is the parameter vector of the model, ℓ ( θ ; ξ ) is the loss function for a data sample ξ drawn from the training distribution ˜ D , e.g., the uniform distribution over the training set.
SGD with learning rate η and batch size
B does the following update at each step, using a batch of B independent ξ t, 1 , . . .
, ξ t,B ∼ ˜ D :
[Equation]
where
[Equation]
Parallel SGD tries to improve efficiency by distributing the gradient computation to K ≥ 2 workers, each of whom focuses on a local batch of B loc :
= B/K samples and computes the average gradient over the local batch.
Finally, g t is obtained by averaging the local gradients over the K workers.
[Figure]
c
[Figure]
H
[Figure] [Figure]
[Figure] [Figure]
T o p -
[Figure]
T e s t
A cc u r a c y
H = 1 H =
[Figure]
H = 4 H =
[Figure]
[Figure]
In the ideal case where the batch size B is large enough, the speedup of parallel SGD over singlemachine SGD (in terms of wall-clock time) can be linear in the number of workers K .
However, large-batch training leads to a significant test accuracy drop compared to a small-batch training baseline with the same number of training steps or epochs [23, 27, 53, 55], which probably stems from a low level of gradient noise (see Appendix A for discussion).
A variant of SGD, called Local SGD [61, 67, 69] has been recently observed to help resolve the generalization degradation in certain regimes of large-batch training.
Perhaps surprisingly, Local
SGD is not designed for better generalization, but for reducing the high communication cost, another important issue that bottlenecks large-batch training [5, 50, 52, 58].
Instead of averaging the local gradients per step as in parallel SGD, Local SGD averages the local model parameters on the K workers whenever they finish H local steps of SGD (reducing the communication by a factor H ).
Formally, Local SGD proceeds in multiple rounds of model averaging, where each round produces a global iterate ¯ θ ( s ) .
In the ( s + 1) -th round, every worker k ∈
[ K ] starts with its local copy of the global iterate θ ( s ) k, 0 ← ¯ θ ( s ) and does H steps of SGD with local batches.
The k -th worker at its t -th local step draws a local batch of B loc := B/K independent samples ξ ( s ) k,t, 1 , . . .
, ξ ( s ) k,t,B loc from a shared training distribution ˜ D and updates as follows:
[Equation]
where
[Equation]
After finishing the H local steps, the workers aggregate the resulting local iterates θ ( s ) k,H and assign the average to the next global iterate:
[Equation]
Lin et al.
[39] discovered that Local SGD can be used as a strong component to improve generalization in large-batch training.
They proposed Post-local SGD , a hybrid method that starts with parallel SGD (equivalent to Local SGD with H = 1 in math) and switches to Local SGD with
H > 1 after a fixed number of steps t 0 .
Following a standard training procedure with momentum and multiple learning rate decays, they showed through extensive CIFAR-10 [31] experiments that
Post-local SGD significantly outperforms parallel SGD in test accuracy, where t 0 is carefully chosen to be the time of the first learning rate decay.
Ortiz et al.
[47] conducted a large-scale empirical study on ImageNet [51] but found that switching at the first learning rate decay can hurt generalization; instead, switching at a later time may improve validation accuracy.
Our experiments on
CIFAR-10 and ImageNet (Figure 1) reproduce the generalization improvement of Post-local SGD.
The success of Post-local SGD suggests that Local SGD induces a generalization benefit if the training starts from a model pre-trained by (parallel) SGD, though Local SGD is designed for a
[Figure]
different purpose.
The current paper tries to understand: Why does Local SGD generalize better?
Under what general conditions does this generalization benefit arise?
Previous theoretical research on Local SGD is mainly restricted to the convergence rate for minimizing a convex or non-convex objective (see Appendix A for a survey), and the generalization aspect of Local SGD is still unclear.
Our Contributions.
In this paper, we provide the first theoretical understanding on why (and when) switching from parallel SGD to Local SGD improves generalization.
We conduct ablation studies on CIFAR-10 and ImageNet and identify that small learning rate and sufficient training time are two important factors that contribute to the generalization improvement of Local SGD over the corresponding (parallel) SGD baseline.
Inspired by a previous analysis of the long-term generalization benefit of SGD [37], we derive an SDE for Local SGD that can track its long-term behavior with a novel proof.
We explain the generalization improvement of Local SGD over SGD through comparison with the corresponding SDEs: increasing the number of local steps H strengthens the drift term of SDE.
We then connect the stronger drift term to a stronger implicit regularization effect.
2. When does Local SGD Generalize Better?
It is still in debate under what general conditions this generalization benefit arises; see Appendix B. for a review.
In response, we focus on a simple yet insightful setup: the learning rate η is constant with time, and we train SGD and Local SGD without additional tricks (e.g., without momentum).
We observed the generalization improvement of Local SGD over SGD even in this setting (Figure 1).
We further conduct ablation studies and identify that small learning rate and sufficient training time are key factors that enable the improved generalization of Local SGD over its SGD counterpart.
Finding 2.1 Given a sufficiently small learning rate η and sufficiently long training time, Local
SGD exhibits better generalization than (parallel) SGD (with the same η, B loc , K ), if the number of local steps H is tuned properly according to η .
This holds for both training from random initialization and from SGD pre-trained models.
Our finding can help to settle the debate to a large extent; see Appendix B.3 for discussion.
Now we go through each point of our finding.
Pretraining is not necessary.
In contrast to previous works claiming the benefits of Post-local SGD over Local SGD [39, 47], we observe that
Local SGD with random initialization also generalizes significantly better than SGD, as long as η is small and the training time is sufficiently long (Figure 2( a )).
The learning rate should be small.
The learning rate is 0 .
32 for Figures 2( a ) 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 when η is large (e.g., η = 3 . 2 ), resulting in inferior final test accuracy.
Even for training from a pre-trained model, the generalization improvement of Local SGD disappears for large learning rates (e.g., η = 1 . 6 in Figure 3( d )).
In contrast, Figure 3( c ) shows that Local SGD achieves a comparable test accuracy with a much smaller learning rate η = 0 .
064 (with H and the training budget set properly).
Training time should be long enough.
In Figures 2( b ) and 2( c ), we extend the training budget for the Post-local SGD experiments in Section 1 and observe that a 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 3( a ) and 3( b ) for an enlarged view.
The number of local steps H should be tuned carefully.
The number of local steps H has a complex interplay with the learning rate η , but generally speaking, the test accuracy first rises as H increases, then it begins to fall when H is too large.
A smaller η needs a higher H to achieve consistent generalization improvement.
See Figures 2( e ) and 2( f ).
3. Theoretical Analysis of Local SGD: The Slow SDE
In this section, we use an SDE-based approach to establish the generalization benefit of Local SGD.
Below, we first identify the difficulty of adapting the conventional SDE framework to Local SGD.
Then, we present our new SDE and explain the generalization benefit.
Notations.
We follow the notations in Section 1.
We also define L ( θ ) :=
E ξ ∼ ˜ D
[ ℓ ( θ ; ξ )] as the expected loss, Σ ( θ ) :
=
Cov ξ ∼ ˜ D [ ∇ ℓ ( θ ; ξ )] as the noise covariance of gradients at θ .
Let { W t } t ≥denote the standard Wiener process.
For a mapping F : R d → R d , denote by ∂F ( θ ) the Jacobian at θ and ∂ 2 F ( θ ) the second order derivative at θ .
Furthermore, for any matrix M ∈ R d × d , ∂ 2 F ( θ )[ M ] = (cid:80) i ∈
[ d ] ⟨ ∂ 2 F i ∂ θ 2 , M ⟩ e i where e i is the i -th vector of the standard basis.
We write ∂ 2 ( ∇L )
( θ )
[ M ] as ∇ 3 L ( θ )[ M ] for short.
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 with the following SDE (3), which we call the conventional SDE approximation .
Below, we discuss why it cannot be directly adopted to characterize the behavior of Local SGD.
[Equation]
It is proved by Li et al.
[34] that this SDE is a first-order approximation to SGD, where each discrete step corresponds to a continuous time interval of η .
By 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 ( η − 1 ) discrete steps [34, 36].
However, this time scaling is too short for the difference to appear between Local
SGD and SGD.
We show in Theorem 9 that they closely track each other for O ( η − 1 ) steps.
See also
Appendix D for Lin et al.
[39]’s attempt to model Local SGD with multiple conventional SDEs and for our discussion on why it does not give much insight.
3.2. SDE Approximation near the Minimizer Manifold
Inspired by a recent paper [37], our strategy to overcome the shortcomings of the conventional
SDE is to design a new SDE that can guarantee a good approximation for O ( η − 2 ) discrete steps.
Following their setting, we assume the existence of a manifold Γ consisting only of local minimizers and track the global iterate ¯ θ ( s ) around Γ after it takes ˜ O ( η − 1 ) steps to approach Γ .
Assumption 3.1 The loss function L ( · ) and the matrix square root of the covariance Σ 1 / 2 ( · ) are
C ∞ -smooth.
Besides, we assume that ∥∇ ℓ ( θ ; ξ ) ∥ 2 is bounded by a constant for all θ and ξ .
Assumption 3.2 Γ is a C ∞ -smooth, ( d − m ) -dimensional submanifold of R d , where any ζ ∈ Γ is a local minimizer of L .
For all ζ ∈ Γ , rank( ∇ 2 L ( ζ ))
= m .
Additionally, there exists an open neighborhood of Γ , denoted as U , such that Γ = arg min θ ∈ U L ( θ ) .
Assumption 3.3 Γ is a compact manifold.
The existence of a minimizer manifold with rank( ∇ 2 L ( ζ ))
= m has also been made as a key assumption in Fehrman et al.
[11], Li et al.
[37], Lyu et al.
rank( ∇ 2 L ( ζ ))
=
m ensures that the Hessian is maximally non-degenerate on the manifold and implies that the tangent space at ζ ∈ Γ equals the null space of ∇ 2 L ( ζ ) .
Our SDE for Local SGD characterizes the training dynamics near Γ .
For ease of presentation, we define the following projection operators Φ , P ζ for points and differential forms.
Definition 1 (Gradient Flow Projection) Fix a point θ null / ∈ Γ .
For x ∈ R d , consider the gradient flow
d
x
( t ) d t = −∇L
( x ( t )) with x (0) =
x .
We denote the gradient flow projection of x as Φ( x ) .
Φ( x ) :=
lim t → + ∞ x ( t ) if the limit exists and belongs to Γ ; otherwise, Φ( x ) = θ null .
Definition 2 For any ζ ∈ Γ and any differential form A d W t + b d t in Itô calculus, where A is a matrix and b is a vector, we use P ζ ( A d W t + b d t ) as a shorthand for the differential form ∂ Φ( ζ ) A d W t + (cid:0) ∂ Φ( ζ ) b + 12 ∂ 2 Φ( ζ )[ AA ⊤ ] (cid:1) d t .
Here P ζ equals Φ( ζ + A d W t + b d t ) − Φ( ζ ) by Itô calculus, which means ζ + P ζ ( A d W t + b d t ) does not leave Γ .
It can be shown that ∂ Φ( ζ ) equals the projection matrix onto the tangent space of Γ at ζ .
We decompose the noise covariance Σ ( ζ ) for ζ ∈ Γ into the tangent space part Σ ∥ ( ζ ) :
= ∂
Φ( ζ ) Σ ( ζ ) ∂ Φ( ζ ) and the rest part Σ ♢ ( ζ ) :
= Σ ( ζ ) − Σ ∥ ( ζ ) .
Now we are ready to state our SDE.
Definition 3 (Slow SDE for Local SGD)
Given η, H > 0 and ζ 0 ∈ Γ , define ζ ( t ) as the solution of the following SDE with initial condition ζ (0) = ζ 0 :
[Equation]
Here (cid:98) Σ ♢ ( ζ ) , (cid:98) Ψ ( ζ )
∈ R
d × d are defined as
[Equation]
where { v i } di =1 is a set of eigenvectors of ∇ 2 L ( ζ ) that forms an orthonormal eigenbasis, λ 1 , . .
, λ d are the corresponding eigenvalues.
Additionally, ψ (0) = 0 and ψ ( x ) :
= e − x − 1+ x x for x ̸ = 0 .
The use of P ζ keeps ζ ( t ) on Γ through projection.
Σ 1 / 2 ∥ ( ζ ) introduces a diffusion term to the
SDE in the tangent space.
The two drift terms involve (cid:98) Σ ♢ ( · ) and (cid:98) Ψ ( · ) , which can be intuitively understood as rescaling the entries of the noise covariance in the eigenbasis of Hessian.
In the special case where ∇ 2 L = diag( λ 1 , · · · , λ d ) ∈ R d × d , we have (cid:98) Σ ♢ ,i,j =
[Equation]
ψ ( ηH ( λ i + λ j ))
λ
i + λ j Σ 0 ,i,j .
ψ ( x ) increases from 0 to 1 as x goes from 0 to infinity (see Figure 7)
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 η 2 instead of an interval of η in the conventional SDE.
This Slow SDE is inspired by Li et al.
Under nearly the same set of assumptions, they proved that SGD can be tracked by an SDE that is essentially equivalent to (4) with K = 1 , namely, without the drift-II term:
[Equation]
We refer to (7) as the Slow SDE for SGD .
We remark that the drfit-II term in (4) is novel and is 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.
Theorem 4 Let Assumptions 3.1 to 3.3 hold.
Let T > 0 be a constant and ζ ( t ) be the solution to (4) with the initial condition ζ (0) =
Φ( ¯ θ (0) ) ∈ Γ .
If H is set to αη for some constant α > 0 , then for any C 3 -smooth function g ( θ ) ,
[Equation]
where ˜ O ( · ) hides log factors and constants that are independent of η but can depend on g ( θ ) .
Theorem 5 For δ
= O (poly( η )) , with probability at least 1 − δ , it holds for all O ( 1 α log 1 η ) ≤ s ≤ Tαη that Φ( ¯ θ ( s ) )
∈ Γ and ∥ ¯ θ ( s ) − Φ( ¯ θ ( s ) )
∥ 2 =
O ( (cid:113) αη log αηδ ) , where O ( · ) hides constants independent of η , α and δ .
Theorem 4 suggests that the trajectories of the manifold projection and the solution to the Slow
SDE (4) are close to each other in the weak approximation sense.
Theorem 5 further states that the iterate ¯ θ ( s ) keeps close to its manifold projection after the first few rounds.
Remark 6 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) can converge to Γ ); (2) give better bounds for smaller η ; (3) characterize a long training horizon ∼ η − 2 .
The need for tuning H will be discussed in Remark 8.
3.3. Interpretation of the Slow SDEs
In this subsection, we compare the Slow SDEs for SGD and Local SGD and provide an 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.
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; the latter one drives the dynamics to move along the negative gradient of 12 B ⟨∇ 2 L ( ζ ) , (cid:98) Σ ♢ ( ζ ) ⟩ projected onto the tangent space, but ignoring the dependency of (cid:98) Σ ♢ ( ζ ) on ζ .
This can be connected to the class of semi-gradient methods which only computes a part of the gradient [4, 45, 60].
In this view, the long-term behavior of SGD is similar to a stochastic semi-gradient method minimizing the implicit regularizer [Equation] , (cid:98) Σ ♢ ( ζ ) ⟩ on Γ .
This argument reveals that SGD has a deterministic trend toward the region with a smaller magnitude of Hessian, which is commonly believed to correlate with better generalization [19, 25, 27, 46].
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 ( 1 B , 12 B ) -Slow SDE corresponds to the Slow SDE for SGD (7).
Definition 7 For κ 1 , κ 2 ≥ 0 , define ( κ 1 , κ 2 ) -Slow SDE to be the following:
[Equation]
Hypothesis 3.1 Starting at a minimizer ζ 0 ∈ Γ , run ( κ 1 , κ 2 ) -Slow SDE and ( κ 1 , κ ′ 2 ) -Slow SDE respectively for the same amount of time T > 0 and obtain ζ ( T ) , ζ ′ ( T ) .
If κ 2 > κ ′ 2 , then the expected test accuracy at ζ ( T ) is better than that at ζ ′ ( T ) .
One motivating example is SGD with label noise regularization.
In this case, the Slow SDE for
SGD turns out as a simple gradient flow on Γ aimed at minimizing tr( ∇ 2 L ) , and larger drift term means flatter minima; see Appendix E. 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).
Local SGD Strengthens the Drift Term in Slow SDE.
Our hypothesis can help us understand why Local SGD generalizes better than SGD.
The Slow SDE for Local SGD (4) has an additional drfit-II term.
Similarly, this term has the effect of driving the dynamics to move along the negative semi-gradient of K − 1 2 B ⟨∇ 2 L ( ζ ) , (cid:98) Ψ ( ζ ) ⟩ (with the dependency of (cid:98) Ψ ( ζ ) 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 a stochastic semi-gradient method minimizing the implicit regularizer
2 L ( ζ ) , (cid:98) Σ
♢ ( ζ ) ⟩ + K − 1 2 B ⟨∇ 2 L ( ζ ) , (cid:98) Ψ ( ζ ) ⟩ on the minimizer manifold of the original loss L .
(cid:98) Ψ ( ζ ) is basically a rescaling of the entries of (cid:98) Σ ⋄ ( ζ ) in the eigenbasis of Hessian, where the rescaling factor ψ ( ηH · ( λ i + λ j )) is between 0 and 1 (see Figure 7 for the plot of ψ ).
When ηH is small, the rescaling factors should be close to 0 , then (cid:98) Ψ ( ζ ) ≈ 0 , leading to almost no additional regularization.
On the other hand, when ηH is large, the rescaling factors should be close to 1 , then (cid:98) Ψ ( ζ ) ≈ (cid:98) Σ ⋄ ( ζ ) .
We can then merge the two implicit regularizers to be K 2 B ⟨∇ 2 L ( ζ ) , (cid:98) Σ
♢ ( ζ ) ⟩ .
The corresponding Slow SDE is approximately the ( 1 B , K 2 B ) -Slow SDE, which is restated below:
[Equation]
Comparing the ( 1 B , 12 B ) -Slow SDE for SGD (7) and the ( 1 B , K 2 B ) -Slow SDE above (9), the difference is that the drift term is amplified by K times.
According to our hypothesis, we can then attribute the generalization improvement of Local SGD to the amplified drift term.
Since the Slow SDE of Local SGD is better approximated by the ( 1 B , K 2 B ) -Slow SDE above (9) when Hη increases, our hypothesis also implies that we should increase H when decreasing η to experience the similar generalization benefit, which is consistent with our empirical finding in
Section 2 that the optimal H for generalization increases as η decreases (Figures 2( e ) and 2( f )).
Remark 8 When η is small but finite, tuning H offers a trade-off between regularization strength and SDE approximation quality.
Larger α := ηH makes the regularization stronger in the SDE, but the SDE itself may lose track of Local SGD, which can be seen from the error bound in Theorem 5.
This matches our finding that tuning H is important for better generalization (Finding 2.1).
Prediction: Increasing the number of workers helps generalization.
In addition to strengthening the drift term of the Slow SDE for SGD, another way to help the corresponding semi-gradient method to optimize the implicit regularizer is to reduce the diffusion term .
We conduct experiments where we keep H and η fixed and gradually increase the number of workers K to reduce the diffusion term.
As shown in Figure 6, a higher test accuracy is achieved for larger K .
4. Conclusions
In this paper, we provide a theoretical analysis for Local SGD that captures its long-term generalization benefit in the small learning rate regime.
We derive the Slow SDE for Local SGD as a generalization of the Slow SDE for SGD [37], and attribute the generalization improvement over
SGD to the larger drift term in the SDE for Local SGD.
Our empirical validation show 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 terms of test accuracy, which requires a much deeper understanding of loss landscape and the Slow SDEs and is left for future work.
Another future work direction is to design a distributed training method that provably generalizes better than SGD based on the insights from Slow SDEs
Proof Outline . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
J Proof Details of Main Theorems
J.
Additional Notations . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
J.
Computing the Derivatives of the Limiting Mapping . . . . . . . . . . . . . . . . .
J.
Preliminary Lemmas for GD and GF . . . . . . . . . . . . . . . . . . . . . . . . .
J.
Construction of working zones . . . . . . . . . . . . . . . . . . . . . . . . . . . .
J.
High Probability Bounds for Phase 1: Iterate Approaching the Manifold . . . . . .
J.5.
Additional notations . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
J.5.
Proof for Subphase 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
J.5.
Proof for Subphase 2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
J.
High Probability Bounds for Phase 2: Iterates Staying Close to Manifold . . . . . .
J.6.
Additional notations . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
J.6.
Proof for the High Probability Bounds . . . . . . . . . . . . . . . . . . . .
J.
Summary of High Probability Bounds . . . . . . . . . . . . . . . . . . . . . . . .
J.
Proof of Theorem 5 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
J.
Computing the Moments for Phase 2 . . . . . . . . . . . . . . . . . . . . . . . . .
J.10 Proof of Weak Approximation . . . . . . . . . . . . . . . . . . . . . . . . . . . .
J.10.1 Preliminaries and additional notations . . . . . . . . . . . . . . . . . . . . {+J.10.
+}
Proof of the approximation in our context . . . . . . . . . . . . . . . . . .
K Deriving the Slow SDE for Label Noise Regularization
L Experimental Details
L. Post-local SGD Experiments in Section 1 . . . . . . . . . . . . . . . . . . . . . .
L.
Experimental Details for Sections 2 and B.2 . . . . . . . . . . . . . . . . . . . . . L.
Experiments on Reducing the Diffusion Term . . . . . . . . . . . . . . . . . . . . L.
Local SGD with Label Noise Regularization . .
Appendix A. Additional Related Works
Optimization Aspect of Local SGD.
Local SGD is a communication-efficient variant of parallel
SGD, where multiple workers perform SGD independently and average the model parameters periodically.
Dating back to Mann et al.
[43] and Zinkevich et al.
[71], this strategy has been widely adopted to reduce the communication cost and speed up training in both scenarios of data center distributed training [6, 48, 59, 68] and Federated Learning [26, 44].
To further accelerate training, Wang and Joshi [61] and Haddadpour et al.
[16] proposed adaptive schemes for the averaging frequency, and Basu et al.
[1] combined Local SGD with gradient compression.
Motivated to theoretically understand the empirical success of Local SGD, a lot of researchers [14, 28, 57, 62, 67] analyzed the convergence rate of Local SGD under various settings, e.g., convex/non-convex objective functions.
The error bound of Local SGD obtained by these works is typically inferior to that of SGD 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 are interested in the generalization aspect of Local SGD, assuming the training loss can be optimized to a small value.
Gradient noise and generalization.
The effect of stochastic gradient noise on generalization has been studied from different aspects, e.g., changing the order of learning different patterns Li et al.
[34], inducing an implicit regularizer in the second-order SDE approximation Li et al.
Smith et al.
Our work follows a line of works studying the effect of noise in the lens of sharpness, which is long believed to be related to generalization Hochreiter and Schmidhuber [19],
Neyshabur et al.
Keskar et al.
[27] empirically observed that large-batch training leads to worse generalization and sharper minima than small-batch training.
Hu et al.
[21], Ma and Ying [41], Wu et al.
[63] showed that gradient noise destabilizes the training around sharp minima, and Ibayashi and Imaizumi [22], Kleinberg et al.
[29], Xie et al.
[64], Zhu et al.
[70] quantitatively characterized how SGD escapes sharp minima.
The most related papers are Blanc et al.
[3], Damian et al.
[7], Li et al.
[37], 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 [8], sharpness still appears to be one of the most promising tools for predicting generalization [13, 25].
Improving generalization in large-batch training.
The generalization issue of the large-batch (or full-batch) training has been observed as early as [2, 33].
As mentioned in Section 1, the generalization issue of large-batch training could be due to the lack of a sufficient amount of stochastic noise.
To make up the noise in large-batch training, Goyal et al.
[15], Krizhevsky [30] empirically discovered the Linear Scaling Rule for SGD, which suggests enlarging the learning rate proportionally to the batch size.
Jastrz˛ebski et al.
[23] adopted an SDE-based analysis to justify that this scaling rule indeed retains the same amount of noise as small-batch training (see also Section 3.1).
However, the SDE approximation may fail if the learning rate is too large [36], especially in the early phase of training before the first learning rate decay [55].
Shallue et al.
[53] demonstrated that generalization gap between small- and large-batch training can also depend on many other training hyperparameters.
Besides enlarging learning rate, other approaches have also been proposed to reduce the gap, including training longer [20], learning rate warmup [15], LARS [65], LAMB [66].
In this paper, we focus on using Local SGD to improve generalization, but adding local steps is a generic training trick that can also be combined with others, e.g., Local LARS [39], Local ExtrapSGD [38].
Appendix B. Supplementary for SectionB.1. The Debate on Local SGD
In this section, we summarize a debate in the literature regarding when to switch the training mode in
Post-local SGD.
As Post-local SGD can be viewed as running Local SGD from an SGD-pretrained model, the discussion around the time point for switching can reveal some information about the conditions for Local SGD to generalize better.
Local SGD generalizes better than SGD on CIFAR-10.
Lin et al.
[39] empirically observed that
Post-local SGD exhibits a better generalization performance than SGD.
Most of their experiments are conducted on CIFAR-10 and CIFAR-100 with multiple learning rate decay, and the algorithm switches from (parallel) SGD to Local SGD right after the first learning rate decay.
We refer to this particular 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 decay.
It has also been observed by Wang and Joshi [62] 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.
[38] showed that adding local steps to Extrap-SGD, a variant of SGD proposed therein, after the first learning rate decay also improves generalization, suggesting that the firstdecay 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 beyond CIFAR-10, Lin et al.
[39] 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, Lin et al.
[39] suggested that the time of switching should be tuned aiming at “capturing the time when trajectory starts to get into the influence basin of a local minimum” in a footnote, but no further discussion or experiments are provided to justify this guideline.
Ortiz et al.
[47] 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.
1 To explain this phenomenon, they conjecture that switching to Local SGD has a regularization effect that is beneficial only in the short-term, so it is 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 TinyImageNet, which is a spatially downscaled subset of ImageNet, the first-decay switching strategy indeed leads to better validation accuracy.
B.2. Additional Experimental Results for Section
In Figure 3, we present additional experimental results for Section 2 to further verify our finding.
Specifically, in Figure 3, (a) and (b) are enlarged views for Figure 2 (a) and (c) respectively, showing that Local SGD can generalize worse than SGD in the first few epochs.
(c) shows that Local SGD can still chieve comparable test accuracy when we use a much smaller learning rate (e.g., η = 0 . 064 ) on the condition that H and the training budget are set properly.
(d) presents the case where, with a large learning rate, the generalization improvement of Local SGD disappears even starting from a pre-trained model.
In (e), the generalization benefit of Local SGD with H = 24 becomes less significant after the learning rate decay at epoch 226 , which is consistent with the observation by
Ortiz et al.
[47] that the generalization benefit of Local SGD usually disappears after the learning rate decay.
But we can preserve the improvement by increasing H to 900 .
Here
[Figure]
we use Local SGD with momentum.
We refer the readers to Appendix L.2 for training details.
[Figure]
B.3. Reconciling previous works
Our finding can help to settle the debate presented in Appendix B.1 to a large extent.
Simultaneously requiring a small learning rate and sufficient training time poses a trade-off when learning rate decay is used with a limited training budget: switching to Local SGD earlier may lead to a large learning rate, while switching later may result in insufficient training time.
It is thus unsurprising that firstdecay switching strategy is not always the best when the dataset and learning rate schedule change.
The need for sufficient training time does not contradict with Ortiz et al.
[47]’s conjecture that
Local SGD only has a “short-term” generalization benefit.
In their experiments, the generalization improvement usually disappears right after the next learning rate decay (instead of after a 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 to tune H after η changes.
In
Figure 3 (e), we reproduce this phenomenon and show that increasing H after learning rate decay retains the improvement.
Appendix C. Accuracy and Loss on the Training Set
This section visualizes the accuracy and loss on the training set for CIFAR-10 experiments in Figures 1, 2 and 3.
The accuracy and loss on the training set are computed on the averaged parameters. In each evaluation [Figure] we first randomly sample 100 batches of training data (batch size = 128 ) without replacement and go through them to estimate the running mean and variance.
Then we compute accuracy and loss on the whole training set.
We omit the plots for training accuracy and loss for
ImageNet experiments since it is computationally expensive to go through the whole training set of
ImageNet.
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
i n i
[Figure]
g
[Figure]
t [Figure] Appendix D. Modeling Local SGD with Multiple Conventional SDEs Several previous works adopt the conventional SDE approximation (3) and connect good generalization to having a large diffusion term (cid:113) ηB Σ 1 / 2 d W t in the SDE [23, 55], because a suitable amount of noise can be necessary for large-batch training to generalize well (see also Appendix A).
Lin et al.
[39] tried to informally explain the success of Local SGD by adopting this argument.
Basically, they attempted to write multiple SDEs, each of which describes the H -step local training process of each worker in each round (from θ ( s ) k, 0 to θ ( s ) k,H ).
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 B loc instead of B :
[Equation]
Lin et al.
[39] then argue that the total amount of “noise” in the training dynamics of Local SGD is larger than that of SGD.
However, it is hard to see whether it 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, a complete modeling of Local SGD following this idea should view the sequence of global iterates { ¯ θ ( s ) } as a Markov process { X ( s ) } .
Let P X
( x , B, t ) the distribution of
X ( t ) in (3) with initial condition X (0)
=
x .
Then the Markov transition should be X ( s +1)
= 1 K (cid:80)
Kk =1 X ( s ) k,H , where X ( s ) 1 ,H , . . .
, X ( s ) K,H are K independent samples from P X
( X ( s ) , B loc , Hη ) , i.e., sampling from (10).
Consider one round of model averaging.
It is true that P X ( X ( s ) , B loc , Hη ) may have a larger variance than the corresponding SGD baseline P X ( X ( s ) , B, Hη ) because the former one has a smaller batch size.
However, it is unclear whether X ( s +1) also has a larger variance than P X
( X ( s ) , B, Hη ) .
This is because X ( s +1) is the average of K samples, which means we have to compare 1 K times the variance of P X ( X ( s ) , B loc , Hη ) with the variance of P X ( X ( s ) , B, Hη ) .
Then it is unclear which one is larger.
In the special case where Hη is small, P X ( X ( s ) , B loc , Hη ) is approximately equal to the following Gaussian distribution:
[Equation]
Then averaging over K samples gives
[Equation]
which is exactly the same as the Gaussian approximation of the SGD baseline.
This means there do exist certain cases where Lin et al.
[39]’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.
Appendix E. The Slow SDE for Training with Label Noise Regularization
In this section, we study the case of training over-parameterized neural nets with label noise regularization to exemplify the generalization benefit of having a larger drift term [3, 7, 37].
For a C -class classification task, the label noise regularization is as follows: every time we draw a sample from the training set, we make the true label as it is with probability 1 − p , and replace it with any other label with equal probability pC − 1 .
When we use cross-entropy loss, the Slow SDE for
SGD turns out to be a simple deterministic gradient flow on Γ (instead of a semi-gradient method) for minimizing the trace of Hessian:
[Equation]
where ∇ Γ f stands for the gradient of the function f projected to the tangent space of Γ .
Checking the validity of Hypothesis 3.1 reduces to the following question: Is minimizing the trace of Hessian beneficial to generalization?
Many previous works provide positive answers, including the line of works we just mentioned.
Blanc et al.
[3] and Li et al.
[37] connect minimizing the trace of Hessian to finding sparse or low-rank solutions for training two-layer linear nets.
Damian et al.
[7] empirically showed that good generalization correlates with a smaller trace of Hessian in training ResNets with label noise.
Besides, Ma and Ying [41] connects the trace of Hessian to the smoothness of the function represented by a deep neural net.
We defer all the proofs to Appendix K.
As for Local SGD, the Slow SDE can be simplified as:
[Equation]
[Equation]
where F ( x ) := (cid:82) x 0 ψ ( y )
d y and is interpreted as a matrix function in (14).
Note that the magnitude of the RHS in (14) becomes larger as H increases.
When H gets large enough, the RHS in (14) is approximately K times of the RHS in (13): [Equation]
[Equation]
Comparing (13) and (15), we can conclude that Local SGD accelerates the process of sharpness reduction, thus leading to better generalization.
Furthermore, the regularization effect gets stronger for larger H .
We also conduct experiments on non-augmented CIFAR-10 with label noise regularization to verify our conclusion.
As shown in Figure 5, adding local steps indeed gives better generalization performance.
[Figure]
Appendix F. Experimental Results on Reducing the Diffusion Term
We conduct experiments on CIFAR-10 with varying K and fixed η , H to justify that reducing the diffusion term is beneficial to generalization.
As shown in Figure 6, we can achieve higher test accuracy for larger K where the diffusion term is smaller.
[Figure]
Appendix G. Discussion
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 a 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 are associated with the diffusion term in the conventional SDE (Slow SDEs become stationary if Σ = 0 ).
Our view identifies the effect of each different component of the noise covariance on generalization, and therefore, goes one step further on the conventional wisdom.
Can Local SGD close the generalization gap between small- and large-batch training?
We remark that the mechanism causing the generalization gap is different from the mechanism of Local
SGD for improving generalization, so it is unclear what a general claim can be made.
As noted by
Smith et al.
[55], the generalization gap occurs early in the training when the learning rate is large and SDE does not give a good approximation.
In contrast, as shown in our paper, Local SGD yields generalization benefits mainly for small learning rate in later phases of training.
Appendix H. Local SGD stays close to SGD for O ( η − 1 ) steps
The following theorem states that Local SGD and SGD closely track each other for O ( η − 1 ) steps.
It suggests that this time horizon is too short for the difference between the two algorithm to appear, necessitating an analysis that lasts for a longer time.
Theorem 9 Assume that the loss function L is C 3 -smooth with bounded second and third order derivatives and that ∇ ℓ ( θ ; ξ ) is bounded.
Let T > 0 be a constant, ¯ θ ( s ) be the s -th global iterate of Local SGD and w t be the t -th iterate of SGD with the same initialization w
= ¯
θ (0) and same η, B loc , K .
Then for any H ≤ Tη and δ
= O (poly( η )) , it holds with probability at least 1 − δ that for all s ≤ [-TηH , ∥ ¯ θ ( s ) − w sH ∥ 
-]T [Equation]
[Equation]
H.1. Proof of Theorem
This subsection presents the proof for Theorem 9.
First, we define some notations that will be used throughout this section.
For the sequence of Local SGD iterates { θ ( s ) k,t : k ∈
[ K ] , 0 ≤ t ≤
H, s ≥ 0 } , we introduce an auxiliary sequence { ˆ u t } t ∈ N which starts from ¯ θ (0) and follows.
Define ˆ ∆ ( s ) k,t := θ ( s ) k,t − ˆ u ( s ) t and ¯ ∆ ( s ) :
= ¯ θ ( s ) − ˆ u ( s ) 0 , which stand for the difference between the Local
SGD iterate and GD iterate.
For convenience, let ˆ u ( s ) t
= ˆ u sH + t and z k,sH + t =
z ( s ) k,t , which will be used interchangeably.
Assume that L is C 3 -smooth with bounded second and third order derivatives.
Let ν 2 :=
sup θ ∈ R
d ∥∇ 2 L ( θ ) ∥ 2 and ν 3 :
= sup θ
∈ R
d ∥∇ 3 L ( θ ) ∥ 2 .
Since ∇ ℓ ( θ ; ζ ) is bounded, the gradient noise z ( s ) k,t is also bounded and we denote by σ max the upper bound such that ∥ z ( s ) k,t ∥ 2 ≤ σ max , ∀ s, k, t .
For each client, define the following sequence { ˆ Z k,t : t ≥ 0 } :
[Equation]
To prove Theorem 9, we will show that both Local SGD iterates ¯ θ ( s ) and SGD iterates w sH will closely track GD iterates ˆ u sH with high probability.
The following lemma establishes the concentration property of ˆ Z k,t .
Lemma 10 (Concentration property of { ˆ Z k,t } )
With probability at least 1 − δ , the following holds simultaneously for all k ∈
[ K ] , 0 ≤ t <
⌊ Tη ⌋ :
[Equation]
where ˆ C 1 := exp( T ν 2 ) .
Proof For each ˆ Z k,t , construct a sequence { ˆ Z k,t,t ′ } tt ′ [-=
-]=0 :
[Equation]
Since ∥∇ 2 L (ˆ u l ) ∥ 2 ≤ ν 2 for all l ≥ 0 , the following holds for all 0 ≤ τ < t − 1 and 0 < t <
⌊ Tη ⌋ :
[Equation]
Notice that { ˆ Z k,t,t ′ } tt ′ =0 is a martingale with ∥
ˆ Z k,t,t ′ − ˆ Z k,t,t ′ − 1 ∥ 2 ≤ ˆ C 1 σ max .
Since ˆ Z k,t = ˆ Z k,t,t , by Azuma-Hoeffding’s inequality,
[Equation]
Taking union bound on all k ∈
[ K ] and 0 ≤ t ≤ ⌊ Tη ⌋ , we can conclude that with probability at least 1 − δ ,
[Equation]
The following lemma states that, with high probability, Local SGD iterates θ ( s ) k,t and ¯ θ ( s ) closely track the gradient descent iterates ˆ u sH for ⌊ THη ⌋ rounds.
Lemma
For δ =
O (poly( η )) , the following inequalities hold with probability at least 1 − δ :
[Equation]
and
[Equation]
where ˆ C 3 is a constant independent of η and H .
Proof According to the update rule for θ ( s ) k,t and ˆ u ( s ) t ,
[Equation]
[Equation]
Subtracting (16) from (17) gives
[Equation]
[Equation]
where ∥ ˆ v ( s ) k,t ∥ 2 ≤ ρ 3 2 ∥ ˆ ∆ ( s ) k,t ∥ 22 .
Applying (18) t times
[Equation]
we have
[Equation]
Here, T can expressed in the following form:
[Equation]
Substituting in t =
H and taking the average
[Equation]
we derive the following recursion:
[Equation]
Applying (20) s times yields
[Equation]
Substitute (21) into (19) and we have
[Equation]
By Cauchy-Schwartz inequality and triangle inequality [Equation] we have
[Equation]
where ˆ C 1 = exp( ν 2 T ) .
Below we prove by induction that for δ
= O (poly( η )) , if
[Equation]
then there exists a constant ˆ C 2 such that for all k ∈
[ K ] , 0 ≤ s < ⌊ TηH ⌋ and 0 ≤ t ≤ H ,
[Equation]
First, for all k ∈ [ K ] , ∥ ˆ ∆ (0) k, 0 ∥ 2 = 0 and hence (24) holds.
Assuming that (24) holds for all ˆ ∆ ( r ) k ′ ,τ where k ′ ∈
[ K ] , 0 ≤ r < s , 0 ≤ τ ≤ H and r = s , 0 ≤ τ < t , then by (22), then for all k ∈
[ K ] , the following holds:
[Equation]
Let ˆ C 2 ≥ 6 ˜ C 21 σ max √ 2 T .
Then for sufficiently small η , (24) holds.
By Theorem 10, (23) holds with probability at least 1 − δ .
Furthermore, notice that ¯ θ ( s ) − ˆ u sH =
1 K (cid:80) k ∈
[ K ] ˆ ∆ ( s − 1) k,H .
Hence we have the lemma.
The iterates of standard SGD can be viewed as the local iterates on a single client with the number of local steps ⌊ Tη ⌋ .
Therefore, we can directly apply Theorem 11 and obtain the following lemma about the SGD iterates w t .
Corollary 12 For δ
=
O (poly( η )) , the following holds with probability at least 1 − δ :
[Equation]
where ˆ C 3 is the same constant as in Theorem 11.
Applying Theorem 11 and Theorem 12 and taking union bound, we have Theorem 9.
Appendix I. Proof Outline of Main Theorems
In this section, we provide the proof outline of Theorem 4.
The proof details are deferred to Appendix J. We first introduce additional notations that will be used throughout Appendix I and Appendix J.
I.1. Additional Notations and Definitions
We first introduce the notion of µ -PL.
Definition 13 (Polyak-Łojasiewicz Condition)
For µ > 0 , we say a function L ( · ) satisfies µPolyak-Łojasiewicz condition (abbreviated as µ -PL) on set U if
[Equation]
We will later show that there exists a neighborhood of Γ where L satisfies µ -PL.
We then introduce the definitions of the ϵ -ball at a point and the ϵ -neighborhood of a set.
For θ ∈ R d and ϵ > 0 , B ϵ ( θ ) := { θ ′ | ∥ θ ′ − θ ∥ 2 < ϵ } is the open ϵ -ball centered at θ .
For a set
Z ⊆ R d , M ϵ :
= (cid:83) θ ∈M B ϵ ( θ ) is the ϵ -neighborhood of M .
I.2. Construction of working zones
We construct four nested working zones (Γ ϵ 0 , Γ ϵ 1 , Γ ϵ 2 , Γ ϵ 3 ) in the neighborhood of Γ .
Later we will show that the local iterates θ ( s ) k,t ∈ Γ ϵ 2 and the global iterates ¯ θ ( s ) ∈ Γ ϵ 0 with high probability after
O (log 1 η ) rounds.
The following lemma illustrates the properties the working zones should satisfy and the proof is deferred to Appendix J.4.
Lemma 14 (Working zone lemma)
There exists constants ϵ 0 < ϵ 1 < ϵ 2 < ϵ 3 such that (Γ ϵ 0 , Γ ϵ 1 , Γ ϵ 2 , Γ ϵ 3 ) satisfy the following properties:
L satisfies µ -PL in Γ ϵ 3 for some µ > 0 .
Any gradient flow starting in Γ ϵ 2 converges to some point in Γ .
Then, by [10], Φ( · ) is C ∞ in Γ ϵ 2 .
Any θ ∈ Γ ϵ 1 has an ϵ 1 neighborhood B ϵ 1 ( θ ) such that B ϵ 1 ( θ ) ⊆
Γ ϵ 2 . 4.
Any gradient descent starting in Γ ϵ 0 with sufficiently small learning rate will stay in Γ ϵ 1 .
I.3. Proof Outline
We are now in a position to state the proof outline.
The general idea is to adopt the framework proposed by Li et al.
[34] to bound the closeness of the manifold projection { Φ( ¯ θ ( s ) ) }
⌊ T/ ( Hη 2 ) ⌋ s = and the solution to SDE (4), { ζ ( t ) : t ∈ [0 , T ] } .
A key component of this framework is to estimate the moments of change over a fixed time interval.
To obtain the estimation of moments for ζ ( t ) , we can directly apply the results in Li et al.
However, the estimation of the moments for Φ( ¯ θ ( s ) ) requires a careful analysis of the limiting dynamics of { ¯ θ ( s ) } ⌊ T/ ( Hη 2 ) ⌋ s = .
The dynamics of { ¯ θ ( s ) } ⌊ T/ ( Hη 2 ) ⌋ s = are divided into two phases: the approaching phase and the drift phase.
For the approaching phase (Phase 1), we show that after O (log 1 η ) rounds, the iterate will reach within ˜ O ( √ η ) from Γ (see Appendix J.5).
For the drift phase (Phase 2), we first prove that, with high probability, both ¯ θ ( s ) and θ ( s ) k,t stay close to Γ with a distance of only ˜ O ( √ η ) for all O (log 1 η ) <
s <
⌊ THη 2 ⌋ .
We also provide a high probability bound on the movement of the manifold projection (see Appendix J.6).
Based on these high probability bounds, we group R grp :=
⌊ 1 αη β ⌋ rounds together and compute the first and second moments of Φ( ¯ θ ( s + R grp ) )
− Φ( ¯ θ ( s ) ) , which is the movement of manifold projection over
R grp rounds (see Appendix J.9).
Here, β is a constant between 0 and 0 . 5 and will be specified later.
Finally, utilizing the estimation of moments in Appendix J.9, we prove that { ζ ( t ) : t ∈ [0 , T ] } following the SDE (4) are weak approximations of each other following Li et al.
[34] in Appendix J.10.
Appendix J. Proof Details of Main Theorems
The detailed proof is organized as follows.
In Appendix J.1, we introduce the notations that will be used throughout the proof.
To establish preliminary knowledge, Appendix J.2 provides explicit expression for the projection operator Φ( · ) and Appendix J.3 presents lemmas about gradient descent (GD) and gradient flow (GF).
Based on the preliminary knowledge, we prove the working zone lemma in Appendix J.4.
Appendices J.5 to J.10 make up the main body of the proof.
Specifically,
Appendices J.5 and J.6 derive high probability bounds for phase 1 and 2 respectively.
Then, we provide a summary of these high probability bounds in Appendix J.7 and the proof of Theoremin Appendix J.8.
Utilizing the high probability bounds, we derive the estimation for the first and second moments of the one step update Φ( ¯ θ ( s + R grp ) ) − Φ( ¯ θ ( s ) ) in Appendix J.9.
Finally, we prove the approximation theorem 4 in Appendix J.10.
J.1. Additional Notations
Let R tot :=
⌊ THη 2 ⌋ be the total number of rounds.
Denote by ϕ ( s ) the manifold projection of the global iterate at the beginning of round s .
Let x ( s ) k,t := θ ( s ) k,t − ϕ ( s ) be the difference between the local iterate and the manifold projection of the global iterate.
Also define ¯ x ( s ) H :
= 1 K (cid:80) k ∈
[ K ] x ( s ) k,H and ¯ x
( s ):=
1 K (cid:80) k ∈
[ K ] x ( s ) k, 0 which is the average of x ( s ) k,t among K workers at step 0 and H .
Then for all k ∈ [ K ] , x ( s ) k, 0 = ¯ x ( s )=
¯ θ ( s ) − ϕ
( s ) .
Finally, Since ∇ ℓ ( θ ; ζ ) is bounded, the gradient noise z ( s ) k,t is also bounded and we denote by σ max the upper bound such that ∥ z ( s ) k,t ∥ 2 ≤ σ max , ∀ s, k, t .
J.2. Computing the Derivatives of the Limiting Mapping
In subsection, we present lemmas that relate the derivatives of the limiting mapping Φ( · ) to the derivatives of the loss function L ( · ) .
We first introduce the operator V H .
Definition 15 For a semi-definite symmetric matrix H ∈ R d × d , let λ j , v j be the j -th eigenvalue and eigenvector and v j ’s form an orthonormal basis of R d .
Then, define the operator V H :
R d × d →
R d × d as
[Equation]
Intuitively, this operator projects M to the base matrix v i v ⊤ j and sums up the projections with weights [Equation]
Additionally, for θ ∈ Γ , denote by T θ and T ⊥ θ the tangent and normal space of Γ at θ respectively.
Lemmas 16 to 19 are from Li et al.
We include them to make the paper self-contained.
Lemma 16 (Lemma C.1 of Li et al.
For any θ ∈ Γ and any v ∈ T θ (Γ) , it holds that ∇ L ( θ ) v = 0 .
Lemma 17 (Lemma 4.3 of Li et al.
For any θ ∈ Γ , ∂ Φ( θ ) ∈ R
d × d is the projection matrix onto the tangent space T θ (Γ) .
Lemma 18 (Lemma C.4 of [37])
For any θ ∈ Γ , u ∈ R d and v ∈ T θ (Γ) , it holds that ∂
[Equation]
∇ [Equation] u ] − ∇ [Equation]
L ( θ ) + ∇ [Equation] L ( θ )[ v , ∂ Φ( θ
[Equation]
u ] .
Lemma 19 (Lemma C.6 of [37])
For any θ ∈ Γ and Σ ∈ span { uu ⊤ | u ∈ T ⊥ θ (Γ) } , (cid:10) ∂ 2 Φ( θ ) , Σ (cid:11) = − ∂ Φ( θ ) ∇ 3 L ( θ )
[ V ∇ 2 L ( θ ) ( Σ )] .
Lemma
For all θ ∈ Γ , u , v ∈ T θ (Γ) , it holds that
[Equation]
Proof This proof is inspired by Lemma C.4 of [37].
For any θ ∈ Γ , consider a parameterized smooth curve v ( t ) , t ≥ 0 on Γ such that v (0)
= θ and v ′ (0)
= v .
Let P ∥ ( t ) =
∂
Φ( v ( t )) ,
P ⊥ ( t ) =
I − ∂ Φ( v ( t )) and H ( t ) = ∇ 2 L ( v ( t )) .
By Lemma C.1 and 4.3 in [37],
[Equation]
Take the derivative with respect to t on both sides,
[Equation]
⇒ P
[Equation]
At t = 0 , we have
[Equation]
WLOG let H (0) =
diag( λ 1 , · · · , λ d ) , ∈ R d × d , where λ i = 0 for all m < i ≤ d .
Therefore
P ⊥ (0) =
(cid:20) I m   (cid:21) , P ∥ (0) =
(cid:20)
I d − m (cid:21) .
Decompose P ′∥ (0) , H (0) and H ′ (0) as follows.
[Equation]
Substituting the decomposition into (26), we have
[Equation]
Therefore, H ′ 22 (0) =
and
[Equation]
Any u ∈ T θ (Γ) can be decomposed as u = (cid:2) 0 , u 2 (cid:3) ⊤ where u 2 ∈ R
d − m .
With this decomposition, we have P ∥ (0) H ′
(0) u
= 0 .
Also, note that H ′ (0) =
∇ 3 L ( θ )[ v ] .
Hence, [Equation]
J.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 ( θ ) is ρ -smooth and µ -PL in an open, convex neighborhood U of a local minimizer θ ∗ .
Denote by L ∗ :
= L ( θ ∗ ) the minimum value for simplicity.
Let ϵ ′ be the radius of the open ϵ ′ -ball centered at θ ∗
such that B ϵ
( θ ∗ ) ⊆
U .
We also define a potential function ˜Ψ( θ ) :
= (cid:112) L ( θ ) − L ∗ .
Consider gradient descent iterates { ˆ u t } t ∈ N following the update rule ˆ u t
+1 =
ˆ u t − η ∇L (ˆ u t ) .
We first introduce the descent lemma for gradient descent.
Lemma 21 (Descent lemma for GD)
If ˆ u t ∈ U and η ≤ 1 ρ , then
[Equation]
and
[Equation]
Proof By ρ -smoothness,
[Equation]
By the definition of µ -PL, we have
[Equation]
Then we prove the Lipschitzness of ˜Ψ( θ ) .
Lemma 22 (Lipschitzness of ˜Ψ( θ ) )
˜Ψ( θ ) is √ 2 ρ -Lipschitz for θ ∈ U .
That is, for any θ 1 , θ 2 ∈
U ,
[Equation]
Proof Fix θ 1 and θ 2 .
Denote by θ ( t ) :
= (1 − t )
θ 1 + t θ 2 the convex combination of θ 1 and θwhere t ∈ [0 , 1] .
Further define f ( t ) :
= ˜Ψ( θ ( t )) .
Below we consider two cases.
Case 1.
If ∀ t ∈ (0 , 1) , f ( t ) > 0 , then f ( t ) is differentiable on (0 , 1) .
[Equation]
By ρ -smoothness of L , for all θ ∈ U ,
[Equation]
Since (cid:112)
L
[Equation]
θ
[Equation] t )) [Equation] L ∗ > 0 for all t ∈ (0 , 1) , ∥∇L ( θ ( t ))
∥ 2 √ L ( θ ( t ))
−L ∗ ≤ √ 2 ρ .
Therefore,
[Equation]
Case 2.
If ∃ t ′ ∈ (0 , 1) such that f ( t ′ )
= 0 , then
[Equation]
Since θ ( t ′ ) minimizes L in an open set, ∇L ( θ ( t ′ ))
= 0 .
By ρ -smoothness of L , for all θ ∈ U ,
[Equation]
Therefore,
[Equation]
Then we have
[Equation]
Combining case 1 and case 2, we conclude the proof.
Below we introduce a lemma that relates the movement of one step gradient descent to the change of the potential function.
Lemma 23 (Lemma G.1 in [40])
If ˆ u t ∈ U and η ≤ 1 /ρ 2 then
[Equation]
Proof
[Equation]
where the two inequalities uses Theorem 21.
By µ -PL, ˜ Ψ(ˆ u t ) ≤ 1 √ 2 µ ∥∇L (ˆ u t ) ∥ 2 .
Therefore, we have ˜Ψ
[Equation]
Based on Theorem 23, we have the following lemma that bounds the movement of GD over multiple steps.
Lemma 24 (Bounding the movement of GD)
If ˆ u 0 is initialized such that ∥ ˆ u 0 − θ ∗ ∥ 2 ≤ 14 (cid:113) µρ ϵ ′ , then for all t ≥ 0 , ˆ u t ∈ B ϵ ′ ( θ ∗ ) and
[Equation]
Proof We prove the proposition by induction.
When t = 0 , it trivially holds.
Assume that the proposition holds for ˆ u τ , 0 ≤ τ < t .
For step t , since ˆ u τ ∈ B ϵ ′ ( θ ∗ ) , we apply Theorem 23 and obtain
[Equation]
Further by ρ -smoothness of L ( · ) ,
[Equation]
Therefore, ∥ ˆ u t − θ ∗ ∥ 2 ≤ ∥ ˆ u t − ˆ u 0 ∥ 2 + ∥ ˆ u 0 − θ ∗ ∥ 2 < ϵ ′ , which concludes the proof.
Finally, we introduce a lemma adapted from Thm.
D.4 of which bounds the movement of GF.
Lemma 25 Assume that ∥ θ 0 − θ ∗ ∥ 2 < (cid:113) µρ ϵ ′ .
The gradient flow θ ( t ) = − d L ( θ ( t ))
d t starting at θconverges to a point in U and
[Equation]
Proof Let T :=
inf { t : θ / ∈ U } .
Then for all t < T ,
[Equation]
By µ -PL, ∥∇L ( θ ) ∥ 2 ≥ (cid:112) 2 µ ( L ( θ ) − L ∗ ) .
Hence,
[Equation]
Integrating both sides, we have
[Equation]
where the second inequality uses ρ -smoothness of L .
Therefore, T = + ∞ and θ ( t ) converges to some point in U .
J.4. Construction of working zones
In this subsection, we provide the proof for Theorem 14.
Note that the notions of Z ϵ , M ϵ 4 , ρ 2 , ρ 3 , ν 1 and ν 2 defined in the proof will be useful in the remaining part of this section.
Proof [Proof of Theorem 14] Let ¯ θ (0) be initialized such that Φ( ¯ θ (0) ) ∈ Γ .
Let Z be the set of all points on the gradient flow trajectory starting from ¯ θ (0) and Z ϵ be the ϵ -neighborhood of Z , where ϵ is a 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 [40], there exists an ϵ 3 -neighborhood of Γ , Γ ϵ 3 , such that L satisfies µ -PL for some µ > 0 .
Let M be the convex hull of Γ ϵ 3 ∪ Z ϵ
and
M ϵ 4 be the ϵ 4 -neighborhood of M where ϵ 4 is a positive constant.
Then M ϵ 4 is bounded.
Define ρ 2 = sup θ ∈M ϵ 4 ∥∇ 2 L ( θ ) ∥ 2 and ρ 3 = sup M ϵ 4 ∥∇ 3 L ( θ ) ∥ 2 .
By Theorem 25, we can construct an ϵ 2 -neighborhood of Γ where ϵ 2 < (cid:113) µρ 2 ϵ 3 such that all GF starting in Γ ϵ 2 converges to Γ .
By Falconer [10], Φ( · ) is C 2 in Γ ϵ 3 .
Define ν 1 = sup θ ∈ Γ ϵ 3 ∥ ∂ Φ( θ ) ∥ 2 and ν 2 = sup θ ∈ Γ ϵ 3 ∥ ∂ 2 Φ( θ ) ∥ 2 .
We also construct an ϵ 1 neighborhood of Γ , Γ ϵ 1 , where ϵ 1 ≤ 12 ϵ 2 < 1 2 (cid:113) µ ρ 2 ϵsuch that all θ ∈ Γ ϵ 1 has an ϵ 1 neighborhood where Φ is well defined.
Finally, by Theorem 24, there exists an ϵ 0 -neighborhood of Γ where ϵ 0 ≤ 14 (cid:113) µρ 2 ϵ 1 such that all gradient descent iterates starting in Γ ϵ 0 with η ≤ 1 ρ 2 will stay in Γ ϵ 1 .
When analyzing the limiting dynamics of Local SGD, we will show that all θ ( s ) k,t stays in Γ ϵ 2 ,
˜ u
( s ) t ∈ Γ ϵ 1 , ¯ θ ( s ) ∈
Γ ϵ 0 with high probability after O (log 1 η ) rounds.
J.5. High Probability Bounds for Phase 1: Iterate Approaching the Manifold
The approaching phase can be further divided two subphases.
In the first subphase, ¯ θ (0) is initialized such that ϕ (0) ∈ Γ .
We will show that after a constant number of rounds s 0 , ¯ θ ( s 0 ) goes to the inner part of Γ ϵ 0 such that ∥ ¯ θ ( s 0 ) − ϕ (0) ∥ 2 ≤ cϵ 0 with high probability, where 0 < c < 1 and the constants will be specified later (see Appendix J.5.2).
In the second subphase, we show that the iterate can reach within ˜ O ( √ η ) distance from Γ after O (log 1 η ) rounds with high probability (see
Appendix J.5.3).
J.5.1. A DDITIONAL NOTATIONS
Consider an auxiliary sequence { ˜ u ( s ) t } where ˜ u ( s )
= ¯ θ ( s ) and ˜ u ( s ) t +1 = ˜
u ( s ) t [Equation] ≤
H − 1 .
Define ˜ ∆ ( s ) k,t := θ ( s ) k,t − ˜ u ( s ) t to be the difference between the local iterate and the gradient descent iterate.
Notice that ˜ ∆ ( s ) k, 0 = 0 , for all k and s .
Consider a gradient flow { u ( t ) } t ≥ 0 with the initial condition u (0)
=
¯ θ (0) and converges to ϕ (0) ∈ Γ .
For simplicity, let u ( s ) t := u ( sα + tη ) be the gradient flow after s rounds plus t steps.
Let s 0 be the smallest number such that ∥ u ( s 0 )− ϕ (0) ∥ 2 ≤ 14 (cid:113) µρ 2 ϵ 0 .
Note that s 0 is a constant independent of η .
In this subsection, the minimum value of the loss in Appendix J.3 corresponds to the loss value on Γ , i.e., L ∗ =
L ( ϕ ) , ∀ ϕ ∈ Γ .
We also define the following sequence { ˜ Z ( s ) k,t } Ht =0 that will be used in the proof.
Define
[Equation]
(cid:32)
[Equation]
l =
[Equation]
( I
[Equation]
J.5.2. P ROOF FOR S UBPHASE
First, we have the following lemma about the concentration of ˜ Z ( s ) k,t .
Lemma 26 (Concentration property of { ˜ Z ( s ) k,t } Ht =0 )
Given ¯ θ ( s ) such that ˜ u ( s ) t ∈ Γ ϵ 3 ∪ Z ϵ for all 0 ≤ t ≤ H , then with probability at least 1 − δ ,
[Equation]
where ˜ C 1 := exp( αρ 2 ) .
Proof For each ˜ Z ( s ) k,t , construct a sequence { ˜ Z ( s ) k,t,t ′ } tt ′ =0 :
[Equation]
Since ˜ u ( s ) t ∈ Γ ϵ 3 ∪ Z ϵ , we have ∥∇ 2 L (˜ u ( s ) t ) ∥ 2 ≤ ρ 2 for all 0 ≤ t ≤ H .
Then, for all τ and t , ∥ t [Equation] 1 (cid:89)
[Equation]
Notice that for all 0 ≤ t ≤ H , { ˜ Z ( s ) k,t,t ′ } tt ′ =0 is a martingale with ∥ ˜ Z ( s ) k,t,t ′ − ˜ Z ( s ) k,t,t ′ − 1 ∥ 2 ≤ ˜ C 1 σ max .
By Azuma-Hoeffding’s inequality,
[Equation]
Taking a union bound on all k ∈ [ K ] and 0 ≤ t ≤ H , we can conclude that with probability at least 1 − δ ,
[Equation]
The following lemma states that the gradient descent iterates will closely track the gradient flow with the same initial point.
Lemma 27 Denote G := sup t ≥
∥∇L ( u ( t ))
∥ 2 as the upper bound of the gradient on the gradient flow trajectory.
If ∥ ˜ u ( s ) t − u ( s ) t ∥ 2 = O ( √ η ) , then for all 0 ≤ t ≤ H , the closeness of ˜ u
( s ) t and u ( s ) t is bounded by
[Equation]
where ˜ C 1 = exp( αρ 2 ) .
Proof We prove by induction that
[Equation]
When t = [Equation] (27) holds trivially.
Assume that (27) holds for 0 ≤ τ ≤ t , then
[Equation]
By smoothness of L ,
[Equation]
Since ρ 22 η 2 G (cid:80) t − 1 τ =0 (1 + ρ 2 η ) τ ≤ ηG (1 + ρ 2 η ) t ≤ exp( αρ 2 ) ηG , then ∥ ˜ u ( s ) t
[Equation]
u
[Equation]
, which implies that ˜ u ( s ) t ∈ M ϵ 4 . Hence, ∥∇L (˜ u ( s ) t ) − L ( u ( s ) t
[Equation]
≤
[Equation]
t
[Equation]
By triangle inequality,
[Equation]
which concludes the induction step.
Appling 1 + ρ 2 η ≤ exp( ρ 2 η ) , we have the lemma.
Utilizing the concentration probability of { ˜ Z ( s ) k,t } , we can obtain the following lemma which implies that the Local SGD iterates will closely track the gradient descent iterates with high probability.
Lemma
Given ¯ θ ( s ) such that ˜ u ( s ) t ∈ Γ ϵ 3 ∪ Z ϵ for all 0 ≤ t ≤ H , then for δ =
O (poly( η )) , with probability at least 1 − δ , there exists a constant ˜ C 3 such that
[Equation]
and
[Equation]
Proof Since ˜ u ( s ) t [Equation] Z ϵ for all 0 ≤ t ≤ H , we have ∥∇ 2 L (˜ u ( s ) t )
∥ 2 ≤ ρ 2 .
According to the update rule for θ ( s ) k,t and ˜ u
( s ) t ,
[Equation]
Subtracting (29) from (28) gives
[Equation]
Here, ˜ v ( s ) k,t = (1 − β ( s ) k,t ) θ ( s ) k,t + β ( s ) k,t ˜ u ( s ) k,t , where β ( s ) k,t ∈ (0 , 1) depends on θ ( s ) k,t and ˜ u
( s ) t .
Therefore, ∥ ˜ v ( s ) k,t ∥ 2 ≤ ρ 3 2 ∥ ˜ ∆ ( s ) k,t ∥
if θ ( s ) k,t ∈ M ϵ 4 .
Applying (30) t times, we have
[Equation]
By Cauchy-Schwartz inequality, triangle inequality and the definition of ˜ Z ( s ) k,t , if for all 0 ≤ τ ≤ t −and k ∈ [ K ] , θ ( s ) k,τ ∈ M ϵ 4 , then we have
[Equation]
Applying Theorem 26 and substituting in the value of H , we have that with probability at least 1 − δ ,
[Equation]
[Equation]
Now we show by induction that for δ
= O (poly( η )) , when (32) holds, there exists a constant ˜ C 2 > 2 σ max √
α ˜ C
such that ∥ ˜ ∆ ( s ) [Equation]
When t = 0 , ˜ ∆ ( s ) k, 0 = 0 .
Assume that ∥ ˜ ∆ ( s ) k,τ ∥ 2 ≤ ˜ C 2 (cid:113) η log 2 αKηδ , for all k ∈ [ K ] , 0 ≤ τ ≤ t − 1 .
Then for all 0 ≤ τ ≤ t − 1 , θ ( s ) k,τ ∈ M ϵ 4 .
Therefore, we can apply (31) and obtain
[Equation]
Given that ˜ C 2 ≥ 2 σ max √ 2 α ˜ C 1 and δ = O (poly( η )) , when η is sufficiently small, ∥ ˜ ∆ ( s ) k,t ∥ 2 ≤ ˜ C 2 (cid:113)
η log 2 αKηδ .
To sum up, for δ
= O (poly( η )) , with probability at least 1 − δ , ∥ ˜ ∆ ( s ) k,t ∥ 2 ≤ ˜ C 2 (cid:113) η log 2 αKηδ for all k ∈ [ K ] , 0 ≤ t ≤ H .
By triangle inequality,
[Equation]
The combination of Theorem 27 and Theorem 28 leads to the following lemma, which states that the Local SGD iterate will enter Γ ϵ 1 after s 0 rounds with high probability.
Lemma
Given ¯ θ (0) such that Φ( ¯ θ (0) )
∈ Γ , then for δ
= O (poly( η )) , there exists a positive constant ˜ C 4 such that with probability at least 1 − δ ,
[Equation]
Proof First, we prove by induction that for δ
= O (poly( η )) , when
[Equation]
the closeness of ¯ θ ( s ) and u ( s ) 0 is bounded by
[Equation]
When s =
[Equation]
¯ θ (0) =
u (0)0 .
Assume that (34) holds for round s .
Then by Theorem 27, for all 0 ≤ t ≤ H ,
[Equation]
Therefore
[Equation]
for sufficiently small η , ˜ u ( s ) t ∈
Z ϵ , ∀ 0 ≤ t ≤ H .
Combing the above inequality with
Theorem
[Equation]
we have
[Equation]
which concludes the induction.
Therefore, when (33) holds, there exists a positive constant ˜ C 4 such that
[Equation]
By definition of u ( s 0 )
[Equation]
Finally, according to Theorem 26, (33) holds with probability at least 1 − δ .
J.5.3. P ROOF FOR S UBPHASE
In subphase 2, we show that the iterate can reach within ˜ O ( √ η ) distance from Γ after O (log 1 η ) rounds with high probability.
The following lemma manifests how the potential function ˜Ψ( ¯ θ ( s ) )
evolves after one round.
Lemma 30 Given ¯ θ ( s ) ∈ Γ ϵ 0 , for δ
= O (poly( η )) , with probability at least 1 − δ ,
[Equation]
and
[Equation]
where ˜ C 5 is a positive constant.
Proof Since ¯ θ ( s ) ∈ Γ ϵ 0 , then for all 0 ≤ t ≤ H , ˜ u ( s ) t ∈
Γ
ϵ 1 by the definition of the working zone.
By Theorem 21, for η ≤ 1 ρ 2 ,
[Equation]
Specially, for t =
H ,
[Equation]
Therefore,
[Equation]
According to the proof of Theorem 28, for δ
= O (poly( η )) , when
[Equation]
(35) there exists a constant ˜ C 3 such that
[Equation]
and ∥
[Equation]
Since ˜ u ( s ) t ∈ Γ ϵ 1 , ∀ 0 ≤ t ≤ H , ¯ θ ( s {++
 )+} ∈
Γ ϵ 2 and ¯ θ ( s ) k,t ∈
Γ ϵ 2 , ∀ 0 ≤ t ≤ H , k ∈
[ K ] .
By Theorem 22, ˜Ψ( · ) is √ 2 ρ 2 -Lipschitz in M ϵ 4 .
Therefore, when (35) holds, there exists a constant ˜ C 5 := √ 2 ρ 2 ˜ C 3 such that
[Equation]
and
[Equation]
Finally, by Theorem 26, (35) holds with probability at least 1 − δ .
We are thus led to the following lemma which characterizes the evolution of the potential ˜Ψ( ¯ θ ( s ) ) and ˜Ψ( θ ( s ) k,t ) over multiple rounds.
Lemma 31 Given ∥ ¯ θ (0) − ϕ (0) ∥ 2 ≤ 12 (cid:113) µρ 2 ϵ 0 , for δ = O (poly( η )) and any integer 1 ≤ R ≤ R tot , with probability at least 1 − δ ,
[Equation]
Furthermore,
[Equation]
Proof We prove induction that for δ
= O (poly( η )) , when
[Equation]
then for all 0 ≤ s ≤ R , (36) and (37) hold.
When s = 0 , ¯ θ (0) ∈
Γ ϵ 0 and (36) trivially holds.
By Theorem 30, (37) holds.
Assume that (36) and (37) hold for round s − 1 .
Then for round s , by Theorem 30, ¯ θ ( s ) ∈
Γ ϵ 2 and
[Equation]
where the second inequality comes from the induction hypothesis.
By Theorem 25,
[Equation]
Here, the last inequality uses ˜Ψ( ¯ θ (0) ) ≤ (cid:113) ρ 2 2 ∥ ¯ θ ( s ) − ϕ (0) ∥ 2 ≤ 12 (cid:113) µ 2 ϵ 0 .
Hence, when η is sufficiently small, ¯ θ ( s ) ∈ Γ ϵ 0 .
Still by Theorem 30, ¯ θ ( s ) k,t ∈
Γ ϵ 2 and
[Equation]
Finally, according to Theorem 26, (38) holds with probability at least 1 − δ .
The following corollary is a direct consequence of Theorem 31 and Theorem 25.
Corollary 32 Let s 1 :=
⌈ 20 αµ log 1 η ⌉ .
Given ∥ ¯ θ (0) − ϕ (0) ∥ 2 ≤ 12 (cid:113) µρ 2 ϵ 0 , for δ
= O (poly( η )) , with probability at least 1 − δ ,
[Equation]
where ˜ C 6 is a constant.
Proof Substituting in R = s 1 to Theorem 31 and applying ∥ ¯ θ ( s 1 ) − ϕ
( s ) ∥ 2 ≤ (cid:113) 2 µ ˜Ψ( ¯ θ ( s 1 ) ) for ¯ θ ( s 1 )
∈
Γ ϵ 0 , we have the lemma.
Finally, we provide a high probability bound for the movement of the projection on the manifold after s 1 rounds ∥
ϕ
( s 1 ) − ϕ
(0) ∥ 2 .
Lemma 33 Let s 1 := ⌈ 20 αµ log 1 η ⌉ .
Given ∥ ¯ θ (0) − ϕ (0) ∥ 2 ≤ 12 (cid:113) µρ 2 ϵ 0 .
For δ = O (poly( η )) , with probability at least 1 − δ ,
[Equation]
Proof From Theorem 31, for δ
= O (poly( η )) , when ∥
[Equation]
then ¯ θ ( s ) ∈
Γ ϵ 0 , for all 0 ≤ s ≤ s 1 .
By the definition of Γ ϵ 0 , ˜ u ( s ) t ∈ Γ ϵ 1 , for all 0 ≤ t ≤ H, 0 ≤ s ≤ s 1 .
By triangle inequality, ∥ ϕ ( s 1 ) − ϕ (0) ∥ 2 can be decomposed as follows. [Equation]
[Equation]
By Theorem 28, when (40) hold , then for all 0 ≤ s < s 1 − 1 ,
[Equation]
This implies that ¯ θ ( s +1) ∈
[Equation]
ϵ 1 (˜ u ( s ) H ) .
Since for all θ ∈ Γ ϵ 2 , ∥ ∂ Φ( θ ) ∥ 2 ≤ ν 1 , then Φ( · ) is ν 1Lipschitz in B ϵ 1 (˜ u ( s ) H ) .
This gives
[Equation]
Then we analyze ∥ ¯ θ ( s +1)
− ˜ u
( s ) H ∥ 2 .
By Theorem 24 and the definition of Γ ϵ 0 and Γ ϵ 1 , there exists ϕ ∈ Γ such that ˜ u ( s ) t ∈ B ϵ 1 ( ϕ ) , ∀ 0 ≤ t ≤ H .
Therefore, we can expand Φ(˜ u ( s ) t +1 ) as follows:
[Equation]
where c ( s ) t ∈ (0 , 1) .
Then we have
[Equation]
∥ [Equation] By Theorem 21, η 2 ∥∇L (˜ u ( s ) t ) ∥ 22 ≤ L (˜ u ( s ) t ) − L (˜ u ( s ) t +1 ) .
Therefore, ∥ Φ(
[Equation]
where the last inequality uses Cauchy-Schwartz inequality and Theorem 31.
Summing up (43)
[Equation]
we obtain
[Equation]
where ˜ C 7 is a constant.
Substituting (42) and (44) into (41) [Equation] for sufficiently small η , we have
[Equation]
where ˜ C 8 is a constant.
Finally, according to Theorem 26, (40) holds with probability at least 1 − δ .
J.6. High Probability Bounds for Phase 2: Iterates Staying Close to Manifold
In this subsection, we show that ∥ x ( s ) k,t ∥ 2 = ˜ O ( √ η ) and ∥ ¯ θ ( s + r ) − ¯ θ ( s ) ∥
= ˜ O ( η 0 .
5 − 0 .
5 β ) , ∀ 0 ≤ r ≤ R grp with high probability.
J.6.1. A DDITIONAL NOTATIONS
Before presenting the lemmas
[Equation]
we define the following martingale { m ( s ) k,t } Ht =0 that will be useful in the proof:
[Equation]
We also define ˜ P : R d → R d × d as an extension of ∂ Φ :
[Equation]
Finally, we define a martingale { Z ( s ) t
[Equation]
J.6.2. P ROOF FOR THE H IGH P ROBABILITY B OUNDS
A direct application of Azuma-Hoeffding’s inequality yields the following lemma.
Lemma 34 (Concentration property of m ( s ) k,t )
With probability at least 1 − δ , the following holds:
[Equation]
where ˜ C 9 is a constant.
Proof Notice that ∥ m ( s ) k,t +1 − m ( s ) k,t ∥ 2 ≤ σ max .
Then by Azuma-Hoeffdings inequality,
[Equation]
Taking union bound on K clients, H local steps and R grp rounds, we obtain that the following inequality holds with probability at least 1 − δ :
[Equation]
Substituting in H = αη and R grp =
⌊ 1 αη β ⌋ yields the lemma.
Again applying Azuma-Hoeffding’s inequality, we have the following lemma about the concentration property of Z ( s ) t .
Lemma 35 (Concentration property of Z ( s ) t )
With probability at least 1 − δ , the following inequality holds:
[Equation]
Proof Notice that ∥ Z ( s ) t +1 − Z ( s ) t ∥ 2 ≤ ν 2 σ max , ∀ 0 ≤ t ≤ H − 1 and ∥ Z ( s +1)
[Equation]
By Azuma-Hoeffding’s inequality,
[Equation]
Taking union bound on R grp rounds, we obtain that the following inequality holds with probability at least 1 − δ :
[Equation]
Substituting in H = αη and R grp =
⌊ 1 αη β ⌋ yields the lemma.
We proceed to present a direct corollary of Theorem 31 which provides a bound for the potential function over R grp rounds.
Lemma 36 Given ∥ ¯ θ (0)
[Equation]
log 1 η where C 0 is a constant, then for δ
= O (poly( η )) , with probability at least 1 − δ ,
[Equation]
[Equation] [Equation]
and
[Equation]
where C 1 is a constant that can depend on C 0 .
Furthermore,
[Equation]
where ˜ C 9 is a constant independent of C 0 .
Proof By ρ 2 -smoothness of L , ˜Ψ( ¯ θ (0) ) ≤ C 0 (cid:113) ηρ 2 2 log 1 η .
Substituting R grp =
⌊ 1 αη β ⌋ and ˜Ψ( ¯ θ (0) ) ≤ C(cid:113) ηρ 2 2 log 1 η into Theorem 31, for δ
= O (poly( η )) , with probability at least 1 − δ , (45) and (46) where C 1 is a constant that can depend on C 0 .
Furthermore, for round ¯ θ ( R grp ) ,
[Equation]
where ˜ C 9 is a constant independent of C 0 .
Lemma 37 Given ∥ ¯ θ [Equation] − [Equation]
η
[Equation]
1 η where C 0 is a constant, then for δ
= O (poly( η )) , with probability at least 1 − δ , for all 0 ≤ s [Equation] < [Equation]
grp
[Equation]
≤
[Equation]
∈
[Equation]
where C 2 is a constant that can depend C 0 .
Furthermore,
[Equation]
where ˜ C 11 is a constant independent of C 0 .
Proof Decomposing x ( s ) k,t by triangle inequality, we have
[Equation]
We first bound ∥ ¯ θ ( s ) − ϕ ( s ) ∥ 2 .
By Theorem 36, for δ = O (poly( η )) , with probability at least 1 − δ 2 ,
[Equation]
and
[Equation]
where C 2 is a constant that may depend on C 0 and ˜ C 10 is a constant independent of C 0 .
When (47) and (49) hold, by Theorem 25,
[Equation]
Then we bound ∥ θ ( s ) k,t − ¯ θ ( s ) ∥ 2 .
By the update rule, we have
[Equation]
Still by triangle inequality
[Equation]
we have
[Equation]
Due to ρ 2 -smoothness of L , when (48) holds,
[Equation]
By Theorem 34, with probability at least 1 − δ 2 ,
[Equation]
Combining (52) and (53), when (48) and (49) hold simultaneously, there exists a constant C 3 which can depend on C 0 such that
[Equation]
By triangle inequality,
[Equation]
Combining (50), (51) and (54), we complete the proof.
Then we provide high probability bounds for the movement of ϕ ( s ) within R grp rounds.
Lemma 38 Given ∥ ¯ θ (0) − ϕ (0) ∥ 2 ≤ C 0 (cid:113) η log 1 η where C 0 is a constant, then for δ =
O (poly( η )) , with probability at least 1 − δ ,
[Equation]
where C 4 is a constant that can depend on C 0 .
Proof By the update rule of Local SGD,
[Equation]
Averaging among K clients gives
[Equation]
By Theorem
[Equation]
for δ = O (poly( η ))
[Equation]
the following holds with probability at least 1 − δ/ 3 ,
[Equation]
When (55) and (56) hold
[Equation]
we can expand Φ( ¯ θ ( s +1) ) as follows:
[Equation]
where a ( s ) ∈ (0 , 1) .
Telescoping from round 0 to s − 1 , we have
[Equation]
From (56), we can bound ∥T ( s ) 3 ∥ 2 by ∥T ( s ) 3 ∥ 2 ≤ 12 ν 2 C 22 η log 3 ηδ .
We proceed to bound ∥T ( s ) 1 ∥ 2 .
When (55) and (56) hold, we have
[Equation]
where b ( s ) k,t ∈ (0 , 1) .
By Theorem 31, with probability at least 1 − δ/ 3 , the following holds:
[Equation]
When (55), (56) and (57) hold simultaneously, we have for all 0 ≤ s < R grp ,
[Equation]
Finally, we bound ∥ (cid:80) s − 1 r =0 T ( r ) 2 ∥ 2 .
By Theorem 35, the following inequality holds with probability at least 1 − δ/ 3 :
[Equation]
[Equation] When (55), (56) and (58) hold simultaneously [Equation] we have
[Equation]
Combining the bounds for ∥T ( s ) 1 ∥ 2 , ∥ (cid:80) sr
=0 T ( r ) 2 ∥ 2 and ∥T ( s ) 3 ∥ 2 and taking union bound, we obtain that for δ
= O (poly( η )) , the following inequality holds with probability at least 1 − δ :
[Equation]
where C 4 is a constant that can depend on C 0 .
J.7. Summary of High Probability Bounds
Based on the results in Appendix J.5 and Appendix J.6, we summarize the dynamics of Local SGD iterates in this subsection.
For convenience, we first introduce the definition of global step and δ -good step .
Definition 39 (Global step) Define an index set I :=
{ ( s, t ) | s ≥ 0 , 0 ≤ t ≤ H } with lexicographical order, which means ( s 1 , t 1 ) ⪯ ( s 2 , t 2 )
if and only if s 1 < s 2 or ( s
= s 2 and t 1 ≤ t 2 ).
A global step is indexed by ( s, t ) which corresponds to the t -th local step at round s .
Definition 40 ( δ -good step)
In the training process of Local SGD, we say the global step ( s, t ) ⪯ ( R tot , 0) is δ -good if the following inequalities hold:
[Equation]
Applying the concentration properties of ˜ Z ( r ) k,τ , m ( r ) k,τ and Z ( r ) H (Lemmas 35, 34 and 26) yields the following theorem.
Theorem 41 For δ =
O (poly( η )) , with probability at least 1 − δ , all global steps ( s, t ) ⪯ ( R tot , 0) are δ -good.
In the remainder of this subsection, we use O ( · ) notation to hide constants independent of δ and η .
Now we are ready to present a summary of the dynamics of Local SGD when ¯ θ (0) is initialized such that Φ( ¯ θ (0) ) ∈ Γ and all global steps are δ -good.
Phase 1 lasts for s
+ s
=
O (log 1 η ) rounds.
At the end of phase 1, the iterate reaches within O ( (cid:113) η log 1 ηδ ) from Γ , i.e., ∥ ¯ θ ( s 0 + s 1 ) − ϕ
( s 0 + s 1 ) ∥ 2 =
O ( (cid:113) η log 1 ηδ ) .
The movement of the projection on manifold over s 0 + s 1 rounds, ∥ ϕ ( s 1 + s 0 ) − ϕ (0) ∥ 2 , is bounded by O (log 1 η (cid:113) η log 1 ηδ ) .
After s 0 + s 1 rounds, the dynamic enters phase 2 when the iterates stay close to Γ with ¯ θ ( s ) ∈
Γ ϵ 2 , ∀ s 0 + s 1 ≤ s ≤ R tot and θ ( s ) k,t ∈ Γ ϵ 2 , ∀ k ∈
[ K ] , ( s 0 + s 1 , 0) ⪯ ( s, t ) ⪯ ( R tot , 0) .
Furthermore, ∥ x ( s ) k,t ∥ 2 and ∥ ¯ x ( s ) H ∥ 2 satisfy the following equations:
[Equation]
Moreover, the movement of the manifold projection within R grp rounds can be bounded as follows:
[Equation]
Finally, we provide a theorem which states that ¯ θ ( s ) stays within ˜ O ( √ η ) from the manifold after
O (log 1 η ) rounds with high probability.
This theorem is a direct consequence of the lemmas in
Appendix J.5 and J.6.
Theorem 42 For δ =
O (poly( η )) , with probability at least 1 − δ , for all O (log 1 η ) ≤ s ≤ ⌊
[Equation]
where O ( · ) hides constants independent of η and δ .
Proof [Proof for Theorem 42] By Lemmas 29, 37 and Theorem 32, for δ =
O (poly( η )) , when all global steps are δ -good, ¯ θ ( s ) ∈
Γ ϵ 2 , ∀ s 0 + s 1 ≤ s ≤ R tot and θ ( s ) k,t ∈ Γ ϵ 2 , ∀ k ∈
[ K ] , ( s 0 + s 1 , 0) ⪯ ( s, t ) ⪯ ( R tot , 0) and ∥ x ( s ) k,t ∥ 2 , ∥ ¯ x ( s ) H ∥ 2 satisfy the following equations:
[Equation]
Hence ∥ ¯ x ( R tot )
[Equation]
L
[Equation]
Theorem 25.
According to Theorem 41, with probability at least 1 − δ , all global steps are δ -good, thus completing the proof.
J.8. Proof of Theorem
In this subsection, we explicitly derive the dependency of the approximation error on α .
The proofs are quite similar to those in Appendix J.5 and hence we only state the key proof idea for brevity.
With the same method as the proofs in Appendix J.5.2, we can show that with high probability, ∥ ¯ θ ( s ) − [Equation] [Equation] µρ 2 after s ′
= O (1) rounds.
Below we focus on the dynamics of Local SGD thereafter.
We first remind the readers of the definition of { ˜ Z sk,t } :
[Equation]
We have the following lemma that controls the norm of the matrix product (cid:81) t − 1 l
= τ +
( I − η ∇ 2 L (˜ u ( s ) l )) .
Lemma
Given ¯ θ ( s ) ∈ Γ ϵ 0 , then there exists a positive constant C ′ 3 independent of α such that for all 0 ≤ τ < t ≤ H ,
[Equation]
Proof Since ¯ θ ( s ) ∈ Γ ϵ 0 , then ˜ u ( s ) t ∈ Γ ϵ 1 for all 0 ≤ t ≤ H .
We first bound the minimum eigenvalue of ∇ 2 L (˜ u ( s ) t ) .
Due to the PL condition, by Theorem 21, for η ≤ 1 ρ 2 ,
[Equation]
Therefore,
[Equation]
Let C ′
= ρ 3 (cid:113) ρ 2 µ .
By Weyl’s inequality,
[Equation]
where the last two inequalities use Lemmas 25 and 22 respectively.
Therefore, for all 0 ≤ t ≤ H and 0 ≤ τ ≤ t − 1 ,
[Equation]
For sufficiently small η , there exists a constant C ′ 2 such that
[Equation]
Substituting (60) into (59), we obtain the lemma.
Based on Theorem 43, we obtain the following lemma about the concentration property of ˜ Z ( s ) k,t , which can be derived in the same way as Theorem 26.
Lemma 44 Given ¯ θ ( s ) ∈ Γ ϵ 0 , then with probability at least 1 − δ ,
[Equation]
where C ′ 3 is defined in Theorem 43.
The following lemma can be derived analogously to Theorem 28 but the error bound is tighter in terms of its dependency on α .
Lemma
Given ¯ θ ( s ) ∈ Γ ϵ 1 , then for δ = O (poly( η )) , with probability at least 1 − δ , there exists a constant C ′ 4 independent of α such that
[Equation]
and
[Equation]
Then, similar to Theorem 31, we can show that for δ
= O (poly( η )) and simultaneously
all s ≥ s ′
+ s ′ 1 where s ′
= O ( 1 α log 1 η ) , it holds with probability at least 1 − δ that ∥ ¯ θ ( s ) − ϕ ( s ) ∥ 2 =
O ( (cid:113) αη log αηδ ) .
Note that to eliminate the dependency of the second term’s denominator on α in (37), we can discuss the cases of α > c
and α < c 0 respectively where c 0 can be an arbitrary positive constant independent of α .
For the case of α < c 0 group ⌈ c 0 α ⌉ 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.
J.9. Computing the Moments for Phase
In this subsection, we compute the first and second moments for the movement of manifold projection after R grp rounds of Local SGD.
Since the randomness in training might drive the iterate out of the working zone, making the dynamic intractable, we analyze a more well-behaved sequence { ˆ θ ( s ) k,t : ( s, t ) ⪯ ( R tot , 0) , k ∈ [ K ] } which is equal to { ˆ θ ( s ) k,t } with high probability.
Specifically, ˆ θ ( s ) k,t equal to θ ( s ) k,t if the global step ( s, t ) is η 100 -good and is set as a point ϕ null ∈ Γ otherwise.
Denote by E ( s ) t the event { global step ( s, t ) is η 100 -good } .
Then ˆ θ ( s ) k,t
= θ ( s ) k,t 1 E ( s ) t + ϕ null 1 ¯ E ( s ) t and { ˆ θ ( s ) k,t } satisfy the following update rule:
[Equation]
By Theorem 41, with probability at least 1 − η 100 , ˆ θ ( s ) k,t = θ ( s ) k,t , ∀ k ∈ [ K ] , ( s, t ) ⪯ ( R tot , 0) .
Similar to { θ ( s ) k,t } , we define the following variables with respect to { ˆ θ ( s ) k,t } :
[Equation]
k ∈ [
[Equation]
[Equation]
[Figure]
Notice that ˆ x ( s ) k, 0 = ˆ x ( s ) avg , 0 for all k ∈
[ K ] .
Finally, we introduce the following mapping Ψ ( θ ) :
Γ → R d × d , which is closely related to (cid:98) Ψ defined in Theorem 4.
Definition 46 For θ ∈ Γ , we define the mapping Ψ ( θ ) :
Γ → R d × d :
[Equation]
where λ i , v i are the i -th eigenvalue and eigenvector of ∇ 2 L ( θ ) and v i ’s form an orthonormal basis of R d .
Additionally, ψ ( x ) := e − x − 1+ x x and ψ (0) = 0 ; see Figure 7 for a plot.
Remark 47 Intuitively, Ψ ( θ ) rescales the entries of Σ ( θ ) in the eigenbasis of ∇ 2 L ( θ ) .
When ∇ L ( θ ) =
diag( λ 1 , · · · , λ d ) ∈ R d × d , where λ i = 0 for all m < i ≤ d , Ψ( Σ 0 )
i,j = ψ ( ηH ( λ i + λ j ))Σ 0 ,i,j .
Note that Ψ ( θ ) can also be written as
[Equation]
where ⊕ denotes the Kronecker sum A ⊕ [Equation] =
A ⊗ I
[Equation]
+ I d ⊗ B , vec( · ) is the vectorization operator of a matrix and ψ ( · ) is interpreted as a matrix function.
Now we are ready to present the result about the moments of ˆ ϕ
( s + R grp ) − ˆ ϕ ( s ) .
Theorem 48 For s 0 + s 1 ≤ s ≤ R tot − R grp and 0 < β < 0 .
5 , the first and second moments of ˆ ϕ
( s + R grp )
− ˆ ϕ ( s ) are as follows:
[Equation]
where ˜ O ( · ) hides log terms and constants independent of η .
Remark 49 By Theorem 41 and the definition of ˆ θ ( s ) k,t , (63) and (64) still hold when we replace ˆ ϕ ( s ) with ϕ ( s ) and replace ˆ ϕ ( s + R grp ) with ϕ ( s + R grp ) .
We shall have Theorem 48 if we prove the following theorem, which directly gives Theoremwith a simple shift of index.
For brevity, denote by ∆ ˆ ϕ ( s ) :
=
ˆ ϕ ( s ) − ˆ ϕ (0) , Σ 0 :
= Σ ( ˆ ϕ (0) ) , Σ 0 , ∥ := Σ ∥ ( ˆ ϕ (0) ) .
Theorem
Given ∥ ˆ θ (0)avg − ˆ ϕ (0) ∥ 2 = O ( (cid:113) η log 1 η ) , for 0 < β < 0 .
5 , the first and second moments of ∆ ˆ ϕ
( R grp ) are as follows:
[Equation]
We will prove Theorem 50 in the remainder of this subsection.
For convenience, we introduce more notations that will be used throughout the proof.
Let H 0 :
= ∇ 2 L ( ˆ ϕ (0) ) .
By Assumption 3.2, rank( H 0 ) = m .
WLOG, assume H 0 =
diag( λ 1 , · · · , λ d ) ∈ R d × d , where λ i = 0 for all m < i ≤ d and λ 1 ≥ λ 2 · · · ≥ λ m .
By Theorem 17, ∂ Φ( ˆ ϕ (0) ) is the projection matrix onto the tangent space
T ˆ ϕ (0) (Γ) (i.e. the null space of ∇ 2 L ( ˆ ϕ (0) ) ) and therefore, ∂ Φ( ˆ ϕ (0) ) =
(cid:20)
[Equation]
I d − m (cid:21) .
Let
P ∥ := ∂ Φ( ˆ ϕ (0) ) and P ⊥ :
[Equation]
Let ˆ A ( s ) avg :=
E [ˆ x ( s ) avg ,H ˆ x ( s ) ⊤ avg ,H ] , ˆ q ( s ) t := E
[ˆ x ( s ) k,t ] and ˆ B ( s ) t :=
E [ˆ x ( s ) k,t ∆ ˆ ϕ ( s ) ] .
The latter two notations are independent of k since for all k ∈ [ K ] , ˆ θ ( s ) k,t ’s are identically distributed.
The following lemma computes the first and second moments of the one round movement.
Lemma
Given ∥ ˆ θ (0)avg − ˆ ϕ (0) ∥ 2 = O ( (cid:113) η log 1 η ) , for 0 ≤ s < R grp , the first and second moments of ˆ ϕ ( s +1)
− ˆ ϕ
( s ) are as follows:
[Equation]
[Equation]
Proof By Taylor expansion
[Equation]
we have
[Equation]
Rearrange the terms and we obtain:
[Equation]
[Equation]
Moreover,
[Equation]
Noticing that ˆ x ( s ) k
[Equation]
H ∆ ˆ ϕ ( s ) ⊤ are identically distributed for all k ∈
[ K ]
[Equation]
we have E [ˆ x ( s ) avg ,H ∆ ˆ ϕ ( s ) ⊤ ] = 1 K (cid:80) k ∈
[ K ] E [ˆ x ( s ) k,
H ∆ ˆ ϕ ( s ) ⊤ ]
= ˆ B ( s ) H .
Then taking expectation of both sides of (67) gives
E
[Equation]
Again taking expectation of both sides of (68) yields
[Equation]
By Lemmas 37 and
[Equation]
the following holds simultaneously with probability at least 1 − η 100 :
[Equation]
Furthermore, since for all k ∈ [ K ] and ( s, t ) ⪯ ( R tot , 0) , ˆ θ ( s ) k,t stays in Γ ϵ 2 which is a bounded set, ∥ ∆ ˆ ϕ ( s ) ∥ 2 and ∥ ˆ x ( s ) avg ,H ∥ 2 are also bounded.
Therefore, we have
[Equation]
which concludes the proof.
We compute ˆ A ( s ) avg , ˆ q ( s ) t and ˆ B ( s ) t by solving a set of recursions, which is formulated in the following lemma.
Additionally, define ˆ A ( s ) t
[Equation]
Lemma
Given ∥ ˆ θ (0)avg − ˆ ϕ (0) ∥ 2 = O ( (cid:113) η log 1 η ) , for 0 ≤ s < R grp and 0 ≤ t < H , we have the following recursions.
[Equation]
Moreover,
[Equation]
Proof We first derive the recursion for ˆ q ( s ) t .
Recall the update rule for ˆ θ ( s ) k,t :
[Equation]
Subtracting ˆ ϕ ( s ) from both sides gives
[Equation]
where the second and third equality perform Taylor expansion.
Taking expectation on both sides gives
[Equation]
By Theorem 41, with probability at least 1 − η 100 , ˆ e ( s ) k,t = 0 , ∀ k ∈ [ K ] , ( s, t ) ⪯ ( R grp , 0) .
Also notice that both ˆ θ ( s ) k,t and ϕ null belong to the bounded set Γ ϵ 2 .
Therefore, ∥ ˆ e ( s ) k,t ∥ 2 is bounded and we have E [ ∥ ˆ e ( s ) k,t ∥ 2 ]
=
O ( η 100 ) .
Combining this with (69) to (71) yields (72).
Secondly, we derive the recursion for ˆ B ( s ) t .
Multiplying both sides of (80) by ∆ ˆ ϕ ( s ) ⊤ and taking expectation, we have
[Equation]
Still by Theorem 41 and (69) to (71), we have (75).
Thirdly, we derive the recursion for ˆ A ( s ) t .
By (80), we have
[Equation]
which establishes (73).
Fourthly, we derive the recursion for ˆ M ( s ) t .
Multiplying both sides of (80) by ˆ x ( s ) l,t +1 and taking expectation, l ̸
= k , we obtain
[Equation]
By a similar argument to the proof of Theorem
[Equation]
we have
[Equation]
which yields (74)
[Equation]
Now we proceed to prove (76) to (79).
By definition of ˆ A ( s ) avg ,
[Equation]
which demonstrates (76)
Then we derive (77).
By definition of ˆ x ( s +1) avg , 0 ,
[Equation]
Hence
[Equation]
By (69) and (71)
[Equation]
we obtain (77).
By (67),
[Equation]
Combining (81) and (82) gives
[Equation]
Therefore,
[Equation]
Finally
[Equation]
we apply Theorem 51 to derive (78).
[Equation]
which concludes the proof
[Equation]
With the assumption that the hessian at ˆ ϕ (0) is diagonal
[Equation]
we have the following corollary that formulates the recursions for each matrix element
[Equation]
Corollary 53 Given ∥ ˆ θ (0)avg − ˆ ϕ (0) ∥ 2 =
O ( (cid:113) η log 1 η ) , for 0 ≤ s < R grp and 0 ≤ t <
H
[Equation]
we have the following elementwise recursions.
[Equation]
otherwise .
[Equation]
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 ∥ MP ∥ in the tangent space of Γ at ˆ ϕ (0) , P ⊥ MP ⊥ in the normal space, along with P ∥ MP ⊥ and
P ⊥ MP ∥ across both spaces.
We first compute the elements of P ⊥ ˆ A ( s ) t P ⊥ and P ⊥
ˆ A ( s )
avg P ⊥ .
Lemma 54 (General formula for P ⊥ ˆ A ( s ) t P ⊥ and P
⊥
ˆ
A ( s ) avg P ⊥ )
Let R 0 :
= ⌈ 10 λ m α log 1 η ⌉ .
Then for 1 ≤ i ≤ m, 1 ≤ j ≤ m and R 0 ≤ s < R grp ,
[Equation]
For s < R 0 , ˆ A ( s ) t,i,j = ˜
O ( η ) and ˆ A ( s ) avg , , i,j = ˜
O ( η ) .
Proof For 1 ≤ i ≤ m, 1 ≤ j ≤ m , λ i > 0 , λ j > 0 .
By
[Equation]
where the second inequality uses (cid:80) t −
[Equation]
where the second equality uses M ( s +1)= A ( [Equation] +
[Equation]
By (86) and (87),
[Equation]
η
[Equation]
Then we obtain
[Equation]
Notice that | 1 − ( λ i + λ j ) η | < 1 and
[Equation]
i
[Equation]
Therefore, s [Equation]
[Equation]
Then we have
[Equation]
Finally, we demonstrate that for s ≥ R 0 , ˆ A ( s ) 0 ,i,j and ˆ A ( s ) avg ,i,j is approximately equal to
[Equation]
KB loc Σ 0 ,i,j .
By (89), when s ≥ R 0 , (1 − ( λ i + λ j ) η ) sH =
O ( η 10 ) , which gives
[Equation]
For s < R 0 , since ˆ A (0)
=
ˆ x ( s ) avg , 0 ˆ x ( s ) ⊤ avg , 0 =
˜ O ( η ) , we have ˆ A ( s ) avg , , i,j = ˜
O ( η ) and ˆ A ( s ) t,i,j = ˜
O ( η ) .
Secondly, we compute P ∥
ˆ A ( s ) t P ⊥ and P ∥
ˆ A ( s )
avg P ⊥ .
Lemma 55 (General formula for P ⊥
ˆ A ( s ) t P ∥ and P
⊥ ˆ A ( s ) avg P ∥ )
For 1 ≤ i ≤ m, m < j ≤ d ,
[Equation]
Proof Note that for 1 ≤ i ≤ m, m < j ≤ d and λ i > 0 , λ j = 0 .
By (83) and (87),
[Equation]
By (84) and (87), ˆ M ( s ) t,i,j = ˜
O ( η 1 . 5 − 0 .
5 β ) .
Then,
[Equation]
Similar to Theorem 55, we have the following lemma for the general formula of P ∥
ˆ A ( s ) t P ⊥ and
P ∥
ˆ A ( s )
avg P ⊥ .
Lemma 56 (General formula for P ∥
ˆ A ( s ) t P ⊥ and P ∥
ˆ A ( s ) avg P ⊥ )
For m < i ≤ d and 1 ≤ j ≤ m ,
[Equation]
Finally, we derive the general formula for P ∥
ˆ A ( s ) t P ∥ and P ∥
ˆ A ( s ) avg P ∥ .
Lemma 57 (General formula for P ∥
ˆ A ( s ) t P ∥ and P ∥
ˆ A ( s ) avg P ∥ )
For m < i ≤ d and m <
j ≤ d ,
[Equation]
Proof Note that for m < i ≤ d , m < j ≤ d and λ i = λ
j = 0 .
(83) is then simplified as
[Equation]
Therefore,
[Equation]
According to (84), ˆ M ( s ) t,i,j = ˜
O ( η 1 . 5 − 0 .
5 β ) for m < i ≤ d and m < j ≤ d .
Combining (84), (87) and (90) yields
[Equation]
Now, we move on to compute the general formula for ˆ B ( s ) t .
Lemma
(The general formula for P ⊥ ˆ B ( s ) t P ∥ ) Note that for 1 ≤ i ≤ m and m < j ≤ d , when R 0 :=
[Equation]
For s < R 0 , ˆ B ( s ) t,i,j = ˜
O ( η ) .
Proof Note that for 1 ≤ i ≤ m , λ i > 0 .
By (85),
[Equation]
Hence,
[Equation]
According to (88),
[Equation]
Then we have
[Equation]
where the second equality uses (89) and the last inequality uses ˆ B (0)
=
ˆ x (0)avg , 0 ∆ ˆ ϕ (0)
= 0 .
For s
[Equation]
[Equation]
[Equation]
which gives
[Equation]
Therefore
[Equation]
For s < R 0 , ˆ A ( s ) avg , , i,j = ˜
O ( η ) and therefore, ˆ B ( s ) t,i,j = ˜
O ( η ) .
Lemma 59 (General formula for the elements of P ⊥ ˆ B ( s ) t P ⊥ )
For 1 ≤ i ≤ m and 1 ≤ j ≤ m , , ˆ B
[Equation]
Proof Note that for 1 ≤ i ≤ m , λ i > 0 .
By (85),
[Equation]
Hence
[Equation]
By (88),
[Equation]
where the last inequality uses ˆ B (0) = 0 .
Lemma 60 (General formula for P ∥
ˆ B ( s ) t )
For m < i ≤ d , ˆ B ( s ) t,i,j = ˜
O ( η 1 .
5 − β ) .
Proof Note that λ i = 0 for m < i ≤ d .
By (85) and (88),
[Equation]
Therefore,
[Equation]
Having obtained the expressions for ˆ B ( s ) t , ˆ A ( s ) t and ˆ A ( s ) avg , we now provide explicit expressions for the first and second moments of one round movements in the following two lemmas.
Lemma
The expectation of one round movement is
E [ ˆ ϕ ( s +1)
[Equation]
where R 0 :
= ⌈ 10 λ m α log 1 η ⌉ .
Proof We first compute E [ ˆ ϕ ( s +1)
− ˆ ϕ
( s ) ] .
By (65), we only need to compute P ∥
ˆ q ( s ) H by relating it to these matrices.
Multiplying both sides of (72) by P ∥ gives
[Equation]
Similarly, according to (78), we have
[Equation]
Combining (92) and (93) yields
[Equation]
By Lemmas 54, 57 and 55, for s ≤ R
=
⌊ 10 λα log 1 η ⌋ , ˆ A ( s ) t
[Equation]
[Equation]
˜
[Equation]
η ) and ˆ B ( [Equation] ) t = ˜
O ( η ) .
Therefore, E [ ˆ ϕ ( s +1) − ˆ ϕ ( s ) ]
=
˜ O ( η ) .
For s > R 0 , ˆ A ( s − 1) avg
[Equation]
Substituting (94) into (65) gives
[Equation]
Below we compute T 1 and T 2 for s > R 0 respectively.
By Theorem 18,
[Equation]
By Theorem
[Equation]
Therefore, for s > R 0 ,
[Equation]
where we apply Theorem 57.
Similarly, for s > R 0 ,
[Equation]
where we apply Theorem 60.
Hence,
[Equation]
We move on to show that
[Equation]
Similar to the way we compute ˆ A ( s ) t , ˆ A ( s ) avg and ˆ B ( s ) t , we compute T 2 by splitting T 3 into four matrices and then substituting them into the linear operator − η P ∥ ∇ 3 L ( ˆ ϕ (0) )
[ · ] one by one.
First, we show that
[Equation]
where ψ ( · ) is interpreted as an elementwise matrix function here.
By Lemmas 54 and 59, for 1 ≤ i ≤ m , 1 ≤ j ≤ m and s > R 0 ,
[Equation]
Therefore,
[Equation]
Then we simplify T 4 .
Notice that
[Equation]
Therefore,
[Equation]
Substituting T 4 back into the expression for (cid:80) H − 1 t =
ˆ A ( s ) t,i,j gives
[Equation]
Combining the elementwise results
[Equation]
we obtain the following matrix form expression:
[Equation]
By Theorem
[Equation]
we have (97)
[Equation]
Secondly, we show that for s > R 0 ,
[Equation]
where ψ ( · ) is interpreted as an elementwise matrix function here.
By symmetry of ˆ A ( s ) t ’s and ∇
[Equation]
Therefore
[Equation]
we only have to evaluate
[Equation]
To compute the elements of (cid:80) H − 1 t =
P ⊥ ( ˆ A ( s ) t
+ ˆ B ( s ) t ) P ∥
[Equation]
we combine Lemmas 55 and 58 to obtain that for 1 ≤ i ≤ m and m < j ≤ d ,
[Equation]
and
H
[Equation]
[Equation]
=
[Equation]
Therefore, the matrix form of (cid:80) H − 1 t =
P ⊥ ( ˆ A ( s ) t
+ ˆ B
[Equation]
P ∥ is
H − 1 (cid:88) t
[Equation]
where ψ ( · ) is interpreted as an elementwise matrix function here.
Furthermore, by Theorem 60, (cid:80) H − 1 t =
[Equation]
[Equation]
− [Equation]
Applying Theorem 18, we have (98).
Finally, directly applying Theorem 20, we have
[Equation]
[Equation]
Notice that ψ ( Σ 0 , ∥ ) = 0 where ψ ( · ) operates on each element.
Combining (97), (98) and (99), we obtain (96).
By (95) and (96), we have (91).
Lemma
The second moment of one round movement is
[Equation]
where R 0 :
= ⌈ 10 λ m α log 1 η ⌉ .
Proof Directly apply Theorem 57 and Theorem 51 and we have the lemma.
With Lemmas 61 and 62, we are ready to prove Theorem 50.
Proof [Proof of Theorem 50.
We first derive E [∆ ˆ ϕ ( R grp ) ] .
Recall that R grp =
⌊ 1 αη β ⌋
=
1 Hη 1+ β + o (1)
where 0 < β < 0 . 5 .
By Theorem 61,
[Equation]
Then we compute E [∆ ˆ ϕ ( R grp )
∆ ˆ ϕ ( R grp ) ⊤ ] .
[Equation]
where the last inequality uses E [( ˆ ϕ ( s +1) − ˆ ϕ ( s ) )
E [( ˆ ϕ ( s ′ +1) − ˆ ϕ ( s ′ ) )
⊤ ] =
˜ O ( η 2 ) .
J.10. Proof of Weak Approximation
We are now in a position to utilize the estimate of moments obtained in previous subsections to prove the closeness of the sequence { ϕ ( s ) }
⌊ T/ ( Hη 2 ) ⌋ s = and the SDE solution { ζ : t ∈ [0 , T ] } in the sense of weak approximation.
Recall the SDE that we expect the manifold projection { Φ( ¯ θ ( s ) ) }
⌊ T/ ( Hη 2 )
⌋ s = to track:
[Equation]
According to Theorem 18 and Theorem 19, the drift term in total can be written as the following form:
[Equation]
Then by definition of P ζ , (100) is equivalent to the following SDE:
[Equation]
Therefore, we only have to show that ϕ ( s ) closely tracks { ζ ( t ) } satisfying Equation (101).
By Theorem 14, there exists an ϵ 3 neighborhood of Γ , Γ ϵ 3 , where Φ( · ) is C ∞ -smooth.
Due to compactness of Γ , Γ ϵ 3 is bounded and the mappings ∂ 2 Φ( · ) , ∂ Φ( · ) , Σ 1 / 2 ( · ) , Σ ( · ) and Ψ ( · ) are all Lipschitz in Γ ϵ 3 .
By Kirszbraun theorem, both the drift and diffusion term of (101) can be extended as Lipschitz functions on R d .
Therefore, the solution to the extended SDE exists and is unique.
We further show that the solution, if initialized as a point on Γ , always stays on the manifold almost surely.
As a preparation, we first show that Γ has no boundary.
Lemma 63 Under Assumptions 3.1 to 3.3, Γ has no boundary.
Proof We prove by contradiction.
If Γ has boundary ∂ Γ , WLOG, for a point p ∈ ∂ Γ , let the
Hessian at p be diagonal with the form ∇ 2 L ( p ) =
diag( λ 1 , · · · , λ d ) where λ i > 0 for 1 ≤ i ≤ m and λ i = 0 for m < i ≤ d .
Denote by x i : j :=
( x i , x i +1 , · · · , x j ) ( i ≤ j ) the ( j − i + 1) -dimensional
vector formed by the i -th to j -th coordinates of x .
Since ∂ ( ∇L ( p ))
∂ p 1: m = diag( λ 1 , · · · , λ m ) is invertible, by the implicit function theorem, there exists an open neighborhood V of p m +1:
d such that ∇L ( v ) = 0 , ∀ v ∈ V .
Then, L ( v ) =
L ( p ) =
min θ ∈ U L ( θ ) and hence V ⊂ Γ , which contradicts with p ∈ ∂ Γ .
Therefore, Γ is a closed manifold (i.e., compact and without boundary).
Then we have the following lemma stating that Γ is invariant for (101).
Lemma 64 Let ζ ( t ) be the solution to (101) with ζ (0) ∈ Γ , then ζ ( t ) ∈ Γ for all t ≥ 0 .
In other words, Γ is invariant for (101) .
Proof According to Filipovi´c [12] and Du and Duan [9], for a closed manifold M to be viable for the SDE d X ( t )
=
F ( X ( t ))
d t + B
( X ( t ))
d W t where F : R d → R d and B : R d → R d are locally
Lipschitz, we only have to verify the following Nagumo type consistency condition:
[Equation]
where D[ · ] is the Jacobian operator and B j ( x ) denotes the j -th column of B ( x ) .
In our context, since for ϕ ∈ Γ , ∂ Φ( ϕ ) is a projection matrix onto T ϕ (Γ) , each column of ∂ Φ( ϕ ) Σ 1 / 2 ( ϕ ) belongs to T ϕ (Γ) , verifying the second condition.
Denote by P ⊥ ( ϕ ) :
= I d − ∂ Φ( ϕ )
the projection onto the normal space of Γ at ϕ .
To verify the first condition, it suffices to show that P ⊥ ( ϕ ) µ ( ϕ ) = 0 .
We evaluate (cid:80) j P ⊥ ( ϕ )D
[ B j ( ϕ )]
B j ( ϕ ) as follows.
[Equation]
where the last inequality uses Theorem 18.
Again applying Theorem 18, we have
[Equation]
Combining (102) and (103), we can verify the first condition.
In order to establish Theorem 4, it suffices to prove the following theorem, which captures the closeness of ϕ ( s ) and ζ ( t ) every R grp rounds.
Theorem 65 If ∥ ¯ θ (0) − ϕ (0) ∥ 2 = O ( (cid:113) η log 1 η ) and ζ (0) = ϕ (
) ∈ Γ , then for R grp =
⌊ 1 αη 0 .
75 ⌋ every test function g ∈ C 3 ,
[Equation]
where C g > 0 is a constant independent of η but can depend on g ( · ) and b > 0 is a constant independent of η and g ( · ) .
J.10.1. P RELIMINARIES AND ADDITIONAL NOTATIONS
We first introduce a general formulation for stochastic gradient algorithms (SGAs) and then specify the components of this formulation in our context.
Consider the following SGA:
[Equation]
where x n ∈ R
d is the parameter, η e is the learning rate, h ( · , · ) is the update which depends on x n and a random vector ξ n sampled from some distribution Ξ( x n ) .
Also consider the following
Stochastic Differential Equation (SDE).
[Equation]
where b ( · ) : R d → R d is the drift function and σ ( · ) : R d × d → R d × d is the diffusion matrix.
Denote by P X ( x , s, t ) the distribution of X ( t ) with the initial condition X ( s ) = x .Define
[Equation]
where X ( n +1) η e ∼ P X ( x , nη e , ( n + 1) η e ) , which characterizes the update in one step.
In our context, we view the movement of manifold projection over R grp :=
[Equation] β [Equation] ( β ∈ (0 , 0 . 5)) rounds as one update step.
Hence the ϕ ( nR grp ) corresponds to the discrete time random variable x n corresponds to and ζ ( t ) corresponds to the continuous time random variable X t .
According to Theorem 48, we set
[Equation]
Due to compactness of Γ , b ( · ) and σ ( · ) are Lipschitz on Γ .
As for the update in one step, ˜ ∆ ( · , · ) is defined in our context as:
[Equation]
where ζ ( n +1) η e ∼ P ζ ( ϕ , nη e , ( n + 1) η e ) and ϕ ∈ Γ .
For convenience, we further define
[Equation]
We use C g,i to denote constants that can depend on the test function g and independent of η e .
The following lemma relates the moments of ˜ ∆ ( ϕ , n ) to b ( ϕ ) and σ ( ϕ ) .
Lemma 66 There exists a positive constant C 0 independent of η e and g such that for all ϕ ∈ Γ ,
[Equation]
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 (101) always stays on Γ almost surely if its initial value ζ (0) belongs to Γ , (ii) b ( · ) and σ ( · ) are C ∞ and (iii) Γ is compact, we can directly apply Lemma B.3 in [42] and Lemma 26 in [34] to obtain the above lemma.
The following lemma states that the expectation of g ( ζ ( t )) for g ∈ C 3 is smooth with respect to the initial value of the SDE solution.
Lemma 67 Let s ∈ [0 , T ] , ϕ ∈ Γ and g ∈ C 3 .
For t ∈ [ s, T ] , define
[Equation]
Then u ( · , s, t ) ∈ C 3 uniformly in s, t .
Proof A slight modification of Lemma B.4 in [42] will give the above lemma.
J.10.2. P ROOF OF THE APPROXIMATION IN OUR CONTEXT
For β ∈ (0 , 0 . 5) , define γ 1 := 1 . 5 − 2 β 1 − β , γ 2 :=
1 1 − β , and then 1 < γ 1 < 1 . 5 , 1 <
γ 2 < 2 .
We introduce the following lemma which serves as a key step to control the approximation error.
Specifically, this lemma bounds the difference in one step change between the discrete process and the continuous one as well as the product of higher orders.
Lemma
If ∥ ¯ θ (0) − ϕ (0) ∥ 2 =
O ( (cid:113) η log 1 η ) , then there exist positive constants C 1 and b independent of η e and g such that for all 0 ≤ n < ⌊ T/η e ⌋ , 1.
[Equation] [Equation]
Proof According to Appendix J
[Equation]
[Equation]
we have
[Equation]
Since γ 1 < 1 . 5 and γ 2 <
[Equation]
we can utilize Theorem 50 and conclude that there exist positive constants C 2 and b independent of η e and g such that
[Equation]
Combining (104) - (106) with Theorem 66 gives the above lemma
[Equation]
Lemma
For a test function g ∈ C 3 , let u l,n ( ϕ ) :
= u ( ϕ , lη e , nη e ) =
E ζ t ∼P ζ ( ϕ ,lη e ,nη e ) [ g ( ζ t )] .
If ∥ ¯ θ (0) − ϕ (0) ∥ 2 =
O ( (cid:113) η log 1 η ) , then for all 0 ≤ l ≤ n − 1 and 1 ≤ n ≤ ⌊ T/η e ⌋ ,
[Equation]
where C g, 1 is a positive constant independent of η and ˆ ϕ ( lR grp ) but can depend on g .
Proof By Theorem 67, u l,n ( ϕ ) ∈ C 3 for all l and n .
That is, there exists K (
· ) ∈ G such that for all l, n , u l,n ( ϕ ) and its partial derivatives up to the third order are bounded by K ( ϕ ) .
By the law of total expectation and triangle inequality,
[Equation]
We first bound A 2 and A 3 .
Since ˆ ϕ ( lR grp ) ∈ Γ , both ˆ ϕ ( lR grp ) + ∆ ( l ) and ˆ ϕ ( lR grp ) + ˜ ∆ ( l ) belong to Γ .
Due to compactness of Γ and smoothness of u l +1 ,n ( · ) on Γ , there exist a positive constant C g,such that A 2 + A 3 ≤ C g, 2 η 100 .
We proceed to bound A 1 .
Expanding u l +1 ,n ( · ) at ˆ ϕ ( lR grp ) and by triangle inequality,
[Equation]
where the remainders R and ˜ R are
[Equation]
] , for some θ, ˜ θ ∈ (0 , 1) .
Since ˆ ϕ ( lR grp ) belongs to Γ which is compact, there exists a constant C g,such that for all 1 ≤ i, j ≤ d, 0 ≤ l ≤ n − 1 , 1 ≤ n ≤ ⌊ T/η e ⌋ ,
[Equation]
By Theorem 68,
[Equation]
Now we bound the remainders.
By Cauchy-Schwartz inequality,
[Equation]
Since ˆ ϕ ( lR grp ) and ˆ ϕ ( lR grp ) + ∆ ( l ) both belong to Γ which is compact, there exists a constant
C g, 4 such that for all 1 ≤ i, j, p ≤ d , 0 ≤ l ≤ n − 1 and 1 ≤ n ≤ ⌊ T/η e ⌋ ,
[Equation]
Combining the above inequality with Theorem 68, we have
[Equation]
Hence, for all 1 ≤ n ≤ ⌊ T/η e ⌋ , 0 ≤ l ≤ n − 1 ,
[Equation]
Similarly, we can show that there exists a constant C g, 5 such that for all 1 ≤ n ≤ ⌊ T/η e ⌋ , 0 ≤ l ≤ n − 1 ,
[Equation]
Combining the bounds on A 1 to A 3 , we have the lemma.
Finally, we prove Theorem 65.
Proof For 0 ≤ l ≤ n , define the random variable ˆ ζ l,n which follows the distribution P ζ ( ˆ ϕ ( lR grp ) , l, n ) conditioned on ˆ ϕ ( lR grp ) .
Therefore, P (ˆ ζ n,n = ˆ ϕ ( nR grp ) )
= 1 and ˆ ζ 0 ,n ∼ ζ nη e .
Denote by u ( ϕ , s, t ) :
=
E ζ t ∼P ζ ( ϕ ,s,t ) [ g ( ζ t )] and T l +1 ,n := u l +1 ,n ( ˆ ϕ ( lR grp ) + ∆ ( l ) , ( l +1) η e , nη e ) − u l +1 ,n ( ˆ ϕ ( lR grp ) + ˜ ∆ ( l ) , ( l + 1) η e , nη e ) .
[Equation]
Noticing that E
[ T l +1 ,n | E ( nR grp )]
=
E [ E [ T
[Equation]
|
[Equation]
E
[Equation] grp )] | E ( nR grp )] [Equation] we can apply
Theorem 69 and obtain that for all 0 ≤ n ≤ ⌊ T/η e ⌋ ,
[Equation]
Notice that η γ 1 e + η γ 2 e = η 0 .
5 − β + η β
and T
, C g, 1 are both constants that are independent of η e .
Let β = 0 .
and we have Theorem 65.
Having established Theorem 65, we are thus led to prove Theorem 4.
Proof [Proof of Theorem 4]
Denote by s cls = s
+ s
=
O (log 1 η ) , which is the time the global iterate ¯ θ ( s ) will reach within ˜ O
( η ) from Γ with high probability.
Define ˜ ζ ( t ) to be the solution to the limiting SDE (101) conditioned on E ( s cls )and ˜ ζ (0) =
ϕ ( s cls ) .
By Theorem 65, we have
[Equation]
where R grp =
⌊ 1 αη 0 .
Noticing that (i) g ∈ C 3 (ii) b , σ ∈ C ∞ and (iii) ζ ( t ) , ˜ ζ ( t ) ∈ Γ , t ∈ [0 , ∞ ) almost surely, we can conclude that given E ( s cls )
[Equation]
Then there exists positive constant b ′ independent of η and g , and C ′ g which is independent of η but can depend on g such that
[Equation]
We can view the random variable pairs { ( ϕ ( nR grp + s cls ) , ζ nη 0 .
+ s cls αη ) :
n = 0 , · · · , ⌊ T/η 0 .
75 ⌋} as reference points and then approximate the value of g ( ϕ ( s ) ) and g ( ζ ( sHη 2 )) with the value at the nearest reference points.
By Lemmas 33 and 38, for 0 ≤ r ≤ R grp and 0 ≤ s ≤ R tot − r ,
[Equation]
Since the values of ϕ ( s ) and ζ are restricted to a bounded set, g ( · ) is Lipschitz on that set.
Therefore, we have the theorem.
Appendix K. Deriving the Slow SDE for Label Noise Regularization
In this this section, we formulate how label noise regularization works and derive the theoretical results in Appendix E.
Consider training a model for C -class classification on dataset D
= { ( x i , y i ) }
Ni =1 , where x i denotes the input and y i ∈
[ C ] denotes the label.
Denote by ∆ C − 1 + the ( C − 1) -open simplex.
Let f ( θ ; x ) ∈ ∆ C − 1 + be the model output on input x with parameter θ , whose j -th coordinate f j ( θ ; x ) stands for the probability of x belonging to class j .
Let ℓ ( θ ; x , y ) be the cross entropy loss given input x and label y , i.e, ℓ ( θ ; x , y )
=
− log f y ( θ ; x ) .
Adding label noise means replacing the true label y with a fresh noisy label ˆ y every time we access the sample.
Specifically, ˆ y is set as the true label y with probability 1 − p and as any other label with probability pC − 1 , where p is the fixed corruption probability.
The training loss is defined as
L ( θ ) =
1 N (cid:80) Ni =1 E [ ℓ ( θ ; x i , ˆ y i )] , where the expectation is taken over the stochasticity of ˆ y i .
Notice that given a sample ( x , y ) ,
[Equation]
By the property of cross-entropy loss, (107) attains its global minimum if and only if f j = pC − 1 , for all j ∈ [ C ] , j ̸ = y and f y = 1 − p .
Due to the large expressiveness of modern deep learning models, there typically exists a set S ∗ :
= { θ | f i ( θ ) =
E [ˆ y i ] , ∀ i ∈
[ N ] } such that all elements of S ∗ minimizes L ( θ ) .
Then, the manifold Γ is a subset of S ∗ .
The following theorem relates the noise covariance Σ ( θ ) :
= 1 N (cid:80) i ∈
[ N ] E [( ∇ ℓ ( θ ; x i , ˆ y i ) − ∇L ( θ )) ( ∇ ℓ ( θ ; x i , ˆ y i ) − ∇L ( θ ))
⊤ ] to the hessian ∇ 2 L ( θ ) for all θ ∈ S ∗ .
Theorem 70 If f ( θ ; x i , ˆ y i ) is C 2 -smooth on R d given any i ∈
[ N ] , ˆ y i ∈
[ C ] and S ∗ ̸ = ∅ , then for all θ ∈ S ∗ , Σ ( θ )
= ∇ 2 L ( θ ) .
Proof Since L ( · ) is C 2 -smooth, ∇L ( θ ) =
for all θ ∈ S ∗ .
To prove the above theorem, it suffices to show that ∀ i ∈
[ N ] , E [ ∇ ℓ ( θ ; x i , ˆ y i ) ∇ ℓ ( θ ; x i , ˆ y i ) ⊤ ] = ∇ 2 L ( θ ) .
W.L.O.G, let y =
and therefore for all θ ∈ S ∗
[Equation]
[Equation]
Additionally, let h ( x ) :=
− log( x ) , x ∈ R + .
The stochastic gradient ∇ ℓ ( θ ; x , ˆ y ) follows the distribution
[Equation]
Then the covariance of the gradient noise is
[Equation]
Now we compute the hessian.
[Equation]
Since (cid:80) j ∈ [ C ] f i = 1 ,
[Equation]
Also, notice that h ′ ( x ) = − 1 x .
Therefore,
[Equation]
Substituting (108) and (109) into the expression of T gives T = 0 , which simplifies ∇ 2 L ( θ ) as the following form:
[Equation]
Again notice that h ′′ ( x ) = h ′ ( x ) for all x ∈ R
+ .
Therefore, ∇ 2 L ( θ ) =
Σ ( θ ) .
With the property Σ ( θ ) = ∇ 2 L ( θ ) , the limiting SDE (4) can be greatly simplified.
Corollary 71 (Slow SDE for Local SGD with label noise)
For C -class classification task with crossentropy loss, the slow SDE of Local SGD with label noise has the following form:
[Equation]
where F ( x ) := (cid:82) x 0 ψ ( y )
d y and is interpreted as a matrix function in (110) .
Additionally, ∇ Γ f stands for the gradient of a function f projected to the tangent space of Γ .
Proof Recall the general form of the slow SDE for Local SGD:
[Equation]
where Ψ ( ζ ) is defined in Definition 46.
Since for ζ ∈ Γ , Σ ( ζ ) = ∇ 2 L ( ζ ) , then
[Equation] [Equation]
Now we show that
[Equation]
Since ∇ 2 L ( ζ ) = Σ ( ζ ) , V ∇ 2 L ( ζ ) [ Σ ] = 12 I .
By Theorem 19,
[Equation]
Finally, we show that
[Equation]
Define ˆ ψ ( x ) :
= xψ ( x )
=
e − x − 1 + x .
By definition of Ψ( ζ ) , when Σ ( ζ ) = ∇ 2 L ( ζ ) , Ψ( ζ )
=
ˆ ψ (2 ηH ∇ 2 L ( ζ )) , where ˆ ψ ( · ) is interpreted as a matrix function.
Since ψ (2 ηH ∇ 2 L ( ζ )) ∈ span { uu ⊤ | u ∈ T ⊥ ζ (Γ) } , by Theorem 19,
[Equation]
By the chain rule, we have (114).
Combining (112),(113) and (114) gives the corollary.
We further have the following corollary as H goes to infinity.
Lemma
As the number of local steps H goes to infinity, the slow SDE of Local SGD with label noise (110) can be simplified as:
[Equation]
Proof We obtain the corollary by simply taking the limit.
By L’Hospital’s rule,
[Equation]
Therefore,
[Equation]
Substituting (116) into (110)yields (115).
Appendix L. Experimental Details
In this section, we specify the experimental details that are omitted in the main text.
Our experiments are conducted on CIFAR-10 [31] and ImageNet [51].
Our implementation of ResNet-56 [18] and
VGG-16 [54] is based on the high-starred repository by Wei Yang
and we use the implementation of ResNet-50 from torchvision 0.3.1.
We run all CIFAR-10 experiments with B loc = 128 on
NVIDIA Tesla P100 GPUs while ImageNet experiments are run on 8 NVIDIA A100 GPUS with
B loc = 32 .
All ImageNet experiments are 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.
[24] 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 B loc to estimate the running mean and variance before evaluation.
Experiments on both datasets follow the standard data augmentation pipeline in
He et al.
[18] except the label noise experiments.
Additionally, we use FFCV [32] to accelerate data loading for ImageNet training.
Slightly different from
the update rule of Local SGD in Section 1, we use the following sampling scheme unless otherwise stated.
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 own 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 to form a batch.
Then another epoch starts with the dataset reloaded to the initial state.
This sampling scheme is standard in practice and is also adopted by Lin et al.
[39] and
Goyal et al.
L.1.
Post-local SGD Experiments in Section
CIFAR-10 experiments.
We simulate 32 clients with B = 4096 .
We follow the linear scaling rule and linear learning rate warmup strategy suggested by Goyal et al.
We first runepochs 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 η = 0 .
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 × 10 − 4 .
As for the initialization scheme, we follow Lin et al.
[39] and Goyal et al.
Specifically, we use Kaiming Normal [17] 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 .
The weights for normalization layers are initialized as 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.
We first runepochs 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 at epoch 50 .
Resuming from epoch 100 , we run Local SGD with η = 0 .
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 ( 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 1 × 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 3 runs.
L.2. Experimental Details for Sections 2 and B.
CIFAR-10 experiments.
We use ResNet-56 for all CIFAR-10 experiments in the two sections.
We simulate 32 workers with B = 4096 and set the weight decay as 5 × 10 − 4 .
For Figure 2 (a) and (b), we set η = 0 .
32 , which is the same as the learning rate after decay in Figure 1 (a).
For
Figure 2 (a), we adopt the same initialization scheme introduced in the corresponding paragraph in
Appendix L.1.
For Figure 2 (b), (e) and Figure 3 (c), we use the model at epoch 250 in Figure 1 (a) as the pre-trained model.
Additionally, we use a training budget of 250 epochs for Figure 2 (e).
In
Figure 3 (e), 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 rate ramps up from 0 .
05 to 1 .
6 linearly in the first 150 epochs.
Note that we conduct grid search for the 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 in the middle of the set.
Resuming from epoch 150 , we run Local SGD H = 1 (i.e., SGD) and 24 with η = 0 .
16 and decay η by 0 .
1 at epoch 226 .
For
Local SGD H = 900 , we resume from the model at epoch 226 of H = 24 with η = 0 .
We report the mean and standard deviation over 3 runs for Figure 2 (a) (b) and Figure 3 (c), and overruns for Figure 2 (e).
ImageNet experiments.
We simulate 256 clients with B = 8192 and set the weight decay as 1 × 10 − 4 .
In Figure 2 (d), both Local SGD and SGD start from the same random initialization.
We warm up the learning rate from 0 .
1 to 3 .
2 in the first 5 epochs and decay the learning rate by a factor of 0 .
1 at epochs 50 and 100 .
For Figure 2 (c) (f) and Figure 3 (d), we use the model at epoch 100 in
Figure 1 (b) as 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 1 (b).
Finally, in Figure 2 (f) and Figure 3 (d), we report the mean and average over 3 runs.
L.3. Experiments on Reducing the Diffusion Term
CIFAR -10 experiments.
The model we use is ResNet-56.
We first run SGD with batch size 128 and learning rate η = 0 .
5 for 250 epochs to obtain the pre-trained model.
The initialization scheme is the same as the corresponding paragraph in Appendix L.1.
Resuming from epochwith η = 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.
We use the model at epoch 100 in Figure 1 (b) as the pre-trained model.
Resuming from epoch 100 with η = 0 .
032 , we run Local SGD with K = 256 until epoch 250 and run all other setups for the same number of iterations.
L.4. Local SGD with Label Noise Regularization
For the label noise experiments, we do not use data augmentation and use sampling with replacement.
We simulate 32 clients with B = 4096 and set the corruption probability as 0 . 1 .
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 [49] 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 p = 0 .
1 for 500 epochs ( 6000 iterations).
We initialize the model by the same strategy introduced in the first paragraph pf Appendix L.1.
Applying the linear warmup scheme proposed by Goyal et al.
[15], we gradually ramp up the learning rate η from 0 .
1 to 3 .
2 for the first 50 epochs and multiply the learning rate by 0 .
1 at epoch 250 .
All subsequent experiments in Figure 5 (a) use learning rate 0 .
The weight decay λ is set as 5 × 10 − 4 .
Note that adding weight decay in the presence of normalization accelerates the limiting dynamics and will not affect the implicit regularization on the original loss function [35].
VGG-16.
We follow Yang’s implementation of the model architecture except that we replace maximum with average pooling and use Swish activation [49] to make the training loss smooth.
We initialize all weight parameters by Kaiming Normal and all bias parameters as zero.
The pre-trained model is obtained by running label noise SGD with total batch size 4096 and corruption probability p = 0 .
1 for 6000 iterations.
We use a linear learning rate warmup from 0 .
1 to 0 .
5 in the firstiterations.
Resuming from the model obtained by SGD, we use learning rate η = 0 .
The weight decay λ is set as zero.