Title: Global Convergence Analysis of Local SGD for Two-layer Neural Network without Overparameterization
Venue: NeurIPS 2023
Accepted: yes
Keywords: convolutional neural network

gaussian input

local SGD

global convergence

non-convex optimization
Abstract: Local SGD, a cornerstone algorithm in federated learning, is widely used in training deep neural networks and shown to have strong empirical performance. A theoretical understanding of such performance on nonconvex loss landscapes is currently lacking. Analysis of the global convergence of SGD is challenging, as the noise depends on the model parameters. Indeed, many works narrow their focus to GD and rely on injecting noise to enable convergence to the local or global optimum. When expanding the focus to local SGD, existing analyses in the nonconvex case can only guarantee finding stationary points or assume the neural network is overparameterized so as to guarantee convergence to the global minimum through neural tangent kernel analysis. In this work, we provide the first global convergence analysis of the vanilla local SGD for two-layer neural networks \emph{without overparameterization} and \textit{without injecting noise}, when the input data is Gaussian. The main technical ingredients of our proof are \textit{a self-correction mechanism} and \textit{a new exact recursive characterization of the direction of global model parameters}. The self-correction mechanism guarantees the algorithm reaches a good region even if the initialization is in a bad region. A good (bad) region means updating the model by gradient descent will move closer to (away from) the optimal solution. The main difficulty in establishing a self-correction mechanism is to cope with the gradient dependency between two layers. To address this challenge, we divide the landscape of the objective into several regions to carefully control the interference of two layers during the correction process. As a result, we show that local SGD can correct the two layers and enter the good region in polynomial time. After that, we establish a new exact recursive characterization of the direction of global parameters, which is the key to showing convergence to the global minimum with linear speedup in the number of machines and reduced communication rounds. Experiments on synthetic data confirm theoretical results.

Global Convergence Analysis of Local SGD for Two-layer Neural Network without Overparameterization

Abstract

Local SGD, a cornerstone algorithm in federated learning, is widely used in training deep neural networks and shown to have strong empirical performance. A theoretical understanding of such performance on nonconvex loss landscapes is currently lacking. Analysis of the global convergence of SGD is challenging, as the noise depends on the model parameters. Indeed, many works narrow their focus to GD and rely on injecting noise to enable convergence to the local or global optimum. When expanding the focus to local SGD, existing analyses in the nonconvex case can only guarantee finding stationary points or assume the neural network is overparameterized so as to guarantee convergence to the global minimum through neural tangent kernel analysis. In this work, we provide the first global convergence analysis of the vanilla local SGD for two-layer neural networks without overparameterization and without injecting noise , when the input data is Gaussian. The main technical ingredients of our proof are a self-correction mechanism and a new exact recursive characterization of the direction of global model parameters . The self-correction mechanism guarantees the algorithm reaches a good region even if the initialization is in a bad region. A good (bad) region means updating the model by gradient descent will move closer to (away from) the optimal solution. The main difficulty in establishing a self-correction mechanism is to cope with the gradient dependency between two layers. To address this challenge, we divide the landscape of the objective into several regions to carefully control the interference of two layers during the correction process. As a result, we show that local SGD can correct the two layers and enter the good region in polynomial time. After that, we establish a new exact recursive characterization of the direction of global parameters, which is the key to showing convergence to the global minimum with linear speedup in the number of machines and reduced communication rounds. Experiments on synthetic data confirm theoretical results.

1 Introduction

Federated learning is a prevalent framework in distributed learning to significantly reduce the communication cost and effectively preserve the privacy of local clients [43, 27]. As the most popular algorithm in federated learning, local SGD has shown great empirical success in training deep neural networks (DNNs) [43, 39]. However, existing literature has not been able to fully explain or characterize the convergence of local SGD in training DNNs. Recently, extensive works are devoted to analyzing the convergence of local SGD and its variants in nonconvex optimization [39, 56, 20, 28,

32, 40]. However, traditional nonconvex analysis only guarantees convergence to a stationary point, and convergence to the global minimum is in general NP-hard [21].

Despite the NP-hardness results for nonconvex optimization, an increasing body of research tries to address structured nonconvex optimization by first-order methods with noise injection. For instance, Ge et al. [17] considered strict-saddle functions and showed that SGD with isotropic noise can find local minima in polynomial time. This motivated several ensuing works [25, 33, 26, 55, 1] on designing different first-order algorithms to improve convergence to local minima by injecting noise. Noise-injecting schemes and their variants (such as, for instance, broadening from isotropic to anisotropic noise [64]) were shown to help convergence to global minima for many problems that satisfy one of two conditions: (i) local minima are global minima, as in matrix completion [18], dictionary learning [48], and certain deep linear networks [29]; or (ii) neural networks with distributional assumption, such as two-layer neural networks with the Gaussian input [63]. It is worth noting that there is also a rich history on noisy GD based on Langevin dynamics (LD) [30, 44, 52, 8, 41, 9]. For instance, recent work [6] proposes Exponential Family Langevin Dynamics (EFLD) to relax the Gaussian noise assumption and include noisy sign-SGD and variants of drop-out as special cases. When assuming the neural network is overparameterized, neural tangent kernel (NTK) analysis [23] guarantees convergence to the global minimum for local SGD [23, 22, 58, 10]. However, the NTK theory is far from sufficient, since neural networks outperform their NTK counterpart in practice [4] and in theory [38].

Despite existing global convergence analyses of first-order methods for solving nonconvex optimization problems such as neural networks, they either require explicitly injecting noise or assume overparameterization such that NTK analysis can apply. Practical federated learning algorithms such as local SGD do not inject any noise and do not belong to the NTK regime, but they can still converge to global minima. For example, McMahan et al. [43] shows that the local SGD algorithm can achieve around 99% accuracy when training neural networks for an image classification task. This motivates us to study the following question in this paper:

Is it possible to formally prove that the local SGD algorithm can find global minima for two-layer neural networks without injecting noise and without overparameterization?

In this paper, we give a positive answer to this question under Gaussian input. We are inspired by a line of work on neural network learning theory with Gaussian input [7, 49, 15, 37, 61, 13] and, in particular address the distributed version of the setting in [13]. Suppose N local machines share the following network:

where Z = ( Z 1 , · · · , Z k ) ∈ R d × k is the input matrix, w ∈ R d is the weight, a ∈ R k is the output weight, and σ ( x ) = max { x, 0 } denotes the ReLU activation function. A good property of this network is positive homogeneity, i.e., f ( Z , c w , a /c ) = f ( Z , w , a ) holds for any c > 0 . We assume each entry of the input Z is independently sampled from a standard Gaussian distribution. Then the response is generated by a noiseless teacher network: y = f ( Z , w ∗ , a ∗ ) . Without loss of generality, we further assume ∥ w ∗ ∥ = 1 . We hope to learn a student network by collaboratively minimizing the following mean square loss among N local machines:

Obviously, ( w ∗ , a ∗ ) is the global minimum of the objective (1) with zero loss. In particular, the loss function also has a spurious local minimum; hence, minimizing loss is a nonconvex optimization problem. Please refer to work in [13] for further details on the landscape of L ( w , a ) .

Despite the special input distribution, to the best of our knowledge, there is currently no work demonstrating the global convergence of vanilla SGD or vanilla local SGD without overparameterization. Work in [13] proved that randomly initialized GD can converge to the global minimum or the local minimum with a constant probability. The initial region where GD can converge to the global minimum is also called the attraction basin, where the gradients of two layers both point in the correct directions to the ground truth. To obtain the global convergence with arbitrary initialization in the same initial region, Zhou et al. [63] proposed a new perturbed GD algorithm by carefully injecting noise to the weights in two layers. Although the convergence of vanilla SGD has not been explored

theoretically, the simulation results in [13, 63] show that vanilla SGD with random initialization can converge to the global minimum with probability 1 when the ratio | 1 ⊤ a ∗ | / ∥ a ∗ ∥ is large. This result motivates us to investigate the global convergence of local SGD without injecting additional noise to escape the local minimum.

In this paper, we analyze the global convergence of the vanilla local SGD for training a two-layer neural network with Gaussian input, whose initialization starts from the same initial region as in [13, 63]. Formally, our main contributions are summarized as follows:

We introduce a new self-correction mechanism of local SGD under a condition on a ∗ (see Assumption 1): the signals from two layers can be corrected in polynomial time even though the initial point comes from a bad region where the gradients point to the wrong direction of ( w ∗ , a ∗ ) . The condition also explains the simulation results in [13]. The self-correction process is very difficult to analyze due to the mutual influence effect of two layers. To address this challenge, we utilize a novel technique by carefully dividing the landscape of the objective (1) into several regions. In each region, the negative effect from one layer to another layer can be controlled to a negligible scale. We notice that Li and Yuan [37] also showed the self-correction phase of SGD for the two-layer network under Gaussian input. However, the network's structure in [37] is different from the network studied in this paper. In addition, Li and Yuan [37] also required bounding the noise of stochastic gradient and so cannot handle the vanilla SGD with Gaussian noise as in our case.

We show the global convergence of local SGD with linear speedup, which indicates that the iteration complexity is divided by the number of machines N . In addition, we also show that the communication complexity of local SGD is reduced compared with the naive parallel version of SGD which needs to communicate at every iteration. The analysis in the convergence stage is very different from the GD in Du et al. [13]. We establish a new recursive dynamic to characterize the direction of the global weight in the first layer. Moreover, the objective is not smooth, since the gradients incorporate the angle between the first layer's weight and the ground truth. Therefore, conventional analysis of local SGD for a general smooth objective [47, 54] cannot be applied in the convergence stage. Due to the inner structure of gradients under Gaussian input, we find that the discrepancy caused by the local updates can shrink as the angle decreases, which enables us to refine the bound of discrepancy to be dominated by the statistical bound of noise.

We conduct several simulations on the two-layer neural network to verify the theoretical results. The experiments demonstrate that local SGD indeed corrects the wrong signals from the initial point and exhibits speedup in the convergence stage, corroborating our theoretical results. The simulation results also verify that the condition imposed on | 1 ⊤ a ∗ | / ∥ a ∗ ∥ is almost necessary to show the convergence with almost arbitrary initialization.

2 Related Work

Federated Optimization There is a wave of studies on federated optimization in different settings. In the convex optimization setting with homogeneous data, one-shot averaging was studied [65, 42, 60], where each machine solves a local optimization problem and the average happens only at the last iterate. Local SGD skips communication rounds, and the convergence analysis is shown for convex [47, 11, 34, 28, 54, 53, 32, 31] and nonconvex optimization problems [62, 24, 51, 39, 20, 56, 34, 28, 45, 59, 32]. There is a line of work which tried to compare minibatch SGD and local SGD in federated learning [54, 53]. However, these optimization algorithms only work for black-box functions and do not utilize the property of neural networks. As neural network loss landscapes are typically nonconvex, these federated optimization algorithm can only guarantee to find a stationary point instead of a global minimum.

Optimization Theory for Neural Networks There is a line of work studying two layer neural networks with Gaussian input [49, 12, 37, 61, 7, 19, 5]. Li and Liang [36] studied two layer neural networks with cross-entropy loss and showed that SGD can find the global minimum when the neural network is overparameterized. Du et al. [16] proved that GD can find the global minimum for two layer overparameterized neural networks under ℓ 2 loss. These results are later extended to deep neural networks by [14, 3, 66, 2] but are not directly applicable to analyzing local SGD in the distributed setting.

Algorithm 1 Local SGD for training two-layer neural network with Gaussian input

