Sketching for Distributed Deep Learning: A Sharper Analysis

We thank the reviewer for the suggestion. We will include a table comparing our results with [Song et al. 2023] in the final paper. Here, we provide brief summary:

Q1. Comparison with [Song et al. 2023]: Before comparing, we briefly summarise the difference in our setup, assumptions and results:

1. Choice of $f$: While [Song et al. 2023] provides results of the $\operatorname{desk}-\operatorname{sk}$ framework for convex/non-convex function $f$ satisfying $\beta-$smoothness, our work looks at the case when $f$ is a overparametrized neural-network and exploits the RSS property of such networks.

2. Assumption on Loss function $\mathcal{L}(\theta)$: [Song et al. 2023] provides results under various assumptions (convex, non-convex, strongly-convex). In our case, motivated by recent works that show neural networks satisfy the PL condition [Zhou et al.],[Hardt et al] , we assume loss function $\mathcal{L}$ satisfies PL condition. As mentioned by the reviewer, our results can be compared with the Strong Convexity results by [Song et al. 2023].

3. Uniform bound on the gradient:: Note that, our analysis for single-step case does not require uniform bound on the gradient showing that we remove the dimensional dependence using RSS. For $K>1$ steps, under PL condition, we require some additional assumptions to show convergence. Motivated by works showing neural losses satisfy PL condition [Zhou et al.],[Hardt et al] we have done the current analysis only for the PL case leaving further extensions for future work.

Under our assumptions and setup (Section 3 in our paper), we provide comparison with [Song et al. 2023] in Table: 1 for K=1 local steps. We achieve a ambient-dimension independent linear convergence under only RSS assumption. We also provide additonal results for non-convex setting and show that we achieve dimension-independence under this setting too. For $K>1$, we compare our results in Table:2.

Table 1:Comparison of iteration complexities and assumptions for K=1 local-steps. $c_s,c_l$ are constants defined in Assumption 3.2 and $\varrho,c_H$ are defined in Lemma C.1,Lemma C.2 respectively. Note that $\varrho$ and $c_H$ are $\mathcal{O}(poly(L))$ and can be seen as independent of ambient dimension for wide networks. Under our assumptions and setup (Section 3), according to Lemma C.4 the loss function $\mathcal{L}$ can be shown to be $\beta-$smooth where $\beta = c_s\varrho^2 + \frac{c_lc_H}{\sqrt{m}}$. $\varepsilon$ is the sketching error and $\epsilon$ refers to the optimization error tolerance.

Table 2: Comparison of iteration complexities and assumptions for $K>1$ local-steps. $c_s,c_l$ are constants defined in Assumption 3.2 and $\varrho,c_H$ are defined in Lemma C.1,Lemma C.2 respectively. Note that $\varrho$ and $c_H$ are $\mathcal{O}(poly(L))$ and can be seen as independent of ambient dimension for wide networks.

---

Q2. Regarding Figure 2 Thank you for your valuable feedback and for understanding the intention behind our experimental design. We appreciate your insightful comments regarding the comparison between the sketching algorithm and FetchSGD. We will revise the main body of the paper to include Figure 2 and provide a comprehensive discussion on the performance comparison.

---

Q3.Typos: On line 592, end of page 15, the matrix norm is the spectral norm. We will add this in the notation for the final draft. We thank the reviewer for pointing out the error in citation and typos. We will update them in the revised manuscript.

[Zhou et al.] Characterization of gradient dominance and regularity conditions for neural
networks, 2017. URL https://arxiv.org/abs/1710.06910. [Hardt et al] M. Hardt and T. Ma. Identity matters in deep learning. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017. URL https://openreview.net/forum?id=ryxB0Rtxx

We thank the reviewer for carefully reading our paper and providing constructive feedback. Below we provide answers to each of the concerns:

Q1: SGD v/s GD: Firstly we point out our theoretical result in Theorem 1 must be compared to Theorem E.1 in [Song et al. 2023] which also studies GD. Further, in Section 3.2 of our paper we explicitly show that any analysis that uses strong smoothness will eventually pick up dimension dependence. Our aim in the paper is to show the inherent ambient dimension dependence in previous sketching based analysis ([Song et al. 2023]) which can be shown in the GD analysis and will also carry over to the SGD analysis. Our main result is to get rid of this dependence and we provide a sharper analysis that depends on the sum of absolute values of eigenvalues of the Hessian, which is potentially much smaller. Such an analysis can be extended to SGD under suitable assumptions.

