SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yajie Bao<sup>1</sup> Amarda Shehu<sup>2</sup> Mingrui Liu<sup>2*</sup>
$^{1}$ School of Mathematical Sciences, Shanghai Jiao Tong University, Shanghai, 200240
$^{2}$ Department of Computer Science, George Mason University, Fairfax, VA 22030
baoyajie2019stat@sjtu.edu.cn, {ashehu,mingruil}@gmu.edu
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.
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:
$$
f (\mathbf {Z}, \boldsymbol {w}, \boldsymbol {a}) = \sum_ {j = 1} ^ {k} a _ {j} \sigma \left(\mathbf {Z} _ {i} ^ {\top} \boldsymbol {w}\right),
$$
where  $\mathbf{Z} = (\mathbf{Z}_1, \dots, \mathbf{Z}_k) \in \mathbb{R}^{d \times k}$  is the input matrix,  $\boldsymbol{w} \in \mathbb{R}^d$  is the weight,  $\boldsymbol{a} \in \mathbb{R}^k$  is the output weight, and  $\sigma(x) = \max\{x, 0\}$  denotes the ReLU activation function. A good property of this network is positive homogeneity, i.e.,  $f(\mathbf{Z}, c\boldsymbol{w}, \boldsymbol{a}/c) = f(\mathbf{Z}, \boldsymbol{w}, \boldsymbol{a})$  holds for any  $c > 0$ . We assume each entry of the input  $\mathbf{Z}$  is independently sampled from a standard Gaussian distribution. Then the response is generated by a noiseless teacher network:  $y = f(\mathbf{Z}, \boldsymbol{w}^*, \boldsymbol{a}^*)$ . Without loss of generality, we further assume  $\| \boldsymbol{w}^* \| = 1$ . We hope to learn a student network by collaboratively minimizing the following mean square loss among  $N$  local machines:
$$
L (\boldsymbol {w}, \boldsymbol {a}) = \frac {1}{2} \mathbb {E} \left[ (y - f (\mathbf {Z}, \boldsymbol {w}, \boldsymbol {a})) ^ {2} \right] = \frac {1}{2} \mathbb {E} \left[ (f (\mathbf {Z}, \boldsymbol {w} ^ {*}, \boldsymbol {a} ^ {*}) - f (\mathbf {Z}, \boldsymbol {w}, \boldsymbol {a})) ^ {2} \right]. \tag {1}
$$
Obviously,  $(\pmb{w}^{*},\pmb{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(\pmb {w},\pmb {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  $|\mathbf{1}^{\top}\mathbf{a}^{*}| / \| \mathbf{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:
1. We introduce a new self-correction mechanism of local SGD under a condition on  $\pmb{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  $(\pmb{w}^{*},\pmb{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.
2. 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.
3. 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  $\left| \mathbf{1}^{\top} \mathbf{a}^{*} \right| / \left\| \mathbf{a}^{*} \right\|$  is almost necessary to show the convergence with almost arbitrary initialization.
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  $\ell_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
Initialize  $\pmb{v}_0\in \mathbb{S}^{d - 1}$  and  $a_0\in \mathbb{B}^k\left(\frac{\left|\mathbf{1}^\top\pmb{a}^*\right|}{\sqrt{k}}\right)$    
for  $r = 0,\dots ,R - 1$  do for  $i = 1,\ldots ,N$  do Synchronization:  $\pmb{a}_{t_r}^i\gets \pmb{a}_{t_r}$  and  $\pmb{v}_{t_r}^i\gets \pmb{v}_{t_r}$  . for  $t = t_r,\dots ,t_{r + 1} - 1$  do Sample  $\mathbf{Z}_t^i$  from the standard Gaussian distribution and generate the response  $y_{t}^{i}$  Update  $\pmb{a}_{t + 1}^{i} = \pmb{a}_{t}^{i} - \eta \nabla_{\pmb{a}}\ell (\pmb{v}_{t}^{i},\pmb{a}_{t}^{i};\mathbf{Z}_{t}^{i},y_{t}^{i})$  Update  $\pmb{v}_{t + 1}^{i} = \pmb{v}_{t}^{i} - \eta \nabla_{\pmb{v}}\ell (\pmb{v}_{t}^{i},\pmb{a}_{t}^{i};\mathbf{Z}_{t}^{i},y_{t}^{i})$  end for end for Update  $\pmb{a}_{t_{r + 1}} = \frac{1}{N}\sum_{i = 1}^{N}\pmb{a}_{t_{r + 1}}^{i}$  and  $\pmb{v}_{t_{r + 1}} = \frac{1}{N}\sum_{i = 1}^{N}\pmb{v}_{t_{r + 1}}^{i}$    
end for
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  $\| \cdot \|$  the Euclidean norm and  $\langle \cdot, \cdot \rangle$  by the inner product.  $\mathbb{S}^{d-1}$  denotes the  $d$ -dimensional unit sphere and  $\mathbb{B}^k(\rho)$  denotes the  $k$ -dimensional ball with center zero and radius  $\rho$ . For two vectors  $\pmb{v}, \pmb{w} \in \mathbb{R}^d$ , denote  $\angle(\pmb{v}, \pmb{w}) \in [0, \pi]$  the angle between  $\pmb{v}$  and  $\pmb{w}$ . The uniform distribution is denoted by  $\mathrm{Unif}(\cdot)$ . Moreover, we use  $\widetilde{O}$  to hide logarithmic factors.
We adopt the weight-normalization technique [46] to the first layer by re-parametrizing  $\boldsymbol{w} = \boldsymbol{v} / \| \boldsymbol{v}\|$ , which leads to the following prediction model:
$$
f (\mathbf {Z}, \boldsymbol {v}, \boldsymbol {a}) = \sum_ {j = 1} ^ {k} a _ {j} \frac {\sigma \left(\mathbf {Z} _ {i} ^ {\top} \boldsymbol {v}\right)}{\| \boldsymbol {v} \|}. \tag {2}
$$
Given any sample  $(\mathbf{Z},y)$ , we denote the empirical loss by
$$
\ell (\boldsymbol {v}, \boldsymbol {a}; \mathbf {Z}, y) = \frac {1}{2} (y - f (\mathbf {Z}, \boldsymbol {v}, \boldsymbol {a})) ^ {2} = \frac {1}{2} \left(f (\mathbf {Z}, \boldsymbol {w} ^ {*}, \boldsymbol {a} ^ {*}) - f (\mathbf {Z}, \boldsymbol {v}, \boldsymbol {a})\right) ^ {2}. \tag {3}
$$
In the distributed environment, suppose there are  $N$  local machines sharing the same teacher model  $f(\mathbf{Z},\boldsymbol{w}^{*},\boldsymbol{a}^{*})$ . It means that given any local input  $\mathbf{Z}^i$ , the response is  $y^{i} = f(\mathbf{Z}^{i},\boldsymbol{w}^{*},\boldsymbol{a}^{*})$ . Let  $\mathcal{I} = \{t_0,\dots,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]:  $\pmb{v}_0\in \mathbb{S}^{d - 1}$  and  $\pmb{a}_{0}\in \mathbb{B}^{k}(|\mathbf{1}^{\top}\pmb{a}^{*}| / \sqrt{k})$ . In each round, the  $i$ -th machine runs  $I$  steps of SGD using the stochastic gradients  $\nabla_{\pmb{v}}\ell (\pmb{v}_{t}^{i},\pmb{a}_{t}^{i};\pmb{\mathbf{Z}}_{t}^{i},y_{t}^{i})$  and  $\nabla_{\pmb{a}}\ell (\pmb{v}_{t}^{i},\pmb{a}_{t}^{i};\pmb{\mathbf{Z}}_{t}^{i},y_{t}^{i})$  computed by the local input  $(\mathbf{Z}_t^i,y_t^i)$ . 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.
We first introduce two auxiliary sequences  $\pmb{v}_t = \sum_{i=1}^N \pmb{v}_t^i / N$  and  $\pmb{a}_t = \sum_{i=1}^N \pmb{a}_t^i / N$ , which often appears in the analysis of local SGD [47, 53, 28]. Notice that  $(\pmb{v}_t, \pmb{a}_t) = (\pmb{v}_t^i, \pmb{a}_t^i)$  for  $i \in [N]$  if  $t$  is the synchronization time, i.e.,  $t \in \mathcal{I}$ . Denote  $\pmb{w}_t^i = \pmb{v}_t^i / \| \pmb{v}_t^i\|$  the local normalized weight and  $\mathbf{P}_t^i = (\mathbf{I} - \pmb{w}_t^i(\pmb{w}_t^i)^\top) / \| \pmb{v}_t^i\|$  the respective projection matrix. Now let  $L(\pmb{w}, \pmb{a}; \mathbf{Z})$  be the empirical version of the loss function  $L(\pmb{w}, \pmb{a})$  defined in (1), that is  $L(\pmb{w}, \pmb{a}; \mathbf{Z}) = \frac{1}{2}(f(\mathbf{Z}, \pmb{w}^*, \pmb{a}^*) - f(\mathbf{Z}, \pmb{w}, \pmb{a}))^2$ . Since  $y = f(\mathbf{Z}, \pmb{w}^*, \pmb{a}^*)$ , we will hide  $y$  and write  $\ell(\pmb{v}, \pmb{a}; \mathbf{Z}) \equiv \ell(\pmb{v}, \pmb{a}; \mathbf{Z}, y)$  hereafter. Notice that  $L(\pmb{w}, \pmb{a}; \mathbf{Z}) = \ell(\pmb{v}, \pmb{a}; \mathbf{Z})$  for  $\pmb{w} = \pmb{v} / \| \pmb{v}\|$ . By recalling the normalized model (2), we have
$$
\nabla_ {\boldsymbol {v}} \ell (\boldsymbol {v} _ {t} ^ {i}, \boldsymbol {a} _ {t} ^ {i}; \mathbf {Z} _ {t} ^ {i}) = \mathbf {P} _ {t} ^ {i} \nabla_ {\boldsymbol {w}} L (\boldsymbol {w} _ {t} ^ {i}, \boldsymbol {a} _ {t} ^ {i}; \mathbf {Z} _ {t} ^ {i}), \quad \nabla_ {\boldsymbol {a}} \ell (\boldsymbol {v} _ {t} ^ {i}, \boldsymbol {a} _ {t} ^ {i}; \mathbf {Z} _ {t} ^ {i}) = \nabla_ {\boldsymbol {a}} L (\boldsymbol {w} _ {t} ^ {i}, \boldsymbol {a} _ {t} ^ {i}; \mathbf {Z} _ {t} ^ {i}).
$$
According to Algorithm 1, the updates of auxiliary sequences can be written as
$$
\boldsymbol {v} _ {t + 1} = \boldsymbol {v} _ {t} - \eta \frac {1}{N} \sum_ {i = 1} ^ {N} \mathbf {P} _ {t} ^ {i} \nabla_ {\boldsymbol {w}} L (\boldsymbol {w} _ {t} ^ {i}, \boldsymbol {a} _ {t} ^ {i}; \mathbf {Z} _ {t} ^ {i}), \quad \boldsymbol {a} _ {t + 1} = \boldsymbol {a} _ {t} - \eta \frac {1}{N} \sum_ {i = 1} ^ {N} \nabla_ {\boldsymbol {a}} L (\boldsymbol {w} _ {t} ^ {i}, \boldsymbol {a} _ {t} ^ {i}; \mathbf {Z} _ {t} ^ {i}).
$$
For ease of technical presentation, we denote the averaged noise terms by  $\pmb{\xi}_{t} = \frac{1}{N}\sum_{i = 1}^{N}\mathbf{P}_{t}^{i}\pmb{\xi}_{t}^{i}$  and  $\pmb{\epsilon}_{t} = \frac{1}{N}\sum_{i = 1}^{N}\pmb{\epsilon}_{t}^{i}$ , where  $\pmb{\xi}_{t}^{i} = \nabla_{\pmb{w}}L(\pmb{w}_{t}^{i},\pmb{a}_{t}^{i};\pmb{Z}_{t}^{i}) - \nabla_{\pmb{w}}L(\pmb{w}_{t}^{i},\pmb{a}_{t}^{i})$  and  $\pmb{\epsilon}_{t}^{i} = \nabla_{\pmb{a}}L(\pmb{w}_{t}^{i},\pmb{a}_{t}^{i},\pmb{Z}_{t}^{i}) - \nabla_{\pmb{a}}L(\pmb{w}_{t}^{i},\pmb{a}_{t}^{i})$  are the local noises in stochastic gradients. In addition, for the iterates  $\pmb{v}_{t}^{i}$  and  $\pmb{v}_{t}$ , we write  $\phi_t^i = \angle (\pmb {v}_t^i,\pmb {w}^*)$  and  $\phi_t = \angle (\pmb {v}_t,\pmb {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  $\mathbf{P}_t = \frac{1}{\|\mathbf{v}_t\|} (\mathbf{I} - \mathbf{w}_t\mathbf{w}_t^\top)$  by the global projection matrix, where  $\mathbf{w}_t = \mathbf{v}_t / \| \mathbf{v}_t\|$ . The proofs of this subsection are deferred to Appendix A.1.
Lemma 1. Let  $\check{\pmb{v}}_{t + 1} = \pmb{v}_t - \eta \left(\mathbf{P}_t\frac{1}{N}\sum_{i = 1}^N\nabla_{\pmb{w}}L(\pmb{w}_t^i,\pmb{a}_t^i) + \pmb{\xi}_t\right)$  and  $\check{\phi}_{t} = \angle (\check{\pmb{v}}_{t},\pmb{w}^{*})$ . The first layer in local SGD satisfies that
$$
\left\| \boldsymbol {v} _ {t + 1} \right\| ^ {2} \sin^ {2} \phi_ {t + 1} = \left(1 - \eta \lambda_ {t} \cos \phi_ {t}\right) ^ {2} \left\| \boldsymbol {v} _ {t} \right\| ^ {2} \sin^ {2} \phi_ {t} - 2 \eta M _ {1, t} + \eta^ {2} M _ {2, t} + H _ {t}, \tag {4}
$$
where  $\lambda_{t} = \frac{1}{N}\sum_{i = 1}^{N}\frac{\pi - \phi_{t}^{i}}{2\pi}\frac{(\pmb{a}_{t}^{i})^{\top}\pmb{a}^{*}}{\|\pmb{v}_{t}\|^{2}}$ ,  $H_{t} = \| \pmb{v}_{t + 1}\|^{2}\sin^{2}\phi_{t + 1} - \| \check{\pmb{v}}_{t + 1}\|^{2}\sin^{2}\check{\phi}_{t + 1}$  and
$$
M _ {1, t} = \left(\boldsymbol {v} _ {t} - \eta \mathbf {P} _ {t} \frac {1}{N} \sum_ {i = 1} ^ {N} \nabla_ {\boldsymbol {w}} L (\boldsymbol {w} _ {t} ^ {i}, \boldsymbol {a} _ {t} ^ {i})\right) ^ {\top} \left(\mathbf {I} - \boldsymbol {w} ^ {*} (\boldsymbol {w} ^ {*}) ^ {\top}\right) \boldsymbol {\xi} _ {t}, M _ {2, t} = \boldsymbol {\xi} _ {t} ^ {\top} \left(\mathbf {I} - \boldsymbol {w} ^ {*} (\boldsymbol {w} ^ {*}) ^ {\top}\right) \boldsymbol {\xi} _ {t}.
$$
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\phi_t^i$ , which cannot characterize the dynamic of global quantity  $\sin \phi_t$  due to the nonlinearity. Here we introduce a new intermediate variable  $\check{v}_{t + 1}$  and find an equality (4) to show the recursive relation between  $\| \pmb{v}_{t + 1}\|^2\sin^2\phi_{t + 1}$  and  $\| \pmb {v}_t\| ^2\sin^2\phi_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 \phi_t$  by upper bounding the discrepancy term  $H_{t}$ . Let us assume the last three terms in (4) are negligible and  $(a_{t}^{i})^{\top}a^{*} > 0$  for any  $i\in [N]$ :
(1) When  $\phi_t > \pi / 2$ ,  $\| \pmb{v}_t \|^2 \sin^2 \phi_t$  will continuously increase since  $\lambda_t \cos \phi_t < 0$ . It indicates that  $\phi_t$  can decrease to  $\pi / 2$  if  $\| \pmb{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.
(2) When  $\phi_t < \pi/2$ ,  $\|\pmb{v}_t\|^2 \sin^2 \phi_t$  will continuously decrease to zero. It indicates that  $\phi_t$  can converge to zero if  $\|\pmb{v}_t\|$  is lower bounded by a constant. The initial region with  $\phi_0 < \pi/2$  and  $\pmb{a}_0^\top \pmb{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 = |\mathbf{1}^\top \mathbf{a}^*|^2 - (\mathbf{1}^\top \mathbf{a}^*)(\mathbf{1}^\top \mathbf{a}_0)$  and  $g(\phi) = (\pi - \phi) \cos \phi + \sin \phi$ . For local SGD algorithm started with the initial point  $(v_0, a_0)$ , the second layer satisfies that