Federated Learning on Neural Networks There is a line of work which studied federated learning algorithms on overparameterized neural networks under the NTK regime [35, 22, 10, 58, 57]. In contrast, our analysis does not fall in the NTK regime: we directly study the dynamics of local SGD over neural networks without overparametrization.

3 Notations and Problem Setup

Denote ∥ · ∥ the Euclidean norm and ⟨· , ·⟩ by the inner product. S d -1 denotes the d -dimensional unit sphere and B k ( ρ ) denotes the k -dimensional ball with center zero and radius ρ . For two vectors v , w ∈ R d , denote ∠ ( v , w ) ∈ [0 , π ] the angle between v and w . The uniform distribution is denoted by Unif( · ) . Moreover, we use ˜ O to hide logarithmic factors.

We adopt the weight-normalization technique [46] to the first layer by re-parametrizing w = v / ∥ v ∥ , which leads to the following prediction model:

Given any sample ( Z , y ) , we denote the empirical loss by

In the distributed environment, suppose there are N local machines sharing the same teacher model f ( Z , w ∗ , a ∗ ) . It means that given any local input Z i , the response is y i = f ( Z i , w ∗ , a ∗ ) . Let I = { t 0 , ..., t R } be the set of synchronization time, where t 0 = 0 , t R = T and t r +1 -t r = I for any r . The detailed procedure of local SGD is presented in Algorithm 1, where the initial point is from the same region in Du et al. [13] and Zhou et al. [63]: v 0 ∈ S d -1 and a 0 ∈ B k ( | 1 ⊤ a ∗ | / √ k ) . In each round, the i -th machine runs I steps of SGD using the stochastic gradients ∇ v ℓ ( v i t , a i t ; Z i t , y i t ) and ∇ a ℓ ( v i t , a i t ; Z i t , y i t ) computed by the local input ( Z i t , y i t ) . At the end of a round, the server aggregates local weights to obtain the global weight and then synchronizes the global weight to each machine.

4 Theoretical Analysis

We first introduce two auxiliary sequences v t = ∑ N i =1 v i t /N and a t = ∑ N i =1 a i t /N , which often appears in the analysis of local SGD [47, 53, 28]. Notice that ( v t , a t ) = ( v i t , a i t ) for i ∈ [ N ] if t is the synchronization time, i.e., t ∈ I . Denote w i t = v i t / ∥ v i t ∥ the local normalized weight and P i t = ( I -w i t ( w i t ) ⊤ ) / ∥ v i t ∥ the respective projection matrix. Now let L ( w , a ; Z ) be the empirical version of the loss function L ( w , a ) defined in (1), that is L ( w , a ; Z ) = 1 2 ( f ( Z , w ∗ , a ∗ ) -f ( Z , w , a )) 2 . Since y = f ( Z , w ∗ , a ∗ ) , we will hide y and write ℓ ( v , a ; Z ) ≡ ℓ ( v , a ; Z , y ) hereafter. Notice that L ( w , a ; Z ) = ℓ ( v , a ; Z ) for w = v / ∥ v ∥ . By recalling the normalized model (2), we have

According to Algorithm 1, the updates of auxiliary sequences can be written as

For ease of technical presentation, we denote the averaged noise terms by ξ t = 1 N ∑ N i =1 P i t ξ i t and ϵ t = 1 N ∑ N i =1 ϵ i t , where ξ i t = ∇ w L ( w i t , a i t ; Z i t ) - ∇ w L ( w i t , a i t ) and ϵ i t = ∇ a L ( w i t , a i t ; Z i t ) -∇ a L ( w i t , a i t ) are the local noises in stochastic gradients. In addition, for the iterates v i t and v t , we write ϕ i t = ∠ ( v i t , w ∗ ) and ϕ t = ∠ ( v t , w ∗ ) , respectively.

4.1 Exact Dynamic of Each Layer

In this subsection, we will give the exact dynamic of each layer in the training process of local SGD, which is the starting point of our analysis. Denote P t = 1 ∥ v t ∥ ( I -w t w ⊤ t ) by the global projection matrix, where w t = v t / ∥ v t ∥ . The proofs of this subsection are deferred to Appendix A.1.

Lemma 1. Let ˇ v t +1 = v t -η ( P t 1 N ∑ N i =1 ∇ w L ( w i t , a i t ) + ξ t ) and ˇ ϕ t = ∠ (ˇ v t , w ∗ ) . The first layer in local SGD satisfies that

Lemma 1 is crucial to control the angle and show the linear speedup in further analysis. For each local machine, Lemma 5.5 in [13] provides the dynamic of sin 2 ϕ i t , which cannot characterize the dynamic of global quantity sin ϕ t due to the nonlinearity. Here we introduce a new intermediate variable ˇ v t +1 and find an equality (4) to show the recursive relation between ∥ v t +1 ∥ 2 sin 2 ϕ t +1 and ∥ v t ∥ 2 sin 2 ϕ t through the global quantities, such as M 1 ,t and M 2 ,t . It is worthwhile noticing that M 1 ,t is the averaged noise in local SGD, whose variance is divided by the number of clients N , namely linear speedup term . In fact, we can control the dynamic of sin ϕ t by upper bounding the discrepancy term H t . Let us assume the last three terms in (4) are negligible and ( a i t ) ⊤ a ∗ > 0 for any i ∈ [ N ] :

When ϕ t > π/ 2 , ∥ v t ∥ 2 sin 2 ϕ t will continuously increase since λ t cos ϕ t < 0 . It indicates that ϕ t can decrease to π/ 2 if ∥ v t ∥ is upper bounded by a constant, which also means the first layer can be corrected and avoided converging to the spurious local minima.

When ϕ t < π/ 2 , ∥ v t ∥ 2 sin 2 ϕ t will continuously decrease to zero. It indicates that ϕ t can converge to zero if ∥ v t ∥ is lower bounded by a constant. The initial region with ϕ 0 < π/ 2 and a ⊤ 0 a ∗ > 0 is also called the attraction basin in [13].

Through the remarks above, we can see that the positive signal of the second layer is crucial to both the self-correction of the first layer and global convergence. Next lemma presents an exact dynamic of the averaged weight in the second layer.

Lemma 2. Let A 0 = | 1 ⊤ a ∗ | 2 -( 1 ⊤ a ∗ )( 1 ⊤ a 0 ) and g ( ϕ ) = ( π -ϕ ) cos ϕ +sin ϕ . For local SGD algorithm started with the initial point ( v 0 , a 0 ) , the second layer satisfies that

where S ( ϵ 0: t -1 ) (defined in Appendix A) is the noise term involving ϵ 0 , · · · , ϵ t -1 and

Notice that the first term g ( ϕ i s ) -1 in B i s is negative whenever ϕ i s < π/ 2 , but the second term in B i s is always positive for ϕ i s ∈ (0 , π ] . In particular, π -g ( ϕ i l ) tends to be larger when ϕ i l gets closer to π (i.e., w i t drifts away from w ∗ ). This insight provides a possibility that local SGD can correct the signal of the second layer by itself, instead of injecting additional noise like [63].

Through carefully inspecting the ingredients of dynamics, we have the following roadmap to show the global convergence of local SGD: (1) For arbitrary initialization ( v 0 , a 0 ) except for a measure zero set, where the angle of the first layer between initialization and the global minimum is π , show that a ⊤ t a ∗ can turn to the positive signal in polynomial time; (2) Show that a ⊤ t a ∗ can be lower bounded by a positive constant value and ϕ t will decrease below π/ 2 . (3) After entering the attraction basin, show that ( w t , a t ) will converge to the ground truth with a linear speedup guarantee.

4.2 Self-correction of Signals in Two Layers

In this subsection, we will show the iterates of local SGD can enter the attraction basin such that a ⊤ t a ∗ > 0 and ϕ t < π/ 2 after ˜ O ( η -1 ) steps. Before that, we introduce an assumption on a ∗ .

Assumption 1. Define α = | 1 ⊤ a ∗ | 2 / ( k ∥ a ∗ ∥ 2 ) . We assume k ≥ 320( π -1) 2 and the ground truth satisfies | 1 ⊤ a ∗ | 2 ≤ k ( π + k -1) 720 π log(4+ k 2 ) and

The conditions on k and | 1 ⊤ a ∗ | 2 are imposed for technical reasons. We believe that a constant lower bound (the right-hand side of (6)) for α is necessary to show the global convergence with almost arbitrary initialization, which is also verified by our simulation results in Table 1. We compute convergence probabilities of local SGD and minibatch SGD under different values of α . The initial points are randomly selected by v 0 ∼ Unif( S d -1 ) and a 0 ∼ Unif( B k ( | 1 ⊤ a ∗ | / √ k )) in each trial. If the convergence probability reaches 1, it means that local SGD or minibatch SGD can converge to the global minima with arbitrary initialization except for a measure zero set. When α ≤ 1 / 8 , both minibatch SGD and local SGD cannot converge to the global minima with probability 1.

Theorem 1 (Self-correction of the second layer) . For any initial point ( v 0 , a 0 ) with ϕ 0 ∈ [0 , π ) , we denote τ a = inf { t ≥ 0 : a ⊤ t a ∗ ≥ γ a } the first time, where γ a = 16( π -1) | 1 ⊤ a ∗ | 2 k ( π + k -1) . Under Assumption 1, if

for a sufficiently large constant c , then τ a ≤ O ( η -1 log k ) with probability at least 1 -δ .

This theorem completes the first step of the roadmap, the self-correction of the signal in the second layer, whose proof is given in Appendix D.1. If | 1 ⊤ a ∗ | 2 / ∥ a ∗ ∥ 2 ≤ 1 / poly( p ) , Du et al. [13] proved GD can converge to the spurious local minima with the initial condition a ⊤ 0 a ∗ < 0 and ϕ 0 > π/ 2 . Therefore, our results do not contradict theirs because we assume | 1 ⊤ a ∗ | 2 / ∥ a ∗ ∥ 2 is large. The simulations of [13, 63] show that the success probability of converging to the global optimum of SGD increases as the ratio | 1 ⊤ a ∗ | / ∥ a ∗ ∥ increasing. Theorem 1 can potentially explain this phenomenon since the condition (6) will be satisfied eventually when | 1 ⊤ a ∗ | / ∥ a ∗ ∥ keeps increasing.

Proof sketch of Theorem 1. The proof is very technical since the angle ϕ i t will affect the sign of B i s in the dynamic (5) of a ⊤ t a ∗ , while controlling ϕ i t also requires bounding the scale and controlling

the sign of a ⊤ t a ∗ . To tackle this challenge: (1) We divide the initial region of ϕ 0 into four regions: (0 , ˜ ϕ o ) , [ ˜ ϕ o , π/ 2) , [ π/ 2 , ˜ ϕ u ) and [ ˜ ϕ u , π ) ; (2) For each of the first three initial regions, we prove that either a ⊤ t a ∗ ≥ γ a in ˜ O (1 /η ) steps or enter the next region in ˜ O (1 /η ) steps; (3) In the last region [ ˜ ϕ u , π ) , we use the condition (6) to show a ⊤ t a ∗ will turn positive in ˜ O (1 /η ) steps. The dividing technique enables us to control the individual angle ϕ i t in four 'nice' regions where a ⊤ t a ∗ can be corrected with negligible disturbance. Here ˜ ϕ o = arccos(1 / 5) is chosen for convenience to avoid dependence on the initial value sin ϕ 0 if ϕ 0 < π/ 2 , and ˜ ϕ u is carefully chosen for technical reasons.

Under the same conditions of Theorem 1. If the learning rate satisfies