---

Q2:On the use of SGD in experimental section: We do not claim that we provide superior performance compared to FetchSGD. We would like to point out that [Rothchild et al. 2020] acknowledges in Appendix B.3 that a $\operatorname{sk}$-$\operatorname{desk}$ based framework is comparable in performance to the uncompressed version, however they are unable to provide a theoretical analysis and thus use an additonal Top-r sparsification step. We fill in this gap through our results. Addionally, the Top-r sparsification step in FetchSGD makes their algorithm non-linear. On the other hand, our experiments show that $\operatorname{sk}$-$\operatorname{desk}$ framework achieves comparable performance with FetchSGD (Figure 2 in the Appendix) while retaining linearity which allows easy integration with Differential Privacy.

---

Q3: Similarity to SGD: We provide the theoretical result for $K-$step case with a constant learning rate. As a result, there is a term that decreases exponentially fast and the variance term that is dependent on the choice of learning rate. Hence, our result is similar to analysis of SGD with constant learning rate due to gradient noise. In the analysis of Theorem 2 in our paper by choosing a decaying learning rate $\eta_t = \mathcal{O}(1/T)$, as is common in SGD analysis, we can show that the result converges to $0$, i.e. $\mathcal{L}(\theta_T) -\mathcal{L}(\theta^{*}) = \mathcal{O}\left(1/T\right)$ and achieve sub-linear convergence rates as is standard in literature[ Karimi et al]. We thank the reviewer for pointing this out and we will add a remark clarifying this in the final draft.

---

Q4. Notation in Assumption 3.3: Yes, the notation is taken from the literature [Banerjee et al.], [Liu et al.]. $\theta_t$ refers to the parameter vector $\theta$(as in (3) in our paper) at iteration $t$ and thus $\theta_0$ refers to the weight at initialization. We will add a clarifying remark in the final draft.

---

Q5. Clarification on the experimental setup: We thank the reviewer for pointing out the confusion around the computing requirement for the experiments. We managed to run our experiments on one GPU by (1) accumulating the model updates on the fly once we compute the model update of a client and (2) moving the model updates and/or accumulated updates from the GPU to the CPU. Accumulating the updates on the fly, instead of computing the updates in parallel and accumulating them all together, is only for the sake of compute efficiency due to the compute constraints of our research environment. Rest assured, this modification does not lead to any change in the empirical results. However, to clarify any confusion around the computing needs, we will add this explanation to the revised paper.

---

Q6.Symbolic notation and citation errors: We indeed re-use some of the results and notations from [Banerjee et al.], [Liu et al.] in lines 137 to 153; We have stated this explictly in line 589. Thank you for pointing the errors in citations, we will add publication details to the references in the final draft.

[Karimi et al] Linear convergence of gradient and proximal-gradient methods under the polyak-ojasiewicz condition, 2020. URL https://arxiv.org/abs/1608.04636

[Banerjee et al.] Restricted strong convexity of deep learning models with smooth activations. In The Eleventh International Conference on Learning Representations, 2023.

[Liu et al.] On the linearity of large non-linear models: when and why the tangent kernel is constant. Advances in Neural Information Processing Systems, 33:15954–15964, 2020.

1. Convergence result of GD (Question 2): The final result in Theorem 2 in our paper has been shown in a form similar to Theorem 5.1 [Song et al. 2023] . Note that, our result can be made arbitrarily small, either by choosing a small step-size (for constant step-size) or a decreasing step-size. We provide one such step size choice for the communication efficiency comparison (Theorem 3 in our paper) to make the final result $\mathcal{L}(\theta_T) -\mathcal{L}(\theta^{*}) \le \epsilon$ and the statement for the decreasing step size in the rebuttal provided above (Question 3).

   Explanation for the residual term in the K-step analysis: In the $K-$local-step GD , note that, the average of local gradients is no longer the true gradient (as in the $1$-local-step analysis). To see this, each client performs multiple GD step on individual loss function . Following the notation in our paper, $\theta_{c,t,k} = \theta_{c,t,k-1} - \eta_{\text{local}}\nabla \mathcal{L_c}(\theta_{c,t,k-1}) ~\forall k=1,\cdots,K$. This, causes a drift in the local parameters $\theta_{c,t,k}$ at each of the clients and thus, $K$-local-step GD is not the same as performing GD on $\mathcal{L}$. Under the bounded gradient assumption and Lipschitz smoothness of the objective function, we can bound this drift and that leads to the residual term in each iteration ( line 652 eq 80 in our paper) and (eq 20 in [Song et al. 2023]). Accumulating the residual over all $T$ iterations, leads to the final result which depends on $\eta$.

   ---