for a large absolute constant c > 0 , we have sin ϕ τ a ≥ ϖ 15 √ k with probability at least 1 -δ .

Recalling the dynamic (4), we know sin ϕ t tends to decrease to 0 when ϕ t > π/ 2 and a ⊤ t a ∗ < 0 . It means that the first layer moves in the direction of spurious local minima during the self-correction process of the second layer. Lemma 3 provides a lower bound for sin ϕ τ a to guarantee the first layer will not converge to the local minima. The proof of Lemma 3 is deferred to Appendix D.2.

Theorem 2 (Self-correction of the first layer) . For the initial point ( v 0 , a 0 ) with a ⊤ 0 a ∗ ≥ γ a (defined in Theorem 1) and sin ϕ 0 > 0 , we denote τ v = inf { t ≥ 0 : ϕ v ≤ ˜ ϕ l } , where cos ˜ ϕ l = 144 π π + k -1 | 1 ⊤ a ∗ | 2 k log(4 + k 2 ) . If the learning rate satisfies

for a large absolute constant c > 0 , it holds that τ v ≤ O ( 1 η k 2 sin 3 ϕ 0 ) and a ⊤ τ v a ∗ ≥ γ a / 32 with probability at least 1 -δ .

This theorem finishes the second step of the roadmap to convergence. The detailed proof is deferred to Appendix D.3. Now we give the iteration complexity spent for self-correction for the worst initialization: a ⊤ 0 a ∗ < 0 and ϕ 0 > max { ˜ ϕ u , π -( ∥ a ∗ ∥ -2 / 20 ∧ 1) } . Since sin ˜ ϕ u = O (1) , plugging the lower bound of sin ϕ τ a in Lemma 3, we can get

Proof sketch of Theorem 2. We divide the initial region into three regions: ( ˜ ϕ f , π ) , [ ˜ ϕ f , ˜ ϕ l ) and [0 , ˜ ϕ l ] , where ˜ ϕ f is adapted to the choices for ˜ ϕ l and ˜ ϕ u . For the first two regions, we show that ϕ t will enter the next region in polynomial time, while the last region is our target. The difficult part to show the correction of ϕ t falls in how to guarantee a ⊤ t a ∗ will be lower bounded by O ( γ a ) for any t ≥ 0 , which is divided and conquered by considering the intermediate quantities such as ˜ ϕ l , ˜ ϕ f and ˜ ϕ u . To better understand the choices for these quantities and γ a , we rearrange B i s and have the following decomposition

Due to the condition (6) in Assumption 1, ζ a is positive and used to correct the sign of a ⊤ t a ∗ and keep it positive. Since ϱ a is still not negligible when s is small, we use γ a to compensate -ϱ a after enrolled summation. The uncertain part ϑ a will be positive whenever ˜ ϕ l ≤ ϕ i l ≤ ˜ ϕ u for l ≤ s .

4.3 Convergence with Linear Speedup

With the correction guarantees in the previous subsection, we are ready to proceed with the convergence analysis of local SGD.

Theorem 3. Suppose the initial point ( v 0 , a 0 ) satisfies a ⊤ 0 a ∗ ≥ γ a / 32 and ϕ 0 ≤ ˜ ϕ l . For any ϵ > 0 , we choose η = 1 ck 2 ∥ a ∗ ∥ 2 Nϵ d log( dN/ϵδ ) for some absolute constant c > 0 . If I ≲ d ϵN min { 1 , k 2 d 1 / 2 N 1 / 2 , k 4 d N 2 ϵ } and ϵ < min { d -1 , dk -2 } , then ℓ ( v T , a T ) ≲ ϵ holds with probability at least 1 -δ where T = ˜ O ( dk 4 Nϵ ) .

We have the following implications about the result in Theorem 3:

To the best of our knowledge, this is the first convergence result with linear speedup on the number of machines N for two-layer neural networks. Besides, our convergence analysis does not rely on the overparameterization of the width of the second layer (i.e., k ).

The dependency on ϵ matches the best-known results of local SGD for strongly convex objective in Woodworth et al. [53]. In fact, the size of the first layer d resembles the variance of stochastic gradients σ 2 in the traditional optimization literature. We can show ⟨∇ a L ( w t , a t ) , a t -a ∗ ⟩ ≥ π -1 2 π ∥ a t -a ∗ ∥ 2 -O ( ϕ t ) ∥ a ∗ ∥ 2 . Therefore, when the first layer converges has converged ( ϕ t ≈ 0 ), a t can converge to a ∗ like the strongly convex regime.

According to the condition on I , we can obtain the communication complexity in the convergence stage as

When N ≲ min { dk 4 , ϵ -1 } , the communication complexity in this stage can be R conv = ˜ O ( k 4 ) , which is significantly reduced compared with the iteration complexity.

We show local SGD's convergence layer by layer. Next lemma ensures that the weight of the first layer can converge to the ground truth in polynomial time.

Lemma 4 (Convergence of the first layer) . Under the settings in Theorem 3. Suppose the initial point satisfies ϕ 0 ≤ ˜ ϕ l and a ⊤ 0 a ∗ ≥ γ a / 32 . With probability at least 1 -δ , we can guarantee that sin 2 ϕ t ≲ ϵ holds for any ˜ O ( k 2 η -1 ) ≤ t ≤ ˜ O ( η -2 ) .

It is worthwhile noticing that ( v τ v , a τ v ) in Theorem 2 satisfies the conditions for initial point in Lemma 4. Therefore, local SGD enters the attraction basin after finishing the self-correction process. In fact, showing the complexity ˜ O ( ϵ -1 ) is not trivial based on the traditional analysis of local SGD [47, 54]. The issue comes from the inner-product noise term M 1 ,t and the discrepancy term H t in Lemma 1, whose enrolled summations can only be bounded by O ( √ ϵ/d ) and O { ( √ Id + I ) ϵ } respectively at the beginning of convergence stage. Thanks to the special structure of gradient, we find that the scales of M 1 ,t and H t can shrink as sin ϕ t decreases. In light of this, we can continuously refine the bound by the following contraction: for ( K +1) T v ≤ t ≤ ˜ O ( η -2 ) it holds that

By taking K = O (log(1 /ϵ )) , we can obtain the target convergence sin 2 ϕ t ≲ ϵ with high probability. More details can be found in Appendix E.1.

Lemma 5 (Convergence of the second layer) . Under the choice for η and conditions for ϵ in Theorem 3. Suppose sin 2 ϕ t ≤ ϵ holds for any 0 ≤ t ≤ ˜ O ( η -2 ) . With probability at least 1 -δ , we can guarantee that ∥ a t -a ∗ ∥ 2 ≲ ϵ holds for any ˜ O ( η -1 ) ≤ t ≤ ˜ O ( η -2 ) .

The lemma stated above guarantees the convergence of the second layer after that of the first layer, whose proof is deferred to Appendix E.2. Equipped with Lemmas 4 and 5, we can prove Theorem 3 by leveraging the closed form of the objective L ( w , a ) (see Lemma 6).

0.2

0.4

0.6

0.8

1.0

cos

t

Local SGD

Minibatch SGD

1.0

0.5

0.0

0.5

1.0

a T

t a

*

0

25

50

75

100

Rounds

0

1

2

3

loss

(a)

a

0

=

-

a

∗

, ϕ

0

> π/

2

0.2

0.4

0.6

0.8

1.0

1.0

0.5

0.0

0.5

1.0

0

20

40

60

80

100

Rounds

0

1

2

3

(b)

a

0

=

-

a

∗

, ϕ

0

< π/

2

0.25

0.00

0.25

0.50

0.75

1.00

0.0

0.2

0.4

0.6

0.8

1.0

0

20

40

60

80

100

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

(c)

a

0

=

0

, ϕ

0

> π/

2

The perturbed GD method in Zhou et al. [63] requires manually adjusting the scale of injected noises and the learning rate between the transition of two phases. In local SGD, we can use a universal learning rate in the self-correction stage and convergence stage. Considering an initial point that is closer to spurious local minima ( ϕ 0 > ˜ ϕ u ) defined in Lemma 3, conditions (7), (8) and (9) on the learning rate can be satisfied if we choose I = O { d/ ( ϵN ) } . Together with (10) and (11), if N ≲ min { dk 4 , ϵ -1 } , we can get the total communication complexity

Therefore, our theory shows that local SGD can correct the signals and converge to the global minima with almost arbitrary initialization except for the case of initializing at the local minima ( ϕ 0 = π ).

5 Experiments

We now report some simulation results on synthetic data. We also compare the performance of two algorithms: local SGD and minibatch SGD. At each round, minibatch SGD updates the model weights by using the stochastic gradients with batch size NI in total to update the model, where each local machine computes I gradients and communicates with other machines. Minibatch SGD and local SGD have the same computation and communication structure [54].

In our first simulation, we set ∥ w ∗ ∥ = ∥ a ∗ ∥ = 1 and | 1 ⊤ a ∗ | = √ k . The results starting from three bad initial regions (i.e., ϕ 0 > π/ 2 or a ⊤ 0 a ∗ ≤ 0 ) are reported in Figure 1. As we can see, the signals of two layers can be both corrected to a good region (i.e., both cos( ϕ t ) and a ⊤ t a ∗ go to 1 and the loss goes to 0 when t increases) even in the worst case when a 0 = -a ∗ and ϕ 0 > π/ 2 . An interesting phenomenon is that local SGD can correct the signals faster than minibatch SGD. The reason is that the statistical error of stochastic gradients is not dominating in the self-correction process, so the effect of large batch size in minibatch SGD (i.e., with batch size NI for 1 iteration) is not as useful as smaller batch size in local SGD with more iterations (i.e., batch size 1 for I iterations on each machine with N machines in total). These results corroborate our theoretical analysis.

1.0

0.8

0.6

0.4

0.2

cos

t

Local SGD

Minibatch SGD

0.4

0.3

0.2

0.1

a T

t a

*

0

250

500

750

1000

Rounds

0.15

0.20

0.25

0.30

0.35

loss

(a)

α

= 1

/

32

1.0

0.8

0.6

0.4

0.2

0.4

0.3

0.2

0.1

0

200

400

600

800

1000

Rounds

0.2

0.3

0.4

0.5

(b)

α

= 1

/

16

0.25

0.00

0.25

0.50

0.75

1.00

0.00

0.25

0.50

0.75

1.00

0

200

400

600

800

1000

Rounds

0.0

0.5

1.0

1.5

(c)

α

= 1

/

4

In the second simulation, we plot the trajectories of local SGD and minibatch SGD under different values of α in Figure 2. Here we set w ∗ = 1 d / √ d and a ∗ = ( 1 αk , 0 , . . . , 0) / √ αk . The bad initial point is fixed as v 0 = ( -1 , 0 , . . . , 0) and a 0 = ( - √ α, 0 , . . . , 0) , where ϕ 0 > π/ 2 and a ⊤ 0 a ∗ < 0 . When α = 1 / 32 , 1 / 16 , we can see that minibatch SGD and local SGD converge to the local minima. When α = 1 / 4 , they can converge to the global minima with the same initial point. In this case, as we can see from Figure 2(c), the signals of two layers can be both corrected to a good region.

In the third simulation, we calculate the probabilities that local SGD and minibatch SGD converge to the global minimum under different values of α . The averaged results are taken over 100 independent repeated simulations. In each trial, we generate initial points by v 0 ∼ Unif( S d -1 ) and a 0 ∼ Unif( B k ( | 1 ⊤ a ∗ | / √ k )) . The results are given in Table 1.

6 Conclusion

We theoretically investigate the convergence of local SGD, a cornerstone algorithm in federated learning with strong empirical performance. We demonstrate convergence to the global minimum for two-layer neural networks without overparameterization , without injecting noise , and when the input data is Gaussian. A new self-correction mechanism guarantees the algorithm reaches a good region even if the initialization is in a bad region. The landscape of the objective is divided into several regions to carefully control the interference of the two layers during the correction process. A new exact recursive characterization of the direction of global parameter provides the key to show convergence to the global minimum with linear speedup in the number of machines and reduced communication rounds. Experiments on simulated data corroborate the theoretical results. To the best of our knowledge, this work is the first to theoretically demonstrate the global convergence of the vanilla local SGD for neural networks without overparameterization.

Acknowledgments and Disclosure of Funding

Wewould like to thank the anonymous reviewers for their helpful comments. Mingrui Liu is supported by a grant from George Mason University. The work of Yajie Bao was done when he was virtually visiting Mingrui Liu's research group in the Department of Computer Science at George Mason University.

References

Zeyuan Allen-Zhu and Yuanzhi Li. Neon2: Finding local minima via first-order oracles. Advances in Neural Information Processing Systems , 31, 2018.

Zeyuan Allen-Zhu and Yuanzhi Li. Can sgd learn recurrent neural networks with provable generalization? Advances in Neural Information Processing Systems , 32, 2019.

Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over-parameterization. In International Conference on Machine Learning , pages 242-252. PMLR, 2019.

Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. Advances in neural information processing systems , 32, 2019.

Ainesh Bakshi, Rajesh Jayaram, and David P Woodruff. Learning two layer rectified neural networks in polynomial time. In Conference on Learning Theory , pages 195-268. PMLR, 2019.

Arindam Banerjee, Tiancong Chen, Xinyan Li, and Yingxue Zhou. Stability based generalization bounds for exponential family langevin dynamics. International Conference on Machine Learning (ICML) , 2022.

Alon Brutzkus and Amir Globerson. Globally optimal gradient descent for a convnet with gaussian inputs. In International conference on machine learning , pages 605-614. PMLR, 2017.

Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina. Entropy-SGD: Biasing gradient descent into wide valleys. In International Conference on Learning Representations , 2017. URL https://openreview.net/forum?id=B1YfAfcgl .

Rishav Chourasia, Jiayuan Ye, and Reza Shokri. Differential privacy dynamics of langevin diffusion and noisy gradient descent. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , 2021. URL https: //openreview.net/forum?id=Nfbe1usrgx4 .

Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Local sgd optimizes overparameterized neural networks in polynomial time. In International Conference on Artificial Intelligence and Statistics , pages 6840-6861. PMLR, 2022.

Aymeric Dieuleveut and Kumar Kshitij Patel. Communication trade-offs for local-sgd with large step size. Advances in Neural Information Processing Systems , 32:13601-13612, 2019.

Simon Du and Jason Lee. On the power of over-parametrization in neural networks with quadratic activation. In International conference on machine learning , pages 1329-1338. PMLR, 2018.

Simon Du, Jason Lee, Yuandong Tian, Aarti Singh, and Barnabas Poczos. Gradient descent learns one-hidden-layer cnn: Don't be afraid of spurious local minima. In International Conference on Machine Learning , pages 1339-1348. PMLR, 2018.

Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International conference on machine learning , pages 1675-1685. PMLR, 2019.

Simon S Du, Jason D Lee, and Yuandong Tian. When is a convolutional filter easy to learn? arXiv preprint arXiv:1709.06129 , 2017.

Simon S Du, Jason D Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. arXiv preprint arXiv:1811.03804 , 2018.

Rong Ge, Furong Huang, Chi Jin, and Yang Yuan. Escaping from saddle points-online stochastic gradient for tensor decomposition. In Conference on learning theory , pages 797-842. PMLR, 2015.

Rong Ge, Jason D Lee, and Tengyu Ma. Matrix completion has no spurious local minimum. Advances in neural information processing systems , 29, 2016.

Rong Ge, Jason D Lee, and Tengyu Ma. Learning one-hidden-layer neural networks with landscape design. arXiv preprint arXiv:1711.00501 , 2017.

Farzin Haddadpour, Mohammad Mahdi Kamani, Mehrdad Mahdavi, and Viveck Cadambe. Local sgd with periodic averaging: Tighter analysis and adaptive synchronization. In Advances in Neural Information Processing Systems , pages 11080-11092, 2019.

Christopher J Hillar and Lek-Heng Lim. Most tensor problems are np-hard. Journal of the ACM (JACM) , 60(6):1-39, 2013.

Baihe Huang, Xiaoxiao Li, Zhao Song, and Xin Yang. Fl-ntk: A neural tangent kernel-based framework for federated learning analysis. In International Conference on Machine Learning , pages 4423-4434. PMLR, 2021.

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems , 31, 2018.

Peng Jiang and Gagan Agrawal. A linear speedup analysis of distributed deep learning with sparse and quantized communication. In Advances in Neural Information Processing Systems , pages 2525-2536, 2018.

Chi Jin, Rong Ge, Praneeth Netrapalli, Sham M Kakade, and Michael I Jordan. How to escape saddle points efficiently. In International conference on machine learning , pages 1724-1732. PMLR, 2017.

Chi Jin, Praneeth Netrapalli, and Michael I Jordan. Accelerated gradient descent escapes saddle points faster than gradient descent. In Conference On Learning Theory , pages 1042-1085. PMLR, 2018.

Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning. Foundations and Trends® in Machine Learning , 14(1-2):1-210, 2021.

Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International Conference on Machine Learning , pages 5132-5143. PMLR, 2020.

Kenji Kawaguchi. Deep learning without poor local minima. Advances in neural information processing systems , 29, 2016.

A. D. Kennedy. The Theory of Hybrid Stochastic Algorithms , pages 209-223. Springer US, Boston, MA, 1990. ISBN 978-1-4615-3784-7. doi: 10.1007/978-1-4615-3784-7_14. URL https://doi.org/10.1007/978-1-4615-3784-7_14 .

Ahmed Khaled, Konstantin Mishchenko, and Peter Richtárik. Tighter theory for local sgd on identical and heterogeneous data. In International Conference on Artificial Intelligence and Statistics , pages 4519-4529. PMLR, 2020.

Anastasia Koloskova, Nicolas Loizou, Sadra Boreiri, Martin Jaggi, and Sebastian Stich. A unified theory of decentralized sgd with changing topology and local updates. In International Conference on Machine Learning , pages 5381-5393. PMLR, 2020.

Kfir Y Levy. The power of normalization: Faster evasion of saddle points. arXiv preprint arXiv:1611.04831 , 2016.

Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Challenges, methods, and future directions. IEEE Signal Processing Magazine , 37(3):50-60, 2020.

Xiaoxiao Li, Meirui Jiang, Xiaofei Zhang, Michael Kamp, and Qi Dou. Fedbn: Federated learning on non-iid features via local batch normalization. arXiv preprint arXiv:2102.07623 , 2021.

Yuanzhi Li and Yingyu Liang. Learning overparameterized neural networks via stochastic gradient descent on structured data. Advances in neural information processing systems , 31, 2018.

Yuanzhi Li and Yang Yuan. Convergence analysis of two-layer neural networks with relu activation. In Advances in Neural Information Processing Systems , pages 597-607, 2017.

Yuanzhi Li, Tengyu Ma, and Hongyang R Zhang. Learning over-parametrized two-layer neural networks beyond ntk. In Conference on learning theory , pages 2613-2682. PMLR, 2020.

Tao Lin, Sebastian U Stich, Kumar Kshitij Patel, and Martin Jaggi. Don't use large mini-batches, use local sgd. arXiv preprint arXiv:1808.07217 , 2018.

Mingrui Liu, Zhenxun Zhuang, Yunwen Lei, and Chunyang Liao. A communication-efficient distributed gradient clipping algorithm for training deep neural networks. arXiv preprint arXiv:2205.05040 , 2022.

Yi-An Ma, Yuansi Chen, Chi Jin, Nicolas Flammarion, and Michael I. Jordan. Sampling can be faster than optimization. CoRR , abs/1811.08413, 2018. URL http://arxiv.org/abs/1811. 08413 .

Ryan McDonald, Keith Hall, and Gideon Mann. Distributed training strategies for the structured perceptron. In Human Language Technologies: The 2010 Annual Conference of the North American Chapter of the Association for Computational Linguistics , pages 456-464. Association for Computational Linguistics, 2010.

Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-efficient learning of deep networks from decentralized data. In Artificial intelligence and statistics , pages 1273-1282. PMLR, 2017.

R. M. Neill. MCMC Using Hamiltonian Dynamics , pages 113-162. Chapman and Hall/CRC, Boston, MA, 2011. ISBN 9780429138508. URL https://www.taylorfrancis.com/chapters/edit/10.1201/b10905-10/ mcmc-using-hamiltonian-dynamics-radford-neal .

Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konecny, Sanjiv Kumar, and H Brendan McMahan. Adaptive federated optimization. ICLR , 2021.

Tim Salimans and Durk P Kingma. Weight normalization: A simple reparameterization to accelerate training of deep neural networks. Advances in neural information processing systems , 29, 2016.

Sebastian U Stich. Local sgd converges fast and communicates little. arXiv preprint arXiv:1805.09767 , 2018.

Ju Sun, Qing Qu, and John Wright. Complete dictionary recovery over the sphere i: Overview and the geometric picture. IEEE Transactions on Information Theory , 63(2):853-884, 2016.

Yuandong Tian. An analytical formula of population gradient for two-layered relu network and its applications in convergence and critical point analysis. In International conference on machine learning , pages 3404-3413. PMLR, 2017.

Roman Vershynin. High-dimensional probability: An introduction with applications in data science , volume 47. Cambridge university press, 2018.

Jianyu Wang and Gauri Joshi. Cooperative sgd: A unified framework for the design and analysis of communication-efficient sgd algorithms. arXiv preprint arXiv:1808.07576 , 2018.

Max Welling and Yee Whye Teh. Bayesian learning via stochastic gradient langevin dynamics. In 28th International Conference on International Conference on Machine Learning , ICML'11, page 681-688, Madison, WI, USA, 2011. Omnipress. ISBN 9781450306195.

Blake Woodworth, Kumar Kshitij Patel, and Nathan Srebro. Minibatch vs local sgd for heterogeneous distributed learning. arXiv preprint arXiv:2006.04735 , 2020.

Blake Woodworth, Kumar Kshitij Patel, Sebastian Stich, Zhen Dai, Brian Bullins, Brendan Mcmahan, Ohad Shamir, and Nathan Srebro. Is local sgd better than minibatch sgd? In International Conference on Machine Learning , pages 10334-10343. PMLR, 2020.

Yi Xu, Rong Jin, and Tianbao Yang. First-order stochastic algorithms for escaping from saddle points in almost linear time. Advances in neural information processing systems , 31, 2018.

Hao Yu, Rong Jin, and Sen Yang. On the linear speedup analysis of communication efficient momentum sgd for distributed non-convex optimization. arXiv preprint arXiv:1905.03817 , 2019.

Yaodong Yu, Alexander Wei, Sai Praneeth Karimireddy, Yi Ma, and Michael Jordan. Tct: Convexifying federated learning using bootstrapped neural tangent kernels. Advances in Neural Information Processing Systems , 35:30882-30897, 2022.

Kai Yue, Richeng Jin, Ryan Pilgrim, Chau-Wai Wong, Dror Baron, and Huaiyu Dai. Neural tangent kernel empowered federated learning. In International Conference on Machine Learning , pages 25783-25803. PMLR, 2022.