2. Notation in Assumption 3.3 (Question 3): We thank the reviewer for pointing this out. We will rectify this in the main draft. For clarification, following eq(3) in our paper,

$$\theta := (({\vec{W}}^{(1)})^\top,\ldots,(\vec{W}^{(L)})^\top, \mathbf{v}^\top )^\top \in \mathbb{R}^p.$$

The iterate at iteration $t$ is give as:

$$\theta_t := (({\vec{W_t}}^{(1)})^\top,\ldots,(\vec{W_t}^{(L)})^\top, \mathbf{v_t}^\top )^\top \in \mathbb{R}^p.$$

Following this, the iterate at initialization $\theta_{0}$ is given as:

$$\theta_0 := (({\vec{W_0}}^{(1)})^\top,\ldots,(\vec{W_0}^{(L)})^\top, \mathbf{v}_0^\top )^\top \in \mathbb{R}^p$$

Assumption 3.3, defines the initialization parameters $\vec{W_0}^{(l)}~\forall l\in L$ and $\mathbf{v_0}$.

---

3. Running FetchSGD with full-batch: Since we provide an analysis for GD, we run FetchSGD with full batch setting for fair comparison. We recognize that FetchSGD is typically designed for mini-batch training, but given that the theoretical framework in [Rothchild et al. 2020] can be extended to GD under similar assumptions, we believe it is appropriate to compare FetchSGD with our algorithm under a full-batch setting. That said, we would be grateful if the reviewer could elaborate on their concerns regarding the use of FetchSGD in a full-batch setting. This would help us better understand and address any issues in our revised manuscript. A goal of our work is to show that the Top-r sparsification and error-feedback mechanism is not required for the analysis of sketched distributed learning algorithms and thus, we show the comparison with FetchSGD. Since [Rothchild et al. 2020] are the first to propose using sketching for distributed learning, we refer to the Top-r + Sketching version of the $\mathop{\text{sk}}-\mathop{\text{desk}}$ algorithm as FetchSGD without Error Feedback in Figure 1. We agree with the reviewer that error-feedback is required for biased compressed algorithms and we will revise the main body of the paper to include Figure 2 and provide a more comprehensive discussion on the performance comparison.

---

4. Usage of notation and assumptions from previous papers: We have cited [28] right at line 137 before using the notation. However, based on the feedback of the reviewer, we will make this more explicit in the subsequent draft. Note that this is only mathematical setup and assumptions, and therefore they would be same across papers that use the same model and assumptions. This is only generic notation that we are re-purposing for consistency and we do not claim any novelty here. We are using the notation from previous papers to make everything streamlined and we have cited any results that we are using appropriately.

[Song et al. 2023] . Song, Y. Wang, Z. Yu, and L. Zhang. Sketching for first order method: Efficient algorithm for low- bandwidth channel and vulnerability. In International Conference on Machine Learning, pages 32365– 32417. PMLR, 2023. [Rothchild et al. 2020] D. Rothchild, A. Panda, E. Ullah, N. Ivkin, I. Stoica, V. Braverman, J. Gonzalez, and R. Arora. FetchSGD: Communication-efficient federated learning with sketching. In H. D. III and A. Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 8253–8265. PMLR, 13–18 Jul 2020

We are thankful to the reviewer for the suggestions and comments. Below we provide answers to each of the concerns:

Q1. Robustness to violation of assumption on eigenvalues: Under the assumption on eigenvalues, i.e. $\kappa = o(1)$, we provide a dimension-independent convergence rate. However, our results hold for any $\kappa$ and thus even for a slower eigendecay such that $\kappa = \mathcal{O}(\sqrt{p})$, our analysis using second order structure shows that $b=\mathcal{O}(\sqrt{p})$ should suffice to achieve faster convergence. To see why this is true, recall from Theorem D.1 in our paper, the sketching dimension $b = \tilde{O}(1/\varepsilon^2)$, thus setting $b=\mathcal{O}(\sqrt{p})$ gives $\varepsilon = \mathcal{\tilde{O}}(1/p^{1/4})$. As a result, leading terms involving $\kappa$ in the expression of learning rate $\eta = \dfrac{(1-\varepsilon)}{\left((1+\varepsilon)^2c_s\varrho^2 +\frac{c_H c_l}{\sqrt{m}}(1 + \kappa(2\varepsilon+\varepsilon^2))\right)}$ are $\frac{\kappa \varepsilon}{\sqrt{m}} = \mathcal{O}(L^{1/4})$ which is quite small compared to the ambient dimension $p=m^2L, \text{for width }m,L \text{ layers}$, leading to faster convergence rate. In contrast, the results from [Song et al. 2023] have learning rate $\eta = \mathcal{O}(b/p)$ which needs $b = \mathcal{O}(p)$ for dimension independent convergence. As a result, our results provide a better interpolation of convergence rate with the assumption violation than [Song et al. 2023] which only provides the blanket pessimistic upper bound of $\mathcal{O}(b/p)$. Furthermore, we present our results on the decay of eigenvalues across a variety of models/datasets in Fig 1 and Fig2 in the attached PDF (in the general response) showing robustness to architecture and dataset choices.

---

Q2. Robustness to violation of assumption on PL condition: Our dimension independent results depend on the Restricted Strong Smoothnes property(RSS) of overparametrized neural networks and are not due to assumption on PL condition. In Table1 , we provide the dimension independent results even for non-convex $\mathcal{L}$.

---

Table 1:Comparison of iteration complexities and assumptions for K=1 local-steps. $c_s,c_l$ are constants defined in Assumption 3.2 and $\varrho,c_H$ are defined in Lemma C.1,Lemma C.2 in our paper respectively. Note that $\varrho$ and $c_H$ are $\mathcal{O}(poly(L))$ and can be seen as independent of ambient dimension for wide networks. Under our assumptions and setup (Section 3), according to Lemma C.4 the loss function $\mathcal{L}$ can be shown to be $\beta-$smooth where $\beta = c_s\varrho^2 + \frac{c_lc_H}{\sqrt{m}}$. $\varepsilon$ is the sketching error and $\epsilon$ refers to the optimization error tolerance.

---

Q3. Performance on tasks beyond image classification and practical challenges:

While we focus on the theoretical analysis for sketching methods, past works [Rothchild et al. 2020],[Melas-Kyriazi et al. 2022] provide extensive empirical results comparing sketching based approaches across vision and language domain. For practical purposes, the large memory footprint of generating and storing large sketch tables and fast computation of sketches is a challenge. Also, due to the high variance of the estimator $R^{\top}R$ (Lemma B.1 in our paper), it is important to apply some sort of variance reduction which can be explored in future work. A practical design decision is trading off this overhead for reduced communication requirements and the downstream benefit of differential privacy implementation.

[Song et al. 2023] . Song, Y. Wang, Z. Yu, and L. Zhang. Sketching for first order method: Efficient algorithm for low- bandwidth channel and vulnerability. In International Conference on Machine Learning, pages 32365– 32417. PMLR, 2023.

[Rothchild et al. 2020] D. Rothchild, A. Panda, E. Ullah, N. Ivkin, I. Stoica, V. Braverman, J. Gonzalez, and R. Arora. FetchSGD: Communication-efficient federated learning with sketching. In H. D. III and A. Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 8253–8265. PMLR, 13–18 Jul 2020

[Melas-Kyriazi et al. 2022] L. Melas-Kyriazi and F. Wang. Intrinsic gradient compression for scalable and efficient federated learning. In B. Y. Lin, C. He, C. Xie, F. Mireshghallah, N. Mehrabi, T. Li, M. Soltanolkotabi, and X. Ren, editors, Proceedings of the First Workshop on Federated Learning for Natural Language Processing (FL4NLP 2022), pages 27–41, Dublin, Ireland, May 2022. Association for Computational Linguistics. doi: 10.18653/v1/2022.fl4nlp-1.4.