Xinwei Zhang, Mingyi Hong, Sairaj Dhople, Wotao Yin, and Yang Liu. Fedpd: A federated learning framework with optimal rates and adaptivity to non-iid data. arXiv preprint arXiv:2005.11418 , 2020.

Yuchen Zhang, John C Duchi, and Martin J Wainwright. Communication-efficient algorithms for statistical optimization. The Journal of Machine Learning Research , 14(1):3321-3363, 2013.

Kai Zhong, Zhao Song, Prateek Jain, Peter L Bartlett, and Inderjit S Dhillon. Recovery guarantees for one-hidden-layer neural networks. In International conference on machine learning , pages 4140-4149. PMLR, 2017.

Fan Zhou and Guojing Cong. On the convergence properties of a k -step averaging stochastic gradient descent algorithm for nonconvex optimization. arXiv preprint arXiv:1708.01012 , 2017.

Mo Zhou, Tianyi Liu, Yan Li, Dachao Lin, Enlu Zhou, and Tuo Zhao. Toward understanding the importance of noise in training neural networks. In International Conference on Machine Learning , pages 7594-7602. PMLR, 2019.

Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects, 2019.

Martin Zinkevich, Markus Weimer, Lihong Li, and Alex J Smola. Parallelized stochastic gradient descent. In Advances in neural information processing systems , pages 2595-2603, 2010.

Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Stochastic gradient descent optimizes over-parameterized deep relu networks. arXiv preprint arXiv:1811.08888 , 2018.

Contents

Appendix

A Preliminaries for Analysis

We write 0 = (0 , . . . , 0) ⊤ , 1 = (1 , . . . , 1) ⊤ and denote I by the identity matrix. The following lemma gives the closed forms of the objective and gradients.

Lemma 6 (Du et al. [13]) . If every entry of Z is i.i.d. sampled from a Gaussian distribution with mean 0 and variance 1, and ∥ w ∗ ∥ = 1 , then the population loss is

where g ( ϕ ) = ( π -ϕ ) cos ϕ +sin ϕ . And the the expected gradient of w and a are

Denote the noises of stochastic gradients in each local machine i ∈ [ N ] by

Then we write the averaged noises as

Let P t = 1 ∥ v t ∥ ( I -w t w ⊤ t ) . We introduce an auxiliary sequence ˇ v t , which updates by

We define the σ -filtration in the training process: F 0 is the trivial filtration and

We also write E t [ · ] = E [ · | F t ] .

A.1 Dynamic of the First Layer

Lemma 1 restated. The first layer in local SGD satisfies that

where λ t = 1 N ∑ N i =1 π -ϕ i t 2 π ( a i t ) ⊤ a ∗ ∥ v t ∥ 2 , H t = ∥ v t +1 ∥ 2 sin 2 ϕ t +1 -∥ ˇ v t +1 ∥ 2 sin 2 ˇ ϕ t +1 and

Then it holds that

Proof. Let ˜ v t = v t -η P t 1 N ∑ N i =1 ∇ w L ( w i t , a i t ) , then we have ˇ v t +1 = ˜ v t -η ξ t . If follows that

From the definition λ t = 1 N ∑ N i =1 π -ϕ i t 2 π ( a i t ) ⊤ a ∗ ∥ v t ∥ 2 and (A.2), we can write

and

where ( i ) holds since P t v t = 0 , ( ii ) is true due to I -w t w ⊤ t is an idempotent matrix, and ( iii ) comes from ∥ w ∗ ∥ = 1 . Combining (A.8) and (A.9), we can get

Plugging (A.10) into (A.7) gives

where E t [ M 1 ,t ] = 0 since P t , ˜ v t ∈ F t and E t [ ξ t ] = 0 , and M 2 ,t ≥ 0 because P t and I -w ∗ ( w ∗ ) ⊤ are both semi positive definite. The result follows from the definition of H t .

A.2 Dynamic of the Second Layer

Lemma 2 restated. Let A 0 = ( 1 ⊤ a ∗ ) 2 -( 1 ⊤ a ∗ )( 1 ⊤ a 0 ) . For local SGD algorithm started with the initial point ( v 0 , a 0 ) , the second layer satisfies that

where

and

Proof. From (A.3), we have

Let A t = ( 1 ⊤ a ∗ ) 2 -( 1 ⊤ a ∗ )( 1 ⊤ a t ) . By the definition of a t +1 , we know

For t ≥ 0 , we also have

Plugging (A.13) into (A.12), then enrolling (A.12) from timestamp 0 to t gives that

In fact, we also used the following summation result in the first equality

A.3 Concentration Inequalities for the Noises

A.3.1 Sub-exponential Martingale Difference

Definition A.1 (Sub-exponential and sub-gaussian random variable) . The random variable x ∈ R is called sub-exponential with norm ∥ x ∥ ψ 1 if

In addition, the random variable x ∈ R is called sub-gaussian with norm ∥ x ∥ ψ 2 if

Lemma 7 (Proposition 2.6.1 in [50]) . Let x 1 , ..., x n be independent mean-zero sub-gaussian random variables. Then ∑ N i =1 x i is also a sub-gaussian random variable with

Lemma 8 (Lemma 2.6.8 in [50]) . If x is a sub-gaussian random variable, then x -E [ x ] is subgaussian too and

Lemma 9 (Lemma 2.7.7 in [50]) . Let x and y be sub-gaussian random variables. Then xy is sub-exponential and

Lemma 10. Suppose { x t } t ≥ 1 ⊆ R are sub-exponential martingale difference sequence with norm {∥ x t ∥ ψ 1 } t ≥ 1 adapted to the filtration {F t } t ≥ 1 , that is E [ x t |F t ] = 0 and ∥ x t ∥ ψ 1 = inf { s ≥ 0 : E [ e | x t | /s |F t ] ≤ 2 } . Then we have

where c > 0 is an absolute constant.

Proof. Invoking Proposition 2.7.1 in [50], we know

Define

It follows from (A.14) that for any t ≥ 1 ,

It further implies that

Using Markov's inequality, for 0 < λ ≤ c max 1 ≤ s ≤ t ∥ x t ∥ ψ 1 , we have

Now taking

and plugging it into (A.15) gives

Similarly, we can also prove

Then the proof is finished.

A.3.2 Scales of Sub-exponential Norm for Noises in Stochastic Gradients

For any ( v , a ) with w = v / ∥ v ∥ , we denote

Lemma 11. Let P = 1 ∥ v ∥ ( I -ww ⊤ ) . The following noise terms satisfy: for ℓ = 1 , 2

For any fixed vector u ∈ R k , u ⊤ ϵ ( ℓ ) is sub-exponential with max ℓ =1 , 2 ∥ u ⊤ ϵ ( ℓ ) ∥ ψ 1 ≤ c ∥ u ∥ max {∥ a ∥ , ∥ a ∗ ∥} .

For any fixed vector u ∈ R d , u ⊤ P ξ ( ℓ ) is sub-exponential with max ℓ =1 , 2 ∥ u ⊤ P ξ ( ℓ ) ∥ ψ 1 ≤ c ∥ a ∥∥ a ∗ ∥ / ∥ v ∥ .

Proof. Recall the definition of ϵ (1) , then we have

Applying Lemma 7 yields

where ( i ) and ( ii ) comes from Lemma 7 and 8 respectively, ( iii ) holds since σ ( Z ⊤ j w ) is a halfnormal random variable with variance ∥ w ∥ 2 / 2 , and ( iv ) holds due to ∥ w ∥ = 1 . Combining (A.16) and Lemma 8, we can verify that

Applying Lemma 9 gives

Invoking Lemma 8, we prove conclusion (1) immediately.

Recall the definitions P = 1 ∥ v ∥ ( I -ww ⊤ ) . Notice that, for any fixed vector u ∈ R d ,

It follows that

where we used ∥ w ∥ = 1 and ∥ P ∥ ≤ 1 / ∥ v ∥ . The counterpart for u ⊤ P ξ (2) follows from similar arguments.

A.3.3 Corollaries of Lemma 11.

Let { α l } l ≥ 0 be a sequence of real numbers. Recall that ξ t = 1 N ∑ N i =1 P i t ξ i t and ϵ t = ∑ N i =1 ϵ i t /N , where { ξ i t } i ∈ [ N ] and { ϵ i t } i ∈ [ N ] are independent random variables given F t . Applying Lemma 10 and Lemma 11(1), we can directly get the following results

and

Specially, choosing u = e j for j ∈ [ k ] , we also have

hold with probability at least 1 -δ .

Similarly, using Lemma 11(2), we also have

and

hold with probability at least 1 -δ .

B Bound the Weight's Norm during Training Process

B.1 The Scale of the Second Layer's Weight

Lemma 12 (The scale of the second layer's weight) . Suppose the learning rate satisfies that

for a large absolute constant c . Then for any 0 ≤ t ≤ ( η ∥ a ∗ ∥ 2 ) -2 , we are guaranteed that

with probability at least 1 -δ .

Proof. Recall the definition a t = 1 N ∑ N i =1 a i t , ϵ t = 1 N ∑ N i =1 ϵ i t and

Then notice that

where the last inequality holds since ∥ 11 ⊤ + ( π -1) I ∥ ≤ π + k -1 and g ( ϕ i t ) ≤ π . Denote A t = ( 1 ⊤ a t ) 2 -( 1 ⊤ a t )( 1 ⊤ a ∗ ) . Together with the definition of a t +1 , we have

Denote the event

Let α l = ( 1 -η π -1 2 π ) l , then it holds that ∑ s -1 l =0 α 2 l ≤ ∑ s -1 l =0 α l ≤ 2 π η ( π -1) . Using concentration (A.17), we have

holds with probability at least 1 -δ . Applying the concentration inequality (A.19) gives

holds with probability at least 1 -δ . In addition, from (A.13), we know

holds with probability at least 1 -δ . Plugging (B.3), (B.4) and (B.5) into (B.2), together with t ≤ η -2 , we can get

where the last inequality holds due to our assumption (B.1) and ∥ a 0 ∥ ≤ ∥ a ∗ ∥ . By setting T = t +1 and K a = 5 ∥ a ∗ ∥ in conclusion (1) of Lemma 15, we have

where we used the assumption (B.1). In conjunction with (B.6), we can verify that

P ( E t ∩ {∥ a t +1 ∥ > 5 ∥ a ∗ ∥} ) ≤ P ( One of (B.3), (B.4), (B.5), and (B.7) does not hold ) ≤ 4 δ.

Therefore, we have

Hence by induction, we can verify that P ( E t ) ≥ 1 -4 tδ for any t ≤ ( ηK 2 a ) -2 . By adjusting the level of δ and P ( E 0 ) = 1 due to ∥ a 0 ∥ ≤ ∥ a ∗ ∥ , we can finish the proof.

B.2 The Scale of the First Layer's weight

Lemma 13 (The scale of the first layer's weight) . Under the conditions of Lemma 12. Let K a = 5 ∥ a ∗ ∥ , we assume the learning rate satisfies

for a large absolute constant c . Then for any 0 ≤ t ≤ (800 cη 2 ( √ Id + I ) √ log( Nd/δ ) K 4 a ) -1 , we are guaranteed that 1 / 2 ≤ ∥ v t ∥ ≤ 3 with probability at least 1 -δ .

Proof. In this proof, we suppose the event Lemma 12 holds and hide the indicator function everywhere. To facilitate the technical presentation, denote ˜ v i t = v i t -η P i t ∇ w L ( w i t , a i t ) . Define the following event

Step 1: show ∥ v t ∥ ≥ 1 / 2 under the event U t -1 . From v t = 1 N ∑ N i =1 v i t and P t -1 v t -1 = 0 , we have

where the inequality follows from dropping positive terms and the inductive assumption ∥ v t -1 ∥ ≥ 1 / 2 . With probability at least 1 -δ , it holds that

where ( i ) holds due to ∥ w i t -1 ∥ = ∥ w t -1 ∥ = 1 ; ( ii ) and ( iii ) hold due to the inductive assumption ∥ v t -1 ∥ ≥ 1 / 2 ; ( iv ) follows from the conclusion (1) in Lemma 14; and ( v ) is true because 6 cη ( √ Id + I ) √ log( Nd/δ ) ≤ 1 / 4 . Substituting (B.9) into (B.8), with probability at least 1 -2 δ , we have

where ( i ) holds because ∥∇ w L ( w i t , v i t ) ∥ ≤ | ( a i t ) ⊤ a ∗ | ≤ 5 K 2 a by using Lemma 12 and ∥ P t -1 ∥ ≤ 1 ∥ v t -1 ∥ ≤ 2 holds under the event U t -1 ; ( ii ) holds due to ∥ v 0 ∥ = 1 ; and ( iii ) follows from t ≤ (800 cη 2 ( d ∨ I ) K 4 a √ log( Nd/δ )) -1 such that

In fact, the bound above holds because 800 cη 2 ( √ Id + I ) √ log( Nd/δ ) K 4 a ≤ 0 . 01 and (1 -x ) 1 x is decreasing in (0 , 1) . According to the definition of ˜ v i s and Lemma 12, we know

Recall that ξ s = 1 N ∑ N i =1 P i s ξ i s . Invoking Lemmas 11 and 10, we have

holds with probability at least 1 -δ . Plugging (B.12) into (B.10), we get

where the first inequality holds since √ Id log( Nd/δ ) ≥ c . Step 2: show ∥ v t ∥ ≤ 3 under the event U t -1 . From the update rule of v i t and v t = 1 N ∑ N i =1 v i t , with probability at least 1 -δ , we have

where ( i ) comes from P t -1 v t -1 = 0 and ∥ v t -1 ∥ ≤ 3 under U t -1 ; ( ii ) comes from ∥ P t -1 ∥ ≤ 1 / ∥ v t -1 ∥ ≤ 2 under U t -1 ; and ( iii ) holds due to (B.9). With probability at least 1 -δ , we have the following concentration

where we used the assumption t ≤ (400 cη 2 ( d ∨ I ) √ log( Nd/δ ) K 4 a ) -1 . Similar to (B.12), we can also guarantee

Plugging (B.15) and (B.16) into (B.14) gives that

holds with probability at least 1 -3 δ . Together with (B.13), we have showed that

P ( U t -1 ∩ { 1 / 2 ≤ ∥ v t ∥ ≤ 3 } ) ≥ P ( U t -1 ) -P ( U t -1 ∩ {∥ v t ∥ ≤ 1 / 2 } ) -P ( U t -1 ∩ {∥ v t ∥ > 3 } ) ≥ P ( U t -1 ) -P ( One of (B.9), (B.10), (B.13), and Lemma 12 does not hold ) -P ( One of (B.14), (B.15), (B.16), and Lemma 12 does not hold ) ≥ P ( U t -1 ) -8 δ.

Using induction and P ( U 0 ) = 1 due to ∥ v 0 ∥ = 1 , we can prove the desired conclusion by adjusting the level of δ .

Remark. In the following proofs, we use c to denote the absolute constant, which does not depend on any variables of the model or training process. For conciseness, we do not distinguish the specific value of c in some contexts.

C Bound the Discrepancy during Training Process

C.1 The Discrepancy in the First Layer

Lemma 14 (The discrepancy in the first layer) . Denote O i,j r,T = max t r ≤ t ≤ T ∥ v i t -v j t ∥ for t r ≤ T ≤ t r +1 -1 . Suppose the learning rate satisfies the conditions in Lemmas 12 and 13.

If max i ∈ [ N ] ,t r ≤ t ≤ T -1 ∥ a i t ∥ ≤ K a holds, with probability at least 1 -δ , we have

In addition, we also have

Proof. Denote the discrepancy term by

Since ∥ w ∗ ∥ = 1 and ∥ a i t ∥ ≤ 5 K a in Lemma 12, we can bound h i,j t as

From the definition of B i s and B j s in Lemma 2, we know for any i, j ∈ [ N ] ,

where the last inequality holds due to g ( ϕ ) ∈ [0 , π ] for ϕ ∈ [0 , π ] . Using concentration (A.18) and t ≤ kη -1 , with probability at least 1 -δ , we can guarantee

Using the dynamic of a ⊤ t a ∗ in Lemma 2, with probability at least 1 -δ , we also have

where the second last inequality holds due to the concentration (C.3). Plugging (C.4) into (C.2) gives

where we used the following concentration

Taking maximum on the left hand side of (C.5) over t r ≤ t ≤ T , we can get

where we also used the assumption η ( I + √ kI ) < 1 .

C.2 The Discrepancy in the Second Layer

Lemma 15 (The discrepancy in the second layer) . Let Ξ i,j r,T = max t r ≤ t ≤ T ∥ a i t -a j t ∥ for t r ≤ T ≤ t r +1 -1 . Suppose the learning rate satisfies the conditions in Lemmas 12 and 13. If max i ∈ [ N ] ,t r ≤ t ≤ T -1 ∥ a i t ∥ ≤ K a holds for K a > ∥ a ∗ ∥ and cηkI < 1 , with probability at least 1 -δ , we have

Proof. Recall the closed form of ∇ a L ( w , a ) , then for i, j ∈ [ N ] , we have

Let t r be the previous synchronization time before t . Using the fact a i t r = a j t r = a t r , we have

where the last inequality holds due to Lemma 11(1). Taking maximum on the both sides over t r ≤ t ≤ T , we get

D Self-correction of Signals

Recalling the dynamic in Lemma 2,

where

Using concentration (A.18), with probability at least 1 -δ , we can guarantee

For any 0 ≤ s ≤ ˜ O ( η -2 ) and any i ∈ [ N ] , with probability at least 1 -δ , it holds that

where

According to Lemmas 14, 12 and relations (D.3), (D.4), (A.20), we know P ( E c ℓ ) ≤ δ for 0 ≤ ℓ ≤ 4 .

D.1 Self-correction of the Second Layer

Theorem 1 restated. For any initial point ( v 0 , a 0 ) with ϕ 0 ∈ [0 , π ) , we denote τ a = inf { t ≥ 0 : a ⊤ t a ∗ ≥ γ a } the first time, where γ a = 16( π -1) | 1 ⊤ a ∗ | 2 k ( π + k -1) . Under Assumption 1, if

for a sufficiently large constant c , then τ a ≤ ˜ O ( η -2 ) with probability at least 1 -δ .

Proof. We first define a sequence of events as

Notice that once we show for some τ > 0 ,

then we can conclude that P ( A τ ) ≤ δ . According to the definition of τ a , we can guarantee τ a ≤ τ with probability at least 1 -δ .

For simplicity, we introduce the following stopping time w.r.t. the angle ϕ during the training process,

Define ˜ ϕ l and ˜ ϕ u such that

Choose ˜ ϕ o = arccos(1 / 5) such that cos ˜ ϕ l ≤ 1 / 5 = cos ˜ ϕ o due to Assumption 1 such that

where ( i ) follows from ∥ w ∗ ∥ = 1 and the triangle inequality; and ( ii ) holds due to Lemma 14. Define a good event

The correction condition (6) is equivalent to

where τ 1 , 1 , τ 1 , 2 and τ 1 , 3 are deterministic and will be defined later.

Initial region: (0 , ˜ ϕ o ) . For any s ≤ min { τ -( ˜ ϕ o ) , ˜ O ( η -2 ) } , under the event E good , it holds that

Since g ( · ) ∈ [0 , π ] is decreasing in [0 , π ] , plugging it into (D.2), we can get the following lower bound

Under the event E good , using the lower bound of A 0 in (B.5), for t ≤ min { τ -( ˜ ϕ o ) , ˜ O ( η -2 ) } , we can guarantee

Here we take

where we used the fact g (0 . 46 π ) -1 ≥ 0 . 2 . Since τ 1 , 1 -1 ≤ ˜ O ( η -2 ) , under the event { τ 1 , 1 -1 < τ -( ˜ ϕ o ) } ∩ E good , it holds that

We denote

where ( i ) holds due to ∥ a ∗ ∥ ≥ | 1 ⊤ a ∗ | / √ k ; ( ii ) follows from k ≥ 320( π -1) 2 ; and ( iii ) holds due to the assumption (D.6). We can conclude that

where the second last inequality holds since

such that

If τ 1 , 1 ≥ τ -( ˜ ϕ o ) , we regard ( v τ -( ˜ ϕ o ) , a τ -( ˜ ϕ o ) ) as the new initial point and step into the analysis of the next region.

Initial region: [ ˜ ϕ o , π/ 2) . Denote t 0 = τ -( ˜ ϕ o ) . From now on, we assume the event Z c 1 happens and hide the indicator. Recalling the definition of h t , we have

where ( i ) holds due to (B.9) and ( ii ) is true because ∥∇ w L ( w i t , a i t ) ∥ ≤ | ( a i t ) ⊤ a ∗ | ≤ K 2 a . In addition, we also have

Here we take

For any t 0 ≤ s ≤ t -1 ≤ t 0 + τ 1 , 2 , under the event A t -1 ∩ E good , invoking (D.12) guarantees that,

where ( i ) holds due to ∥ v s ∥ ≥ 1 / 2 , ∥ v t 0 ∥ ≥ 1 / 2 and the definition of A t -1 ; ( ii ) comes from t 0 ≤ τ 1 , 1 and E 1 ; ( iii ) follows from ϕ t 0 ≥ ˜ ϕ l , and E 2 ; and ( iv ) holds due to the assumption (D.6). Further, under the event A t -1 ∩ E good , we have for any i ∈ [ N ] and t 0 ≤ s ≤ t -1 ≤ t 0 + τ 1 , 2 -1

Under the event A t -1 ∩ E good ∩ { τ -( π/ 2) -t 0 ≥ τ 1 , 2 } , it implies that for any t 0 ≤ s ≤ t -1 ≤ t 0 + τ 1 , 2 -1 ,

Together with (D.1) and (D.3), we can guarantee that

Choosing t = t 0 + τ 1 , 2 , it holds that

where we used the assumptions (D.6) and k ≥ 320( π -1) 2 .

Recalling that Z 2 = { τ -( π/ 2) ≥ τ -( ˜ ϕ o ) + τ 1 , 2 } , hence we have verified that

If τ 1 , 2 ≥ τ -( π/ 2) , we regard ( v τ -( π/ 2) , a τ -( π/ 2) ) as the new initial point and step into the analysis of the next region.

Initial region: [ π/ 2 , ˜ ϕ u ) . From now on, we assign t 0 = τ -( π/ 2) and consider the case { τ -( ˜ ϕ o ) ≤ τ 1 , 1 } ∩ { t 0 < τ -( ˜ ϕ o ) + τ 1 , 2 } . Take

where the inequality holds due to our assumptions k ≥ 320( π -1) 2 and (D.8). Similar to (D.14), for any t 0 ≤ s ≤ t -1 ≤ ˜ O ( η -2 ) , under the event A t -1 ∩ E good , we have

where ( i ) holds due to ϕ t 0 ≥ π/ 2 and the definition of A t -1 ; and ( ii ) holds due to s -1 -τ ≤ ˜ O ( η -2 ) . Under the event A t -1 ∩ E good , we can guarantee that

where we used the assumption (D.6). It implies that min i ϕ i s ≥ ˜ ϕ l under the event A t -1 ∩ E good . In addition, we notice that cos ˜ ϕ l = -cos ˜ ϕ u , which implies ˜ ϕ l + ˜ ϕ u = π and sin ˜ ϕ l = sin ˜ ϕ u . Under the event A t -1 ∩ E good ∩ { t -1 ≤ min { τ -( ˜ ϕ u ) , ˜ O ( η -2 ) }} , for any s ≤ t -1 , we have

where the last inequality holds due to the condition (D.7). Together with the definition of B i s , we have

We know ζ a > 0 due to the condition (D.8). Under A t -1 ∩E good ∩{ t -1 ≤ min { τ -( ˜ ϕ u ) , ˜ O ( η -2 ) }} , we can guarantee

where we used | a ⊤ t 0 a ∗ | ≤ ∥ a ∗ ∥ · K a = 5 ∥ a ∗ ∥ 2 . Now let t = t 0 + τ 1 , 3 , under A t -1 ∩E good ∩{ t 0 + τ 1 , 3 -1 ≤ τ -( ˜ ϕ u ) } , we have

where we used the assumption (D.8) and (D.6). Since Z 3 = t 0 + τ 1 , 3 -1 ≤ τ -( ˜ ϕ u ) , we have

If τ -( ˜ ϕ u ) < τ -( π/ 2) + τ 1 , 3 (that is Z c 3 ), we regard ( v τ -( ˜ ϕ u ) , a τ -( ˜ ϕ u ) ) as the new initial point and step into the analysis of the next region.

Initial region: [ ˜ ϕ u , π ) . Now we assign t 0 = τ -( ˜ ϕ u ) and choose τ 1 , 4 = τ 1 , 3 . Recall that cos ˜ ϕ u = -144 π π + k -1 ( 1 ⊤ a ∗ ) 2 k log(4 + k 2 ) . Similar to what we showed in (D.17), for any t 0 ≤ s ≤ t -1 ≤ t 0 + τ 1 , 4 -1 , under the event A t -1 ∩ E good , we have,

where ( i ) holds due to ∥ v m ∥ ≤ 3 , ϕ 0 > π/ 2 and the definition of A t -1 ; and ( ii ) holds due to ∥ v s ∥ ≤ 3 . Applying the bound | cos ϕ s -cos ϕ i s | in (D.4) and the assumption D.6, we have

which means min i ϕ i s ≥ π 2 . It follows that for any s ≤ t -1 ,

Plugging (D.22) and (D.3) into (D.1), under the event A t -1 ∩ E good , we can get

Plugging t = t 0 + τ 1 , 4 , according to (D.23) and (D.19), we can guarantee that ϕ τ a ≥ π/ 2 and

Conclusion. Combining (D.10), (D.15), (D.20) and (D.24), we have

Therefore, we have proved τ a ≤ ∑ 4 q =1 τ 1 ,q with high probability for ϕ 0 ∈ [0 , ˜ ϕ o ) . The conclusion for other initial regions can be obtained in similar arguments. From the definitions in (D.9), (D.13) and (D.16), we have

D.2 Proof of Lemma 3

Lemma 3 restated. Denote ˜ ϕ u = arccos ( -144 π π + k -1 ( 1 ⊤ a ∗ ) 2 k log(4 + k 2 ) ) and

Under the same conditions of Theorem 1. If the learning rate satisfies

for a large absolute constant c > 0 , we have sin ϕ τ a ≥ ϖ 15 √ k with probability at least 1 -δ .

Proof. If ϕ τ a ≤ ˜ ϕ u , the lower bound trivially holds. Next we prove the lower bound for sin ϕ τ a starting from the last initial region, that is ϕ 0 ∈ [ ˜ ϕ u , π ) . From (D.16) and | 1 ⊤ a ∗ | ≥ ∥ a ∗ ∥ , we know that

Let C = log(1+2 k 2 ) 4 , then e 2 C ≤ (1 + 2 k 2 ) 1 2 ≤ 2 k . Denote ϖ 0 = sin ˜ ϕ u 1 ϕ 0 < ˜ ϕ u +sin ϕ 0 1 ϕ 0 ≥ ˜ ϕ u . From the definition of h t and (B.9), we know

where we also used ∥∇ w L ( w i t , a i t ) ∥ ≤ 5 K 2 a . Recall that ˇ v t -v t = η h t -1 , then it follows that

where we used ∥ v t ∥ ≤ 3 in Lemma 13 and ∥ ˇ v t ∥ ≤ ∥ v t ∥ + η ∥ h t ∥ < 4 . For any t , we define a variable

If τ t ( K a ) does not exist, that is min i ∈ [ N ] ϕ i s > π -( (20 K 2 a ) -1 ∧ 1 ) holds for any 0 ≤ s ≤ t -1 , we let τ t ( K a ) = 0 . Then for any τ t ( K a ) ≤ s ≤ t -1 , we always have min i ∈ [ N ] ϕ i s > π -( (20 K 2 a ) -1 ∧ 1 ) such that

Without loss of generality, we assume arg min i ∈ [ N ] ϕ i τ t ( K a ) = 1 . It follows that

Notice that if t ≤ 4 C/η , we have

It holds because η/ 8 ≤ 1 2 and f ( x ) = (1 -x ) 1 /x is decreasing in [0 , 1] . Invoking Lemma 1, we know

where ( i ) holds due to (D.27) and (D.29); ( ii ) follows from (D.30) and (D.26); ( iii ) comes from (D.31) and Lemma 11; and ( iv ) holds due to the assumption D.25. Plugging the choice of C =

log(1 + 2 k 2 ) / 4 , together with ∥ v t ∥ ≤ 3 , we can prove the conclusion

D.3 Self-correction of the First Layer

Theorem 2 restated. For the initial point ( v 0 , a 0 ) with a ⊤ 0 a ∗ ≥ γ a and sin ϕ 0 > 0 , we denote τ v = inf { t ≥ 0 : ϕ v ≤ ˜ ϕ l } , where cos ˜ ϕ l = 144 π π + k -1 ( 1 ⊤ a ∗ ) 2 k log(4 + k 2 ) . If the learning rate satisfies

for a large absolute constant c > 0 , it holds that τ v ≤ O ( 1 η k 2 sin 3 ϕ 0 ) and a ⊤ τ v a ∗ ≥ γ a / 32 with probability at least 1 -δ .

Proof. Find ˜ ϕ f ∈ ( π/ 2 , π ) such that cos ˜ ϕ f = -π π + k -1 ( 1 ⊤ a ∗ ) 2 k log(4 + k 2 ) . We introduce the following stopping time w.r.t. the angle ϕ during the training process,

Initial angle: ϕ 0 ∈ ( ˜ ϕ f , π ) . For any s ≤ τ + ( ˜ ϕ f ) , under the event E good , we have

It means that ϕ i s ≥ π/ 2 holds for any 0 ≤ s ≤ min { τ + ( ˜ ϕ f ) , ˜ O ( η -2 ) } . As a consequence, we have

Notice that if t ≤ 1 η 2 π π -1 , we have

which holds because η π -1 2 π ≤ 1 2 and (1 -x ) 1 /x is decreasing in [0 , 1] . If t ≥ 1 η 2 π π -1 , we also have

because (1 -x ) 1 /x ≤ e -1 for x ∈ [0 , 1] . Together with the dynamic (D.1) and a ⊤ 0 a ∗ ≥ γ a , under the event E good , we have

where ( i ) holds due to (D.8); ( ii ) holds since γ a = 16( π -1) k ( π + k -1) and ( iii ) follows from the condition (D.6). Under the event E good , we know for any s ≤ min { τ + ( ˜ ϕ f ) , ˜ O ( η -2 ) }

holds with probability at least 1 -δ . Recall that ˇ v t -v t = η h t -1 , then it follows that

where ( i ) holds due to I -w ∗ ( w ∗ ) ⊤ is idempotent; and ( ii ) holds due to ∥ w ∗ ∥ = 1 . Invoking Lemma 1 and (D.26), for any s ≤ t -1 ≤ min { τ + ( ˜ ϕ f ) , ˜ O ( η -2 ) } , under the event E good , we can get

where ( i ) holds due to (D.34) and cos ϕ s < 0 for s ≤ τ + ( ˜ ϕ f ) ; ( ii ) follows from assumption (D.32). It implies that sin 2 ϕ s ≥ sin 2 ϕ 0 / 54 for any s ≤ t -1 ≤ min { τ + ( ˜ ϕ f ) , ˜ O ( η -2 ) } . Then we can lower bound sin 2 ϕ i s by

Together with the dynamic (D.12) and lower bound (D.34), we can guarantee

Now we take

where the inequality holds since arcsin(sin ϕ 0 / 6) -arcsin(0) ≥ sin ϕ 0 / 6 due to arcsin ′ ( x ) = 1 √ 1 -x 2 . Plugging it into (D.36), we have

which implies

Initial angle: ϕ 0 ∈ ( ˜ ϕ l , ˜ ϕ f ] . Now assign t 0 = τ + ( ˜ ϕ f ) . Denote

Clearly, P ( ˜ A t 0 ) ≥ 1 -7 δ due to (D.33). Similar to (D.34), we can verify that min i ∈ [ N ] ( a i m ) ⊤ a ∗ ≥ γ a / 32 for any m ≤ s under ˜ A s ∩ E good . In addition, under the event E good , for any ℓ ≥ 0 we have

It means that ˜ ϕ l ≤ ϕ ℓ ≤ ˜ ϕ f for any ℓ ≤ s ≤ min { τ + ( ˜ ϕ l ) , ˜ O ( η -2 ) } under the event ˜ A s ∩ E good . Hence we know (D.18) holds. Consequently, the lower bound (D.22) also holds for 1 N ∑ N i =1 B i s . Therefore, similar to (D.33), we can further verify that

where we used a ⊤ t 0 a ∗ ≥ γ a / 4 . This yields that for any t 0 ≤ s ≤ t -1 ,

which implies that

It follows that

Similar to (D.34), we can guarantee that min i ∈ [ N ] ( a i s ) ⊤ a ∗ ≥ γ a / 64 holds for any s ≤ t -1 ≤ min { τ + ( ˜ ϕ l ) , ˜ O ( η -2 ) } under the event ˜ A t ∩ E good . It follows that

Taking

Plugging t = t 0 + τ 2 , 2 into (D.39), under { t 0 + τ 2 , 2 -1 ≤ τ + ( ˜ ϕ l ) } ∩ ˜ A t 0 + τ 2 , 2 ∩E good , we can get

where we used the condition (D.6). Invoking (D.38), we can have

Recalling t 0 = τ + ( ˜ ϕ f ) and (D.37), we have

By adjusting the level of δ , we can verify that τ v ≤ τ 2 , 1 + τ 2 , 2 such that ϕ τ v ≤ ˜ ϕ l and a ⊤ τ v a ∗ ≥ γ a / 32 with high probability. Now recalling that

we have

where the last inequality holds due to (D.8).

E Convergence with Linear Speedup

E.1 Convergence of the First Layer

Lemma 4 restated. Under the settings in Theorem 3. Suppose the initial point satisfies ϕ 0 ≤ ˜ ϕ l and a ⊤ 0 a ∗ ≥ γ a / 32 . With probability at least 1 -δ , we can guarantee that sin 2 ϕ t ≲ ϵ holds for any ˜ O ( k 2 η -1 ) ≤ t ≤ ˜ O ( η -2 ) .

Proof. Denote the event C t = { a ⊤ s a ∗ ≥ γ a / 128 · ( ∥ a ∗ ∥ 2 ∧ 1) : ∀ s ≤ t } . For any s ≤ t -1 , with probability at least 1 -3 δ , we have

where we used the assumption I ≲ d ϵN min { 1 , k 2 d 1 / 2 N 1 / 2 , k 4 d N 2 ϵ } and ϵ < dk -2 . Together with (D.4), we can guarantee cos ϕ i t ≥ ( 1 ⊤ a ∗ ) 2 k ( π + k -1) holds for any i ∈ [ N ] with probability at least 1 -δ . It follows that for any s ≤ t -1 ,

where ( i ) holds since g ′ ( x ) = -( π -x ) sin x is decreasing in (0 , π/ 2] and ( 1 ⊤ a ∗ ) 2 k ( π + k -1) ≤ √ 3 2 due to Assumption 1; and ( ii ) holds since arccos ′ ( x ) = -1 √ 1 -x 2 . Using the dynamic in Lemma 2, under the event C t -1 , we can get

By adjusting the level of δ , we can guarantee

Next we assume C ˜ O ( η -2 ) happens and hide the indicator.

Invoking (E.1), we have for any s ≤ ˜ O ( η -2 )

where we used ϕ i t < π/ 2 . Invoking Lemma 1, we have

Recall the choices

According to (D.35) and (D.26), with probability at least 1 -δ , we have

where we used η ˜ γ a / 12 < 1 / 2 . Using Lemma 11 and recalling the definition of M 1 ,t , we have the following concentration

holds with probability at least 1 -δ . Moreover, it holds that

which yields with probability at least 1 -δ such that

Plugging (E.6), (E.7) and (E.8) into (E.4), we can guarantee that with probability at least 1 -3 δ , for any T v ≤ t ≤ ˜ O ( η -2 )

where we used I ≲ d ϵN min { 1 , k 2 d 1 / 2 N 1 / 2 , k 4 d N 2 ϵ } and ϵ < d -1 . Then we further have

With this upper bound, we can improve the concentration (E.7) through

Denote H = cηK 4 a ( √ Id + I ) √ log( Nd/δ ) . Using (E.9), we can further refine the bound (E.6) by

It follows that

Combining (E.10), (E.11) and (E.8), we can show that for any 2 T v ≤ t ≤ ( ηK 2 a ) -2

Repeating (E.10)-(E.12) for K = log(1 /ϵ ) times, we can guarantee that for any ( K +1) T v ≤ t ≤ ˜ O ( η -2 )

E.2 Convergence of the Second Layer

Lemma 5 restated. Under the choice for η and conditions for ϵ in Theorem 3. Suppose sin 2 ϕ t ≤ ϵ holds for any 0 ≤ t ≤ ˜ O ( η -2 ) . With probability at least 1 -δ , we can guarantee that ∥ a t -a ∗ ∥ 2 ≲ ϵ holds for any ˜ O ( η -1 ) ≤ t ≤ ˜ O ( η -2 ) .

Proof. Let h i t = v i t -v t . Similar to (D.35), we can verify that for any i ∈ [ N ]

where we used Lemma 14 and sin 2 ϕ t ≤ ϵ . It further implies that

Further, for any 0 ≤ t ≤ ˜ O ( η -2 ) and i ∈ [ N ] we have

where ( i ) holds since g ′ ( x ) = -( π -x ) sin x ; and ( ii ) holds due to arcsin ′ ( x ) = 1 √ 1 -x 2 . Applying this bound, we can get

where the last inequality holds due to Lemma 12 such that ∥ a t -a ∗ ∥ ≤ 6 ∥ a ∗ ∥ . In addition, it holds that

where we used the relation (E.14). Denote ˜ a = a -a ∗ -η 1 ∑ N ∇ L ( w i , a i ) , then we have

Using Lemma 12 again, we can verify that

Similar to (A.17) with α s = ( 1 -η ( π -1) 4 π ) t -s , we can guarantee

with probability at least 1 -δ . Now take T a = 1 η 8 π π -1 log ( 1 ϵ ) . Plugging (E.19) into (E.17), together with the concentration inequality (A.19), we have for any T ≤ t ≤ O ( η -2 )

With this upper bound, we can refine (E.18) by

Consequently, the bound (E.19) can be improved to

Then we can guarantee that for any 2 T a ≤ t ≤ ˜ O ( η -2 ) ,

Repeating the refinement from (E.20) to (E.21) for at most log(1 /ϵ ) times, we can show finish the proof.

E.3 Conclusion

Theorem 3 restated. Suppose the initial point ( v 0 , a 0 ) satisfies a ⊤ 0 a ∗ ≥ γ a / 32 and ϕ 0 ≤ ˜ ϕ l . For any ϵ > 0 , we choose η = 1 ck 2 ∥ a ∗ ∥ 2 Nϵ d log( dN/ϵδ ) for some absolute constant c > 0 . If I ≲ d ϵN min { 1 , k 2 d 1 / 2 N 1 / 2 , k 4 d N 2 ϵ } and ϵ < min { N -1 , d -1 , dk -2 } , then ℓ ( v T , a T ) = O ( ϵK 2 a ) holds with probability at least 1 -δ where T = ˜ O ( dk 4 Nϵ ) .

Proof. According to the dynamic of a t , we know

It follows from (E.14) that

Using Lemma 11(1), it holds that

Plugging (E.23) and (E.24) into (E.22), and taking T 1 = 1 η 2 π π + k -1 log ( 1 ϵ ) , we can have

where we used | 1 ⊤ a 0 | ≤ | 1 ⊤ a ∗ | and ϵ < 1 /k . Invoking Lemma 6 and the normalized network in (2), we know

Similar to (E.14), for any t ≥ T v , we also have

Invoking Lemma 4 and 5 to (E.26), together with (E.27), we conclude that for T = ˜ O ( dk 4 Nϵ ) ,

with probability at least 1 -δ .

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the global convergence of local SGD for One-hidden-layer Convolutional Neural Networks. The authors present solid theoretical understanding. They show that without overparameterization and injecting noise,  local SGD have global convergence by new proof techniques and new understanding.\n\nSTRENGTHS:\nThe paper is clear and well-written. \nUnderstanding the optimization dynamics of gradient-based method is a significant theoretical issue. \nThis paper provides a solid theoretical analysis without highly over-parameterization and injecting noise, which results in the highly non-convexity.\nTheoretically, the authors provide a novel understanding of the training dynamics by dividing them into two phases: self-correction and convergence.\n\nWEAKNESSES:\nThe article's significance is limited due to the restriction of input data to Gaussian distributions. However, this limitation is not a major concern since the problem itself is non-convex, even in this simplified scenario. Additionally, I have significant concerns regarding Assumption 1 about the target function, which I will discuss in detail in the Questions section.\n\nQUESTIONS:\nMy main question pertains to Assumption 1 about the teacher traget. \nIt appears that this assumption suggestes that stronger conditions on $\\boldsymbol{a}^*$ are required for larger values of $k$. \nI am curious to know whether this requirement is essential or merely technical. \nIn other words, if one use a sufficiently large $k$ that contradicts this assumption, do the training dynamics in the last region (in the proof of Theorem 1) fundamentally change? \nThe authors could attempt to demonstrate this aspect experimentally or theoretically.\n\nLIMITATIONS:\nThe limitations of this article lie in the requirements for data and target Assumption. \nHowever, these limitations should not be grounds for rejection. \nGiven our limited understanding of the training dynamics of nonconvex optimization, studying such simple setups can still provide valuable insights and contribute to the field.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-04"}

comments: SUMMARY:
This paper provides a rigorous theoretical justification on how the local SGD (without the injection of noise) can find the global minima for CNN without relying on the NTK-type analysis (i.e., overparametrization) under federated learning framework.

STRENGTHS:
A.	The contributions of the paper are clearly written.  
B.	Nice simulation results which are consistent with the theoretical predictions.
C.	Like the part where the authors divide the landscape of the objective function into several regions. This type of analysis is not something that I can find in the current literature.

WEAKNESSES:
A.	Writing is little bit technical. It would be great if the authors can draw some cartoons of dynamics of each layer for the global convergence.

QUESTIONS:
A.	Why is the f (Z, w, a) CNN? It looks to me like a shallow fully connected neural network. What is the exact structure of $w \in \mathbb{R}^{d}$? 
B.	In line 90, the paper says the algorithm starts from the random initialization from the same initial region as in [13, 63]. In line 203, it says arbitrary initialization. I was bit confused while reading the paper on this point. And it becomes clear while reading line 312~313. 
C.	Just to clear my understanding, so your result is saying that weights in the first layer converges (in polynomial time), and then the weights in the second layer converges afterward?

LIMITATIONS:
This work has no negative societal impact.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-04

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper asserts that local SGD achieves convergence to the global minimum in the presence of Gaussian input. The experimental setup involves a two-layer student model, where the ground truth is generated by a two-layer teacher model. The authors highlight their main contribution as a proof that does not necessitate noise injection.\n\nSTRENGTHS:\nThe paper proved global convergence of local SGD without noise injection.\n\nWEAKNESSES:\n1. The presentation of the results is unclear and the proof is difficult to follow. It would be beneficial for the authors to provide a more intuitive interpretation of the terms, lemmas, and equations. For instance:\n- Could the authors clarify the purpose and significance of the intermediate step $\\check{v}_{t+1}$?\n- How does Theorem 1 demonstrate the self-correction of the second layer?\n\n2. The realism of the assumptions made by the authors is questionable. For instance:\n- What is the probability that a randomly selected $a^*$ satisfies equation (5)? It appears to be very low in my estimation.\n\n3.  I'm uncertain why the authors chose to consider CNN instead of linear layers. It seems to add unnecessary confusion to the problem.\n\n4. What sets \"local\" SGD apart from regular SGD? It would be helpful for the authors to clarify the distinctive aspects (in terms of proof) of \"local\" SGD in comparison to regular SGD.\n\nQUESTIONS:\nPlease review the suggested revisions for the mentioned weaknesses:\n\n\nMinor comments:\n\nIn line 167, could you please clarify the distinction between L(w, a; Z) and l(w, a; Z)?\nIn line 169, \"wights\" should be corrected to \"weights\".\nIn line 215, \"learning\" should be changed to \"learning rate\".\n\nLIMITATIONS:\nI believe it is important for the authors to provide justification for the assumptions made. Without such justification, the scope and applicability of the proposed results may be limited.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-05"}

comments: SUMMARY:
This paper considers a federated learning setting for a  one-hidden-layer convolutional neural network without overparameterization. It is proven that vanilla local SGD (where in each iteration each node updates the model by the local gradient and then all nodes synchronize the model parameters) can converge to a global minimum.

STRENGTHS:
1. The convergence analysis does not require over-parameterization (i.e., in the NTK region) and injected noise.    
     
2. Based on a careful landscape study, the paper proposes a self-correction mechanism that ensures the algorithm enters a good landscape region in polynomial time. This explains why local SGD can converge in practice even though the landscape is not well-conditioned.

WEAKNESSES:
1. The online learning assumption, i.e., each node samples data i.i.d. from the same distribution in each iteration, is crucial to the analysis of this paper, but it deviates from the practical settings, where each node computes its gradient according to a given set of local data.

2. Why is the proposed model a CNN? It seems a fully-connected network.

3. The weight-normalization technique twists the vanilla SGD.

QUESTIONS:
Please see "weakness".

LIMITATIONS:
I don't think the authors adequately addressed the limitations.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-08-02