Sketching for Distributed Deep Learning: A Sharper Analysis
Authors: Mayank Shrivastava
Berivan Isik
Qiaobo Li
Sanmi Koyejo
Arindam Banerjee
Venue: NeurIPS 2024 poster (2024)
Keywords: sketching
distributed learning
federated learning
optimization
Primary area: optimization_for_deep_networks
TL;DR: We provide a tighter analysis of sketching in distributed learning that eliminates the dimension dependence without imposing unrealistic restrictive assumptions in the distributed learning setup.

Abstract: The high communication cost between the server and the clients is a significant bottleneck in scaling distributed learning for overparametrized deep models. One popular approach for reducing this communication overhead is randomized sketching. However, existing theoretical analyses for sketching-based distributed learning (sketch-DL) either incur a prohibitive dependence on the ambient dimension or need additional restrictive assumptions such as heavy-hitters. Nevertheless, despite existing pessimistic analyses, empirical evidence suggests that sketch-DL is competitive with its uncompressed counterpart, thus motivating a sharper analysis. In this work, we introduce a sharper ambient dimension-independent convergence analysis for sketch-DL using the second-order geometry specified by the loss Hessian. Our results imply ambient dimension-independent communication complexity for sketch-DL. We present empirical results both on the loss Hessian and overall accuracy of sketch-DL supporting our theoretical results. Taken together, our results provide theoretical justification for the observed empirical success of sketch-DL.

Sketching for Distributed Deep Learning: A Sharper Analysis
Mayank Shrivastava
Berivan Isik Google berivan@google..om
University of Illinois Trbban-Chhappige mayanks40111 saranks401llinois
Qiaobo Li
Sanmi Koyejo Stanford University sanmi0cs wtw.ccsstaaooo....ooooo edu
University of Illinois Srbanna-hampiia qiaobol28illinois.
Arindam Banerjee
University of Illinois Urbana-Champaign
arindamb0i illbbi@lllli wlliiiiseeee edu
Abstract
The high communication cost between the server and the clients is significant bottleneck in scaling distributed learning for overparametrized deep models. One popular approach for reducing this communication overhead randomized sketch- ing. However, existing theoretical analyses for sketching-base distributed learning (sketch-DL.) either incur a prohibitive dependence on the ambient dimension [62 need additional restrictive assumptions such heavy-hitters (59] Nevertheless,
despite existing pessimistic analyses, empirical evidence suggests that sketch-DI is competitive with its uncompressed counterpart, thus motivating sharper analysis. In this work, we introduce a sharper ambient dimension-ine inmeensioonindeeeddenl convergence analysis for sketch-DL using the second-order geometry specified by the loss Hessian. Our results imply ambient plexity for sketch-DL. We present empirical results both on the loss Hessian and overall accuracy of sketch-DL supporting our theoretical results. Taken together, our results provide theoretical justification for the observed empirical success of sketch-DI..
communication com-
Introduction
Distributed learning popular framework for training machine learning models, often deployed to support large-scale deployments and to support privacy, among other systems goals 20 58
To this end, distributed learning generally employs both server and client devices. Most standard implementations are set up in rounds: in each round, clients participate by performing (multiple) local model updates using their local data, then share the model updates with server. The server aggregates the collected local updates and broadcasts either the aggregate or the updated global model to the clients for the next round. As result the system setup, distributed learning generally requires high communication overhead due to the frequent communication between the server and the clients [42]. Federated learning, popular distributed learning setup, further extends the generic setting where clients' data distribution may be non-i.i.d. and the clients may change over time, thus suffering from the same costly communication. This communication overhead has motivated research compressing
Part of the work done while BI Stanford University.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

the model updates 335 43 through sparsification 44 73 7 30, 31], quantization 4 74 10,, low-
27]. Among these, linear sketching mechanisms have attracted significant attention due to their efficient and simple integration with existing distributed and federated learning frameworks. For instance, unlike many sparsification approaches 44 [30], sketching is an unbiased operation that does not require bias correction with error feedback mechanisms [63] especially as error feedback typically increases the memory cost and the complexity of integrating differential privacy [23]. Moreover, the linearity of the sketching operation makes it compatible with secure aggregation 12 D as opposed to the popular Top-
rank projection [7]
67 68
8 and sketching (59
62
50
33
sparsification 44 or quantization methods [4] that break the linearity in aggregation.
Despite the rising interest in sketching for distributed and federated learning, the existing bounds on the convergence error of sketching scale with the ambient dimension of the model [62] (i.c., in contrast to the lower sketching dimension) which may extremely large for modern overparameterized deep models [39 20 I1. This dimension dependence has been limitation of the sketching mechanisms at scale, making them less attractive for modern deep models. We identify and demonstrate the root cause of the dimension dependence of the standard optimization convergence analysis as associated with the assumption of isotropic smoothness of the loss, commonly made in such analyses [62 [9]. Several modifications have been proposed to get rid of the dimension dependence. For instance, Rothchild et al. [59] assume that the model updates at the clients have heavy hitters and, further, apply Top-r sparsification to break the dimension dependence. However, their heavy hitter assumption may not hold in general, and the Top-r operation introduces a bias that is, in turn, eliminated via error feedback. While their analysis requires these restrictive assumptions and modifications to eliminate the dimension dependence, they also note that they actually do not observe an adverse effectt of dimension dependence empirically even without these restrictions, indicating that the theoretical analysis of their sketching algorithm is not quite explaining the empirical success. Quoting from Appendix B.3 of Rothchild et al. [59], "However, this dimensionality dependence does not reflect our observation that the algorithm performs competitively with uncompressed SGD in practice, motivating our assumptions and analysis." Motivated by this mismatch between the existing analyses of sketching for distributed learning (which impose dimension dependence) [59] 62] and its empirical success in practice, we provide a substantially tighter analysis than the prior work and climinate the dimension dependence in sketching without imposing any unrealistic restrictive assumptions in the setup.
Our sharper analysis avoids the global isotropic smoothness assumption and utilizes the anisotropic restricted strong smoothness (RSS) property of overparameterized deep models 72 (2) [3] 52 [5],
and results in terms of the second order geometry of the loss. We present optimization results for single-local step as well as K-local step distributed learning, and present bounds on communication complexity based on the optimization results. We refer the reader Table [ for an overview of our results. We do not restrict our analysis to any specific sketching matrix, and our results hold for any symmetric nab-Gaussian sketching matrix. For instance, for computational benefits, one can use the popular Count- cketch [14] or Hadamard sketch (2) approaches, which are both examples of sub- nub-Gaussian sketching. Our contributions can be summarized as follows:
1. We identify and demonstrate the widely used global smoothness assumption of loss functions as the root cause of the dimension dependence of sketching methods.
2. We provide novel analysis for sketching in distributed and federated learning that eliminates the ambient dimension dependence in the convergence error.
3. We are the first to break this dimension dependence without making restrictive assumptions such as the heavy hitter assumption by Rothchild et al. [59]
4. We are again the first to do this without Top-r sparsification step in the framework which would have required additional measures to eliminate the bias.
5. We are the first to more precise second-order properties of the loss of deep models, e.g.. approximate restricted strong smoothness, eigenspectrum of Predictor Hessian, to analyze distributed and federated learning frameworks.
"We use Top-r instead of the more commonly used Top-k since we reserve the notation & to refer to the local steps clients.
2

---

Analysis
Iteration Complexity
Assumption
This work
8 + ke uym
log
PL. RSS
(Theorem 4.1
Song et al.
62
o
( + p&c
log
SC, f-smoothness
(Theorem F.
(a) K=1 local steps
Analysis
Iteration Complexity
Assumption
This work
o
((1e
(13e
Ke
log t
PL. RSS, Bounded Gradient
(Theorem 4.2)
Vmuute
Song al. 62
o p9K log
SC, f-smmothhness
(Theorem F.10 (a=0))
(b) K>1 local steps
Table I: Comparison of iteration complexities and assumptions with prior work. Our results scale with & which captures the intrinsic dimension (29) of the model rather than ambient dimension p. SC refers to strong-convexity and PL refers PL-condition. μ refers to the PL constant as well as the strong convexity parameter under SC assumption. According to Lemma C.2 the loss function L can be shown be 3-smooth and m refers the width of the neural network.
Related Work
Conmunication-Effiicee Distributed Learning. The high cost of communicating model updates between the clients and the server has motivated a recent interest in improving communication efficiency in distributed and federated learning. One common strategy called FedAvg (50) enables less frequent communication by letting the client perform multiple local iterations at every round. Another common approach is to compress the model updates before the communication reduce the cost of each round. These efforts can broadly be categorized into sparsification
3 73 44],
quantization (64
4 74), low-rank factorization (67 51 711, sketching (59
62 34 and sparse
47 32 While some of these compression methods are already unbiased [4), many are biased and have to be combined with other mechanisms reduce the bias for better convergence [44] [59]. Linearity is another desired feature, as it simplifies the techniques like secure aggregation (in the compressed dimension). Among the general compression approaches mentioned, sketching stands out as simple linear and unbiased operation, allowing for computations the reduced dimension before desketching. We do not propose a new compression method in this work but instead provide substantially improved convergence analysis for sketching- frameworks that breaks the dimension dependence-whicl explains why sketching would not explode the converge error in large models.
subnetwork training techniques 31 40
41
implementation of distributed learning with security-promotiin
Sketching. Over many years, sketching has been a fundamental tool for many applications, even before the surge of deep learning in 2010s [18 25 36) for low-rank approximation [65], graph sparsification [1]. and least squares regression 21 More recently, sketching has also found use in distributed and federated learning frameworks to reduce the dimension of the model updates for
62 27 The linearity of these sketching based frameworks has led to their successful integration with secure aggregation and differential privacy as well [I6]
communication efficiency [4 33
59
66 62] [7]. Despite the empirical success of these sketching based applications in distributed and federated learning, the existing upper bounds on convergence errors have dependence on the ambient dimension of the model-which limits their scalability to larger models. In this work, we provide a tighter convergence analysis and get rid of the dimension dependence, suggesting the promise of sketching at scale. The closest to our work is by Rothchild et al. [59] who also provide
3

---

dimensionn-nddepedeen independent convergence bound for their sketching algorithm, called FetchSCD (which a combination of Sount-Sketor projection. Top-r sparsification, and bias reduction with error feedback), under the assumption that the model updates have heavy-hitters. We note that both the heavy-hitter assumption and the Top-r step (and the error feedback coming along to minimize the bias) are necessary to get rid of the dimension dependence in their convergence analysis. We avoid both of these restrictions. Our key contributions over Fe are: (1) We do not make the heavy-hitter assumption since it does not hold in general. (2) We do not require Top-r sparsification to break the dimension dependence - this way, we have an unbiased sketching mechanism without the need for error feedback, which would increase the memory cost and make the integration with differential privacy mechanisms complicated.
dimension
Notation. For quantity we use the notation T2 to refer to the global variable at round i shared by all the clients and the server, and Te, to refer to the local variable for client e at round For positive integer n. [n] {1, .n... We use 12 the p X p identity matrix. We use H[]] for expectation. For a vector x, we use ||x|| or ||x||2 to denote its L2 norm. For matrix A, we use ||||| or || for the spectral norm of A. We use VoL(0) denotes its sub-Gaussian norm [70]. The notations (()),((()) (t) are the same the common (1)) (2(),, (() but they hide the dependence logarithmic terms. polylog(n) denotes Olog" (n:) for some kt.
ac
3 0=00-- For a random vector y, ||yl
3 Sketching for Distributed Learning
We consider distributed learning framework with C' clients, each client c 1...., C' having 1n local dataset De = {Z 3i of size ne and a local loss Le R' + R defined as (0) =
nhe Ei= E(Y,, Di.c). where E IP is the parameter vector, Ui.e := f(0; xx.c) is the prediction of the model for input Xs,c. and É x IR + E is a loss function that measures the error between Jt.c and the groundtruth Vi.c. Our goal is to minimize the empirical loss L(0):
1
C
[(0) = C Σ Co)
(1)
Unlike much of the existing literature [59] 33 62), our analysis does not ignore the fact that the predictor is deep learning model - i.c., the loss Le is not just an arbitrary smooth loss. The motivation behind this focus is that many of the modern models being used in distributed learning indeed deep learning models. Interestingly, the losses associated with the deep learning models have a certain second-order structure beyond basic smoothness which will be the key to our sharper analysis. We also demonstrate that just assuming that the loss is smooth, as is typically done in most of the existing literature [59 33 62]. does not avail one of the sharper analysis we introduce - making the dimension dependence unavoidable.
54 45 6), we consider a fully-connected feed-forward neural network f of depth I., with widths m and activations (!)) for each layer I E [7.] = (1,...,4. L.). described as:
Following standard literature
22
a (x) =x,
a (0)(x) = ob
w. - Well [4]
f(0:x) = a [x)+)((x)
1 v v. (x),
mmL
where W.C Xm, is the layer-wies wise weight matrix for layer E [L] and v. E R is the last layer vector at iteration t, (6-) is the smooth (pointwise) activation function, and my dim(x) = d.
We denote the total set of parameters as
0,
(vV)T ERP
For simplicity, we will assume that the width of all the layers is the same, i.e., my = m for all 1 E (L.,, and thus p = (I. - 1) 1)m² + md + m. We consider deep models with only one output, i.e., outputs.
f00:x) x) E IR, but our results can be extended to

---

we make the following assumptions regarding the activation function o, loss function f. and the initialization parameter fn which hold true for the commonly used activation functions, loss functions, and initializations used in practice.
As prevalent in literature [22 6
Assumption 3.1 (Activation function). The activation & is 1 1-Lisschitz and B Ba-smooth.i |0'| v 1
and [6d1
B.-
Assumption 3.2 (Initialization). We initialize 0o with Wo.if o(0, for E [L] where do 0, and Vo is random unit vector with ||vol12 = 1.
(1)
01
Assumption 3.3 (Loss function). The loss L Elys, 9c) with E ie., |6] I c,, and (ii) E and smooth e S Cs for some Ct, Cs > O.
de, dgr
P
St. is (i) Lipsehitz, dg:
3.1 Sketching-Based Distributed Learning
is a compression technique that uses random projections to reduce the dimensionality and helps speed up computations. These random linear mappings can be represented by sketching matrices R 12xp where typically b < p. Examples of sketching matrices include Count-Sket [14.. Subsampled Randomized Hadamard Transforms (SRHT) [13), and sparse John- son-Lindenstrauss (JL) transforms [19]. In this work, we use sketching matrices to compress local updates before sending them to the server and refer to the operation of recovering true gradient vectors from the sketched updates "'deskecchnn".
Random sketching 77 48
We outline the sketching- based distributed learning framework in Algorithm [ Each client receives a random seed from the server to initialize the local parameters 0c,1, and generate sketching matrix R. At each local step k E [1, [1,... K]. each client performs local gradient descent (GD) over their local dataset De- At each communication round, the client accumulates the changes over K- steps, sketches the local updates, and sends the sketched update to the server. The server then aggregates the sketched changes and sends the aggregated sketched updates back to the clients. To update the local parameters, each client needs to recover an unbiased estimate of the true vector from the aggregated sketched update. We call this the desk (desketch) operation (Line 9), for which we use the transpose of the sketching matrix R. Each client then desketches the received aggregated sketched updates by applying desk and updates their local parameters. We refer to the sketching and desketching operations using the sk and desk operators defined as:
sk == R E bxp
(Sketching),
(4)
desk = R T E ppd
(Desketching).
(5)
While we use the same sketching matrix across communication rounds "....,T. ,, in general, using different matrices for each round does not affect the analysis.
Choice of sketching matrix: We use (luVD)-sub-Gaussian matrix as the choice of sketching matrix. We say R E Rx is a (1/vb)-sub-Cr) Sub--aussiinn matrix [69] if each row R; is an independent mean- zero, sub-Gaussian isotropic random-vector such that || |Rll 1/Vb. We assume E[R R] 4ppp-
From the above definition, we can see that for g 92 E IRP
R(g1 + 92) Rgg + Rg2 (Linearity),
E [ Rgl = B-1
(Unbiasedness).
(6)
3.2 Limitations of the Existing Analyses
When analyzing the convergence rates of the sketching bbased based distributed learning frameworks, previous works 62 59] assume that the loss function L is 3-smooth, i.e.
3 [(0') < C()) + (VC(0), 0) + 0112.
In the model updates based on sketching followed by desketching, the term ||0'- 0 effectively yields term of the form || Rgll2, where g E IR" stands for suitable gradient on the full model. While B Rg] g, i.c., the desk-sk operation is unbiased, we unfortunately have || Rgll = @(
5

---

Algorithm Sketching- Baach Distributed Learning.
Hyperparameters: server learning rate Tiglobal, local learning rate Mlxal-
Inputs: local datasets 'D of size ne for clients c = 1, 1,..., C, number of communication rounds T. Output: final model By-
Broadcast random SEED the clients.
2: for = ,,. do
3: On Client Nodes:
C. 4: for c lorc= do
5: if then
6:
Receive the random SEED from the server. Initialize the local model Oe,1 IR" and generate the sketching matrix R E Rxxp (hence sk, desk) using the random SEED.
7:
else
8:
Receive sk( 11) from the server.
9:
Desketch and update the model parameters 02 f- 0:-1 + (skk(sk(( (A, 1)).
10:
Assign the local model's parameters De + Og to be updated locally. end if
11:
12:
for k -...., K do
13:
0e,1 t- Be,t Necal VoLc(fe,1)
14:
end for
15:
Ac.t t + -Bae- 0,
16: Send sketched updates sk( (k)),,)) the server.
17:
end for
18:
19:
On the Server Node:
20: Receive sketched updates sk( (Ae,s) from clients c c.....,
21:
Aggregate: sk( ((( )) + Nglobal à E sk(A),t)
22:
Broadcast sk) to the clients.
23:
end for
B. Appendix). Thus, such an analysis picks up dimension dependence (p)) which can be neutralized only if the sketching dimension is b = s(p).. However, such high- gigh-dimensional projection will be ill-conceived as we will not get the benefits of the sketching projection. We note that prominent recent work has all hit this dimension dependence. For instance, Song et al. [62] have the dimension dependence all their results including communication complexity, and Rothchild et al. [59] discuss the dimension dependence in Appendix B.3 and get around the dependence by using Top-r components the gradient vector, with an analysis having to rely on heavy-hitter assumptions. Interestingly, Rothchild et al. (59) noted that the sketching- brrseidng-baeed distributed deep learning approach seemed to work fine empirically without getting the adverse effect of dimension dependence despite what their theoretical results (based smoothness) suggest. Our work sheds light on this discrepancy and proves (see Section why the sketching-bad based distributed learning approach in its simplest form (see Algorithm I does not pick up the dimension dependence.
with high probability (see Lemma
3.3 Restricted Strong Smoothness (RSS)
In this section, describe the RSS property, interesting property of deep learning losses that we use to derive dimension-inneepedeen convergence guarantees for sketch-DL. Using Taylor's expansion, the loss at = can be written as:
((0') L())) + ((' 0r. VoC(O.) 2 - V 02).
(8)
where 9 e (1 c),,cc c e [0. and ViC(0) the Hessian of the loss. Several prior works 77 61 60 have studied the loss Hessian by decomposing the Hessian into the Gauss-Newton matrix (G) and averaged Hessian of the predictors as (H):
ViC(O) = G H.H,
(9)
G - 0 1
Σ
(10)
1m1

---

Rc v2
nC Σ Th
i.
S2,00 fhe
Hipe
(11)
Th i=l
where Vic = Vef(0; x.) and - =4 v ( respectively. Recent works
45
6
22
have derived sharp upper bounds on the spectral norm of the predictor Hessian: 0(1/vm). This leads to the following restrcited smoothness property of deep
||
Amar V
learning losses:
1
1 Tce
nce L. fr. -0,)² Co A
C(8') L(0)) + (0' 0s, VoC(O.))
0 0.
2C
In contrast the 3-smoothness assumption which is common in optimization literature, this new property shows that deep learning losses exhibit strong smoothness in restricted set of directions. The effect of strong smoothness suitably manifests for such that cos(e' 8,, V fi.)) I 2 K > 0 for restricted set of directions. In other directions, the strong smoothness constant is O mm). i.e., tiny value. For our analysis, we use an even sharper perspective on RSS based on the eigenvalues of predictor Hessian, Hi,e- ( V fr.e- While the smoothness perspective based on the spectral norm of (LL///) is promising, it implicitly assumes all directions have this level smoothness which would impact the analysis since there are p O(mm²) directions. Based on empirical evidence [55 26), most directions have smoothness much smaller than 1/ 1/vmi, and an analysis based on the eigenvalues of the Hi,e captures this sharper perspective opposed to picking up dependence the ambient dimension p.
4 RSS-based analysis for Sketching-based Distributed Learning
In this section, we analyze the convergence of the sketching- based distributed learning approach, summarized in Algorithm T using the RSS property of deep learning losses and provide novel peeension-indepedeer convergence results. First, we state our key assumptions in Section [.1 and explain why they well-supported by recent work. Then, we provide eliminates the dimension dependence for the single-local step (K = 1) case in Section 4.2 Next in Section 4.3 we extend our analysis the multiple-local step (K > 1) case. We conclude in Section 4.4 by deriving the communication cost of Algorithm [under our novel convergence analysis.
novel analysis that
4.1 Assumptions
Before we introduce the assumptions, we would like to recall a key property.
Definition 1 (PL condition). Consider loss function L RP - IK and the solution set X* = {0' 8 E argmin L(6)) and we L denote the corresponding minimum value. Then, L is said to satisfy the Mojas--Lojasiewicz (PL) condition with constant if
1
VoL(O)]
L()) L
21
Remark 4.1. PL condition [57] can be used establish linear convergence of gradient descent while still being weaker than strong convexity. One can show that if function f is ju-strongly convex then also satisfies PL condition with the same constant p. Recent literature has shown that wide neural networks can be shown to satisfy some variant of PL condition [37 46).
Relying on this recent evidence, we make the following assumption our analysis: Assumption 4.1. Loss function L(-) satisfies the PL condition with constant μ..
We further assume the following upper bound on the sum of eigenvalues of the Hi,e follows: Assumption 4.2 (Predictor Hessian Eigenspectrumm) Let values of the predictor Hessian Hie,, = t/ 2 ((e; for E [T] and Amax = max |A,, then there exists
AA,,,, be the eigen-
O()) such that Σ/=1 Aicccsl S K.maas
7

---

Remark 4.2. Several works 24 49 76 have shown that the spectrum of loss Hessian follows bulk and outliers structure where the bulk can be attributed H and follows power law trend [Figure 3(b), Figure in Papyan [56]]. motivating our assumption. Further empirical evidence is presented in Appendix G (refer Figure 2) showing that for common networks, &: is much smaller than p.
4.2 Single-Local Step (K 1)
We now analyze the simpler setting where clients communicate their local updates after running a single local update step (K 1). In this case, there is no local drift, i.e., all clients share the same local parameter vector. As a result, we can show that local update (Line 9 in Algorithm [ can be written as 0g- 0,-1 1y desk (o((VoL(0,)))
We follow similar analysis Song al. [62]. but we exploit the second order structure of the deep learning losses, which helps avoid picking up dimension dependence due to the desk(sk(-)) operator. We state the main theorem below with the full proof in Appendix D.I
Theorem 4.1 (Informal version of Theorem D.I. Set K=1 in Algorithm Under Assumptions 4.1
and 4.2 for suitable constants E, 5 < 1, learning rate Nglobal Mocal independent p, and
b = Ω polylog P¹NTT) with probability at least 1 - ō, we have:
((Or) - [(0)) = (C((g) - (()))) e ((1-))T
(12)
where 8 is minimizer
(T]
In Theorem 4.1 (which also implies linear convergence), the sketching dimension, b, depends on p only polylogarithmically, which explains the competitive performance of the sketching based distributed learning frameworks with uncompressed GD without requiring any additional costly steps such Top-r sparsification (as done by Rothchild al. [59), which also introduces bias that needs to be suitably corrected.
4.3 Multiple-Local Step (K>)
Now, we provide the convergence result for Algorithm 1 for the more general case of multiple local steps (K > 1) at each round. Unlike the single-local step, here, the local clients' parameters drift during K local steps and thus, we need to have additional assumptions guarantee convergence. For purpose of our analysis, we assume that the gradient norms are bounded, i.e., [))(()) I refer the reader to Theorem D.2 in the Appendix for more formal statement. We now present an informal version of our theorem for the convergence of the multiple-local step case:
We
Theorem 4.2 (Informal version of Theorem D.2). Let [|V)(()) < G. Under Assumptions 4.1 and 4.2 for suitable constants e, er and e < I, learning rate r = Agloball < 2nk and
b = Ω polylog( TTN))).
)). with probability least 1 5, have:
[(OT) L()") < (C(00) - -)e -21 2(1-=))pnKT C22E, m, KG
(13)
24
where 8" is minimizer of the problem [ N is the number of training samples, and
C2(e, m. o
CH
+
SKCH
(14)
vm
We refer the reader to Appendix D.2 for the full proof.
4.4 Communication Efficiency
As direct consequence of our analysis, we now provide an improved communication complexity for Algorithm [ compared to prior works (62) that build their analysis solely based on the smoothness assumption without taking advantage of the RSS property of the deep learning losses. With this,
we manage to eliminate the linear dependence of the communication complexity on the ambient dimension p- providing substantial improvement over the analysis of Song et al. [2]. Unlike Rothchild et al. (59). we break this dependence without requiring a Top-r sparsification step or heavy-hitter assumption. the theorem below, we state the required number of communication bits to achieve e-optimal solution based on sharper convergence analysis in the previous sections.

---

Theorem 4.3 (Informal version of Theorem E.
Under Assumptions 4.1 and 4.2
Algorithm
obtains an optimal solution satisfying the error
[(Or) (()")
S
(15)
using Ö (c (ax Cale.m.x) (1-(1-
log ((06)-((00)
bits of communication.
We provide the full proof in the Appendix
E
We state the result by Song et al. 62 for comparison:
o
(
BC (
G2
/3
0
max
P.
log
(16)
JuC
Song al. [62] assumes that the loss function L is 3-smooth and the expectation is the randomness of sketching. Note that our bound depends on x instead the ambient dimension p. presenting sharper bound. This shows the efficacy of our approach in improving the overall communication complexity over prior works.
Experimental Results
In this section, we provide a comparison of the sketching approach in Algorithm with other common approaches such as local Top-r [4] and FetchSGD [59]. We note that both Local Top-r r (as outlined in Algorithm 2n Appendix F and FetchSGD are biased algorithms and they are typically used with error feedback mechanisms to correct the bias. As the sub-Gaussian sketching matrix Algorithm [ we use Count-Sketch This means the only difference between FetchSGD without error feedback and Count-Sketch is the extra global Top-r step the server. We conducted our experiments on NVIDIA Titan X GPUs
internal cluster server, using GPU per one run.
We train ResNet-18 18
23 on CIFAR-IO dataset 38 that i.i.d. distributed 100 clients. Each client performs local gradient descent iterations (i.e., using full-batch of size 500) at every round. Figure I shows that Count-Sken ch-bach distributed learning approach in Algorithm perrithmmprrrfrmms compet- itively with FetchSGD. This result highlights the potential of sketching alone, without additional modifications in FetchSGD, maintain competitive accuracy. Additionally, the error-feedback free approach enables compatibility with Differential Privacy(DP)) techniques which we leave as future work.
0.92
0.90
Top-r with error feedback
Accuracy
0.88
Local Top-r without error feedback
0.86
FetchSGD with error feedback
0.84
FetchSGD without error feedback
0.82
Count-Sketch
0.78
Dimension Reduction
Communication Efficiency. Count-Sketor algorithm in Algorithm I against local Top-r 44 and FetchSGD [59], with and without eтror feedback. 100 clients run local itera- tions with full-batch every round.
Figure 1:
We note that do not claim novelty of any of the methods discussed in this section they have been introduced [44 (34,, improved (59). and analyzed (62) extensively in prior work. We present these empirical comparisons for completeness and provide support for why tight analysis of sketching is important, given the mismatch between its empirical success and the dimension dependence in existing convergence analysis as highlighted in Appendix B.3 of Rothchild al. [59]
6 Discussion and Conclusion
We provide a significantly improved convergence analysis for sketching- distributed learning frameworks by exploiting the properties of the deep learning losses, such as restricted strong smooth-

---

ness. This allows us to break the dimension dependence in the convergence error, and consequently. the communication cost a milestone prior work could not achieve due to relying only on the smoothness assumption of the losses, i.e., ignoring the more special properties of "deep models." By breaking this dimension dependency the convergence analysis and communication cost, we hope to motivate the use of sketching for larger models. One of the many exciting future extensions is to revisit the privacy analysis of private sketching mechanisms using our findings.
We provide a sharper analysis for sketching in distributed learning. In future work, we plan extend our analysis federated learning by allowing client dropout.
Limitations.
Broader Impact. This paper provides a tighter analysis for sketching- distributed learning and federated learning frameworks. We expect this work to be helpful for the community it contributes to the efforts in making machine learning models more decentralized, accessible, and trustworthy.
Acknowledgements
The authors would like to thank the anonymous reviewers of NeurlPS for their valuable feedback and suggestions. AB acknowleges support in part by the National Science Foundation (NSF) through awards IIS 21 21-31335, OAC 21 -300853, DBI 20-21898, as well C3.ai research award. BI was supported in part by a Google PhD Fellowship. SK acknowledges support by NSF III 2046795, IIS 1909577, CCF 1934986, NIH (RO1MH11622--01 NIFA award 2020-67021-32799, the Alfred P. Sloan Foundation, and Google Inc.
References
[1] K. Ahn, S. Guha, and A. McGregor. Graph sketches: sparsification, spanners, and subgraphs. In Proceedings of the 31 ACM JIGMODSSGACCTSSGN symposium on Principles of Database Systems, pages 5 14, 2012.
[2] N. Ailon and B. Chazelle. Approximate nearest neighbors and the fast Sohnson--indenstraate transform. In Proceedings of the eiihty- annual ACM symposium on Theory of computing, pages 557-563, 2006.
[3] A. Aji and K. Heafield. Sparse communication for distributed gradient descent. In EMNLP 2017: Conference on Empirical Methods in Natural Language Processing. Association for Computational Linguistics (ACL), 2017.
[4] D. Alistarh, D. Grubie, Li, R. Tomioka, and M. Vojnovic. Qsgd: Communication-efficel sgd via gradient quantization and encoding. Advances in neural information processing systems, 30, 2017.
[5] A. Banerjee, S. Chen, F. Fazayeli, and V. Sivakumar. Estimation with norm regularization. Advances in neural information processing systems, 27, 2014.
[6] A. Banerjee, Cisneros- Velarde, L. Zhu, and M. Belkin. Restricted strong convexity of deep learning models with smooth activations. In The Eleventh International Conference on Learning Representations, 2023.
[7] I.. P. Barnes, H. A. Inan, B. Isik, and A. Özgür. rtop-k: A statistical estimation approach to distributed sgd. IEEE Journal on Selected Areas in Information Theory, 1(3):897-907, 2020.
[8] R. B. Basat, S. Vargaftik, A. Portnoy, G. Einziger, Y. Ben-Itzhake and M. Mitzenmacher. QUICK- FL: Quick unbiased compression for federated learning. arXiv preprini vrXiv:2205.13341 2022.
[9] J. H. Bell, K. A. Bonawitz, A. Gascón, T. Lepoint, and M. Raykova. Secure single-server aggregation with (poly) logarithmic overhead. In Proceedings of the 2020 ACM SIGSAC Conference on Computer and Communications Security, pages 1253-1269, 2020.
[10] J. Bernstein, Y.-.. Wang, K. Avizzadenesheli, and A. Anandkumar. signsgd: Compressed optimisation for non-convex problems. In International Conference on Machine Learning, pages 560-569, PMLR, 2018.
10

---

[11] R. Bommasani, D. A. Hudson, E. Adeli, R. Altman, S. Arora, S. von Arx, M. S. Bernstein, J. Bohg, A. Bosselut, E. Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprini arXiv:2108.072555 2021.
[12] K. Bonawitz, V. Ivanov, B. Kreuter, A. Marcedone, H. B. McMahan, S. Patel, D. Ramage, A. Segal, and K. Seth. Practical secure aggregation for federated learning on user-held data. arXiv preprint 1rXiv:1611.04440
2016.
[13] C. Boutsidis and Gittens. Improved matrix algorithms via the subsampled randomized hadamard transform. SIAM Journal on Matrix Analysis and Applications, 34(3): 1113)333001-00 101--1301, 2013. doi: 10. 137/1 20874540. URL URL httph:tpp//// htttpp/://oi..rrr/1.113/////////
[14] M. Charikar, K. Chen, and M. Parach-Colton. Finding frequent items in data streams. In International Colloquium on Automata, Languages, and Programming, pages 693-703. Springer, 2002.
[15] S. Chen and A. Banerjee. Structured estimation with atomic norms: General bounds and applications. Advances in Neural Information Processing Systems, 28, 2015.
[16) W.-N. Chen, A. C. Choo, P. Kairouz, and A. T. Suresh. The fundamental price of secure aggregation in differentially private federated learning. In International Conference on Machine Learning, pages 3056-3089 PMLR, 2022.
[17] W.-N. Chen, B. Isik, P. Kairouz, A. No. S. Oh, and Z. Xu. Improved communication- privaacy trade-offs in $1 2S mean estimation under streaming differential privacy. In Forty-first Inter- national Conference on Machine Learning, 2024. URL : //tp///eepe/eeeee//////e/ee wttp:evvwwwwwweefreeeeeeee net/forum? id=x1C7ieRgRd
[18] G. Cormode and S. Muthukrishnan. An improved data stream summary: the count-min sketch and its applications. Journal of Algorithms, 55(1):58-74 75, 2005.
[19] A. Dasgupta, R. Kumar, and Sarlos. A sparse johnson: Lindenstrauss transform. In Proceed- ings of the Forty-Second ACM Symposium on Theory Computing, STOC '0, page 341-350, New York, NY, USA, 2010. Association for Computing Machinery. ISBN 978 4503005506 doi: 10.1 18.145/1806689. 1806337. URL wttph:tt:///ttt //doi. org/10...... 1011/11110000 1806689. 1806737
[20] Dean, G. Corrado, R. Monga, K. Chen, M. Devin, M. Mao, M. Ranzato, A. Senior, Tucker, K. Yang, et al. Large scale distributed deep networks. Advances in neural information processing systems, 25, 2012.
[21] E. Dobriban and S. Liu. Asymptotics for sketching in least squares regression. Advances in Neural Information Processing Systems, 32, 2019.
[22] S. Du, J. Lee, H. L. Wang, and X. Zhai. Gradient descent finds global minima of deep neural networks. International conference on machine learning, pages 1675-1685. PMLR, 2019.
[23] C. Dwork, A. Roth, et al. The algorithmic foundations of differential privacy. Foundations and Trends@ in Theoretical Computer Science, 9(3-4):211-4077 2014.
[24] B. Ghorbani, S. Krishnan, and Xiao. An investigation into neural net optimization via hessian eigenvalue density. In International Conference on Machine Learning, pages 2232-2241. PMIR, 2019.
[25] M. Greenwald and S. Khanna. Space-efficient online computation of quantile summaries. ACM SIGMOD Record, 30(2):58-66 66, 2001.
[26] G. Gur-Ari, D. A. Roberts, and E. Dyer. Gradient descent happens in tiny subspace, 2018.
[27] Haddadpour, B. Karimi, P. Li, and X. Li. Fedsketch: Communication-effficinn and private federated learning via sketching. arXiv preprint arXiv:2008. 04975, 2020.
[28] K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770- 778, 2016.
11

---

[29] I. C. F. Ipsen and A. K. Saibaba. Stable rank and intrinsic dimension of real and complex matrices, 2024. https:////xxxiioorrr wtw.:/rg//bbb////eee abs/2407 21594
[30] B. Isik, T. Weissman, and A. No. An information-theorrticc justification for model pruning. In International Conference on Artificial Intelligence and Statistics, pages 3821-3846. PMLR, 2022.
[31] B. Isik, F. Pase, D. Gunduz, T. Weissman, and Z. Michele. Sparse random networks for communication--eficcennn federated learning. In The Eleventh International Conference on Learn- ing Representations, 2023. URL /ecttp:://pperevviie net formm? wwt.ormmimd.ddakri1r.r....
[32] B. Isik, F. Pase, D. Gunduz, S. Koyejo, T. Weissman, and M. Zorzi. Adaptive compression in federated learning via side information. In International Conference on Artificial Intelligence and Statistics, pages 487-495. PMLR, 2024.
[33] N. Ivkin, D. Rothchild, E. Ullah, I. Stoica, R. Arora, et al. Communication-effin distributed sgd with sketching. Advances in Neural Information Processing Systems, 32, 2019.
[34] J. Jiang, F. Fu, T. Yang, and B. Cui. Sketchml: Accelerating distributed machine learning with data sketches. In Proceedings of the 2018 International Conference on Management of Data, pages 1269-1284, 2018.
[35] P. Kairouz, H. B. McMahan, B. Avent, A. Bellet, M. Bennis, A. N. Bhagoji, K. Bonawitz, Z. Charles, G. Cormode. R. Cummings, et al. Advances and open problems federated learning. Foundations and Trends@ in Machine Learning, 14(1-2)11-20,, 2021.
[36] D. M. Kane and J. Nelson. Sparser johnson-lindenstraus transforms. Journal of the ACM (JACM), (16:1:11122 2014.
[37] H. Karimi, J. Nutini, and M. Schmidt. Linear convergence of gradient and proximal-gradies methods under the polyak-Lojasiewiez condition, 2020.
(38] A. Krizhevsky, G. Hinton, et al. Learning multiple layers of features from tiny images.
2009.
[39] Y. LeCun, Y. Bengio, and G. Hinton. Deep learning, nature, 521(7553):436 444, 2015.
[40] A. Li, J. Sun, B. Wang, L. Duan, S. Li, Y. Chen, and H. Li. Lotteryfl: Personalized and communication--ffiiiee federated learning with lottery ticket hypothesis on non-iid datasets. arXiv preprinl
[41] A. Li, J. Sun, X. Zeng, M. Zhang, H. Li, and Y. Chen. Fedmask: Joint computation and communication--ffiiiee personalized federated learning via heterogeneous masking. In Pro- ceedings of the 19th ACM Conference on Embedded Networked Sensor Systems, pages 42-55, 2021.
[42] M. Li, D. G. Andersen, J. W. Park, A. J. Smola, A. Ahmed, V. Josifovski, J. Long, E. J. Shekita, and B. Su. Scaling distributed machine learning with the parameter server. In / 1th USENIX Symposium on operating systems design and implementation (OSDI 14), pages 583-598, 2014.
[43] M. Li, D. G. Andersen, A. J. Smola, and K. Yu. Communication efficient distributed machine learning with the parameter server. Advances in Neural Information Processing Systems, 27, 2014.
[44] Y. Lin, S. Han, H. Mao, Y. Wang, and B. Dally. Deep gradient compression: Reducing the communication bandwidth for distributed training. In International Conference on Learning Representations, 2018.
[45] C. Liu, Zhu, and M. Belkin. On the linearity of large non-linen oin-lrnear models: when and why the tangent kernel is constant. In Advances in Neural Information Processing Systems, 2020.
[46] C. Liu, L. Zhu, and M. Belkin, 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.
12

---

[47] Y. Liu, Y. Zhao, G. Zhou, and K. Xu. Fedprune: Personalized and communication- -eafficaatto-effiiinntt fed- erated learning on non-iid data. In International Conference on Neural Information Processing. pages 430-437. Springer, 2021.
[48] M. W. Mahoney. Randomized algorithms for matrices and data. Foundations and Trends@ in Machine Learning, 3(2); 1232 2011. ISSN 1935-8237. doi: 10.1561/22000000355 URL attp://áz: doi. org/ 10. 15611220000000055555
[49] M. W. Mahoney et al. Randomized algorithms for matrices and data. Foundations and Trends in Machine Learning, 323):223- 2011.
[50] H. B. McMahan, F. Yu, P. Richtarik, A. Suresh, and D. Bacon. Federated learning: Strategies for improving communication efficiency. In Proceedings of the 29th Conference on Neural Information Processing Systems (NIPS), Barcelona, Spain, pages 5-10, 2016.
[51] A. Mohtashami, M. Jaggi, and S. Stich. Masked training of neural networks with partial gradients. In International Conference on Artificial Intelligence and Statistics, pages 5876- 5890. PMLR, 2022.
[52] S. Negahban and M. J. Wainwright. Restricted strong convexity and weighted matrix completion: Optimal bounds with noise. The Journal of Machine Learning Research, 1361):1665-1697, 2012.
[53] S. Negahban, B. Yu, M. J. Wainwright, and Ravikumar. A unified framework for high- dimensional analysis of m-estimators with decomposable regularizers. Advances in neural information processing systems, 22, 2009.
[54] Q. Nguyen. On the proof of global convergence of gradient descent for deep relu networks with linear widths, 2021.
[55] V. Papyan. Measurements of three-level hierarchical structure in the outliers in the spectrum of deepnet hessians, 2019.
[56] V. Papyan. The full spectrum of deepnet hessians at scale: Dynamics with sgd training and sample size, 2019.
[57] B. Polyak. Gradient methods for the minimisation of functionals. USSR Computational Mathematics and Mathematical Physics, 3(4):864 87888644888,,, 1963. ISSN 0041-5553. doi: https://doi. org/10. (1111000006600011555333333366300 atricleeppi//00415553333000
UtLhhtts:://wwx sciencedirect.. con/science)
[58] D. Povey, X. Zhang, and S. Khudanpur. Parallel training of deep neural networks with natural gradient and parameter averaging. In ICLR (Workshop), 2015.
[59] D. Rothchild, A. Panda, E. Ullah, N. Ivkin, I. Stoica, V. Braverman, J. Gonzalez, and R. Arora. Fetchsgd: Communication-efficle federated learning with sketching. In International Confer- ence on Machine Learning, pages 2265. PMLR, 2020.
[60] L. Sugun, L. Bottou, and Y. LeCun. Eigenvalues of the hessian in deep learning: Singularity and beyond. arXiv preprint arXiv:1611.07476 2016.
[61] L. Sagun, U. Evei, V. U. Guney, Y. Dauphin, and L. Bottou. Empirical analysis of the hessian of over-parametrized neural networks, 2018.
[62] Z. 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, PMLR, 2023.
pages
[63] S. U. Stich, J.-B. Cordonnier, and M. Jaggi. Sparsified sgd with memory. Advances in Neural Information Processing Systems, 31, 2018.
[64] A. T. Suresh, X. Y. Felix, S. Kumar, and H. B. McMahan. Distributed mean estimation with limited communication. In International Conference on Machine Learning, pages 3329-3337. PMLR, 2017.
13

---

[65] J. A. Tropp. A. Yurtsever, M. Udell, and V. Cevher. Practical sketching algorithms for low-rank matrix approximation. SIAM Journal on Matrix Analysis and Applications, 38(4):1454-1444 1485, 2017.
[66] E. Ullah, C. A. Choquette-Choo. P. Kairouz, and S. Oh. Private federated learning with autotuned compression. In International Conference on Machine Learning, pages 34668-34700 PMLR, 2023.
[67] S. Vargaftik, R. Ben-Basat, A. Portnoy, G. Mendelson, Y. Ben-Itzhak, and M. Mitzenmacher. Drive: One-bit distributed mean estimation. Advances in Neural Information Processing Systems, 34:362-377, 2021.
[68] S. Vargaftik, B. Basat, Portnoy, G. Mendelson, Y. B. Itzhak, and M. Mitzenmacher. Eden: Communication- effiicient and robust distributed mean estimation for federated learning. In International Conference on Machine Learning. pages 21984 22014. PMLR, 2022.
[69] R. Vershynin. Introduction to the non-asymptotic analysis of random matrices, page 210-268.
Cambridge University Press, 2012.
[70] R. Vershynin. High--Diensional Probability: An Introduction with Applications in Data Science. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 2018.
[71] T. Vogels, S. P. Karimireddy, and M. Jaggi. Powersgd: Practical low-rank gradient compression for distributed optimization. Advances in Neural Information Processing Systems, 32, 2019.
[72] M. J. Wainwright. High-dimensional statistics: A non-asymptotic viewpoint, volume 48.
Cambridge university press, 2019.
[73] H. Wang, S. Sievert, S. Liu, Z. Charles, D. Papailiopoulos, and S. Wright. Atomo: Communication-fficieen learning via atomic sparsification. Advances in neural information processing systems, 31, 2018.
[74] W. Wen, C. Xu, F. Yan, C. Wu, Y. Wang, Y. Chen, and H. Li. Terngrad: Ternary gradients to reduce communication in distributed deep learning. Advances in neural information processing systems, 30, 2017.
[75] D. P. Woodruff. Computational advertising: Techniques for targeting relevant ads. Foundations and Trends Theoretical Computer Science, 10(1-2):1-157, 2014. ISSN 1551-3068. doi: 10. /0.156//0400000000, URL http://dd doi. org 10. 1561/0500000000
[76] 7. Xie, Q.-Y. Tang, Y. Cai, M. Sun, and P. Li. On the power-law hessian spectrums in deep learning, 2022.
[77] C. Xing, D. Arpit, C. Tsirigotis, and Y. Bengio. walk with sgd. arXiv preprint arXiv:1802.08770, 2018.
[78] X. Zhang, Z. Bu, S. Wu, and M. Hong. Differentially private sgd without clipping bias: An enror-feedback approach, 2024.
14

---

A Sketching Guarantee
Lemma A.I. For randomised matrix R E Rbxp with i.idd subgaussian rows Ra i.e. [illles &
1/Vb and R = -PoxP and for U, V _ R,, such that b = 22( log
(+lU11V1/8) with probability
6²
atleast 8, we have:
sup wEU
(Ru, Ru) (u, u) S e 2
(17)
Proof. We use the following lemma from Song et al. 66
Lemma A.2 (Lemma D.24 from 66
Let R Roxp denote a random Gaussian matrix. Then for any fixed vector h E R² and any fixed vector R2, the following property holds:
log 1.5 (p/0) ||||a|1اl|| (
Pr R-ll
[6
R Rh) - (g (T)h))
vb
Following a similar proof Song et al. (62. we can get sub-Gaussian version of this lemma:
Lemma A.3. Let R E Rx be a random sub- Gaussian matrix, with 4/2 norm of each entry bounded
by I. Then we have:
Vb
<vlog (n/5) vb
Pr max (Ra,i, Ro.i)1 (*)
<
Proof. Note for + j. Ro Ra.j are two independent sub-Gaussian vectors. Let zд RR and (Ru,is Ra Then according to the definition of sub-Gaussian random variables,
SE is sub-exponential random variable with absolute constant e₁. Thus, we have
2k=1 2h E SE c at) by sub-exponential concentration Lemma B.7 in Song al. 62
have:
Pr |||| 2 t 2exp(-cşb (t),t).
Picking
03 /lo /8) /b, we have:
(Rasis Rajil evlog (p/8)
I 5/p²
vb
Taking the union bound over all (i,j)= [p] x [p) and # j. complete the proof.
Then following the same proof with the only difference to apply Lemma A.3 instead of Lemma D.18 Song et (62), we can get the following sub-Gaussian version of Lemma D.24 Song et [62
Lemma A.4. Let R E bixy denote a random Sabb-Gassian matrix, with the norm of each row bounded by Jb Then for any fixed vector E TRP and any fixed vector g R2, the following properties hold:
1.5 clog2.5
Pr R-1l
[1g
R T
Rh) - (g Tg)
Molaalbi1a
s O(6).
Vb
A.4 taking union bound over U, v V and setting 22 C Clog"('|U|]]//
Based on Lemma get the claim.
15

---

B Dimension dependence
We make use of the following lemma from Banerjee et al. [5] to show dimensional dependence in || Rg
Lemma B.1 (Theorem 5 in Banerjee et al. (5). Let A C SP and let X E RR be design matrix with independent isotropic sub- aub-Gaussian rows, i.e., Xillvz S R. and || = pxp Then, for
absolute constants 77, c O, with probability least 2 exp(-1rpu (A)), we have:
w(A)
sup wEA
|xu| 2 = sup uEA
Σ (x,, u)² -1 1=1
1 c
(18)
vn
where w(A) denotes the Gaussian width of any set A given by w(A) Bg (sup (a, g) where an isotropic Gaussian random vector, i.e., g ~ -N (N00, Ipxp).
1 Remark B.1. By choosing X to be a /n sub-Gaussian matrix, we get the form
Xu|| - 1 w(A) E
sup
vn
which we will use next.
Lemma B.2 (Dimension Dependence in aau-Gausiian sketching). Fix a vector g E Rr. For sketching matrix R defined earlier, we have for some constants Th 12 with probability at least 1 2 2exp(-m1e25) 2 exp(-n12")
| Rgll (1 (1- (1 1191..
(19)
Proof. Let G C IRP Gaussian width of G, w(A) Ol log (9)) For E G using Lemma B.l
w(AAA
Pr
sup LgEg
|Rll3 - |||||
2 ci
|| < 2 eexp(-mu2(0)
(20)
db
Similarly, we can define B = {h h = Rg Vg 9). Gaussian width w(B) = O|] /log (
Now setting the design matrix X for Lemma B.1 be /b/p RT:
b - || R
w (
Pr sup OEG
Rgll
||]]
2 C2
Rig 2
< 2 -eeeppppppe 2exp(-n(²(BB)
(21)
p
Taking e max { ci #(1)
02 wB
vb
Pr
sup GE
[Rll3
2 <||g113
(22)
b R Rgll3
and, Pr
sup LEEG
||gl|
S 2eepp -2eppp-r1e> pp)
(23)
Succintly, we can write with probability 2exp(-m)
| Rg 2 (1 - c)2 2 ||ا12
(24)
C Background
We state some useful results from Banerjee et al.
66
16

---

Lemma C.1 (Predictor Hessian Spectral Norm Bound, Theorem 4.1 and Lemma 4.1 in Banerjee
al. 6 ) Under Assumptions 3.1 and 3.2 for any x E X, 0 E BSec (6g), with probability least P.Pl
2((+1) , have
[a](0x))]1 1
en
and
(25)
m
where,
CH L + Ly + 1). (1+ D
max [E[L]
7
LT
max h()) (E(2)
1-1
7 + mm
h(1) = - + +6(0))
Σ
VH
max 164
Bah(li)² h(11)
Bo 80((72 + +h(12)9)) +1(1)))+))
BBn (o07 B(1hh()))))))
2
-M(++ 1))3 (1 p1)² (1(1)3724-11
1
m
Lemma C.2 (Local Smoothness. Theorem 5.2 Banerjee et with probability at least (1 2( V0, 8' E Frobb (6o).
Under Assumptions 3.2 and 3.3
3 = 0 E(0') < (()) + 10' 8, VoL(O)) +
CHCl with B cs
(26)
vm
D Analysis of Convergence
We will bound T1, T2 and T3, which we introduce below:
1 C
((0)
Σ
Σ
1
(yr.o f(0;
(27)
C
o=1
o=1
i=1
1
1
1
nee
Vol(O)
(oL.
l. (0: xie)
(28)
C
C
The
c=1
1 ViL(()
Σ
4. Σ i=1
Le (0:
(29)
C
The
Hs,
where 6, E the first and second order derivatives of the point-wise losses. For an iterate 0.-1, by the second-order Taylor expansion of the empirical loss around 0,-1 we have (ignoring h.o.t.).
=(0.)= = C(0.1) - - 0-11) + 20: - 1.-1) T 0.-1) (30)
Using
((0.) = C(0,-1) (V 0
2C
The
T3
c
1=1
1
O0
0.-1) 1) Hi. -1em1) 02-1)
(31)
2C
The
cm
iml
Ta
where, Hic (1.((0,-1) for brevity.
17

---

D.1 Single-step scheme (K=1)
Theorem D.1. Under assumptions
and
set K = 1 in Algorithm
1 For a suitable constant
211 (1-E)
1) Tioca ((1 c,,0 exci/mm(1 + x(22 =?)))
and 52 ( logo(y7) (yT/6))
with probability at least - 8, we have:
[(Or) - _() =((0)) - [()*)) e -m(2-e)mf
(32)
where 0 is the minimizer of the problem
|
Proof.
0, 1--0 0 Nglobal desk
1
Σ sk (Mocal (oCc(((-1)) =nddesksk VoL(0,-1)
(33)
C
where 7) Voocal for notation simplicity.
D.1.1 Bounding Ti
(0c - 02-1, VoL(Oe-1)) = (deskksk VoC(0)))))))
(34)
(35)
S -- + MEΙ Vo VoL
(36)
where, follows from Lemma A.1 for some for suitable choice of 61-
D.1.2 Bounding T3
Σ 10 1,- Vef(0r-1;xx.)))
(37)
<n
1=1
(38)
$ 1/ Ne
(39)
1 (E² 222) VoL(O,-1)
(40)
< <1 Σ V. + 4d (6) + 282) V L(C,,)))
(41)
1
(42)
Thus,
n Ee (V.)(0--iix)))))---.
22 (1+ 82) 2 c, [v[(0,-))]1 (43)
<
2C
The
for some for suitable choice of E2
18

---

D.1.3 Bounding T3
Using eigen-decomposition of Hi,e.
01-1) -0-1) - A0 V3,,o)
(44)
where Asse V1L. for j E [],-.. .ll are the eigen-values and eigen-vectors of Hs,e respectively. Now,
(0, 0,-, Vsste) -n/desksk voll
(45)
S L- v...c) + nes V
(46)
P
Σ
Aseeel0 - 00-11 vonce)
(47)
Σ
Asria Vo) 2e3 (20 2,00 (3ac) +
Vo 1)12
1=1
(48)
$
Asiie (VoL((l-1), + (63 + 22x)) V
(49)
A Vote)Voo(()-))) +
Σ
Agsiel
(eš + (
(50)
j=1
< Amxxi
VoL(O,-1)) viice)2
Ajricl
(E3 + V [o
(51)
1=1
QCCH
Ve[(O-11)1+
KACH
(E3 + 22a) [o[(0==1)1
(52)
mm
where, Amax,i,c
||H 16 V2 fl2
scgu vm
Thus,
nee 00-
ncnc
Hr H.,.)
01-1) 5
(1 + k<2E3 + +
20
m
(53)
for some for suitable choice of E3-
D.1.4 Combining Tı, T2, T3 Combining Ti, T2, T3. we get:
C(0.) < ((0,-)) - - vll + nEn VoL(6--1)1
(54)
1/2 + 2
(1 =z)²² 2 1 + 72CH4
(1 + x + + +
(55)
2/m
72 2
n²CHrC
((((-))
17 7=1
(1 + 52) co02
(1 + <(28> + +=31)
1 (
2/m
(56)
19

---

Using PL condition,
[(0,) S ((0=-1)
(57)
"1² (1 + 22) cse²²
11 1 1181
(1 x(233 o
24
(c(00_1))) L())))
2/m
(58)
C(0,)- < -1- 8) (c(0)-1)- -
(59)
where,
B = 2pm 1 (1-
!) ((1 E2) c.e²
CHQI vm (1+ (1+*(2es + =3))
(60)
2(1-E1) E1)
Setting 7 < (1+ @² + cM (1 + k: k((Eg + cz))
gives contraction map. we iterate this
recursion, we get:
[(6:) _(0") s (1 B) (C(0)) L()"))
(61)
Set
2(1
E1 82 =
6, 17 ((1 + cc + cHCs/ /m(1 (1 + k2le + +2)))
(62)
we get
[(0r) -(0') (1 uu1 - =)m)²² (C(06) 2(0))
(63)
(C()0) (((*)) e -- (()-5)97
(64)
D.1.5 Choice sketching dimension
Using A.1 for each step t [0, -0....T), T), union bounding over Ti, T2 and Ty and over all T time-steps,
Union bounding over all T time-steps, we get
b Ω { (
(65)
where EC ne the number of training samples.
D.2 Multi-step scheme
Theorem D.2. Let II (o((()) < G. Under Assumption 4.1 for suitable constant < 1j
Oglobal Tocal < 20k(1-5) and [ (p2 (T) with probability least we
have:
-(0r)-((0)) s (E(0n) - -(0"))e- 2)1- n022e, m, KG
(66)
241(1 -6)
where 8" is the minimizer of the problem
And,
Ca(e,m, - [
(coe""
CHCI vm
(1 (++)) + cc (1 + +0)² 2 0²
1
QCH 2/m
+ 2
1
(2e + e2)
QRCH
Vm
(67)
where, and eи defined in Lemma C.I
20

---

Proof. We assume that the gradients at each time step t are bounded i.e ( S G for some suitable constant G.
First, we introduce some notation: In the K-step case since the parameters are shared only in the syne step, define lettt E TRP as the local parameters for the client time step and step k E 21 ... K.. Based on the sync step in Alg: [ we can see that e,,,0 is the same for all clients E C and thus, denote it as Oc, i.c. 0c,2,0 0.VCE C. Thus, we can write the update in the sync step as:
C
K-1
02 00-1 =
-lllobal desk
0
Σ
sk
Tlocal
VeLo(Oor-LAk)
(68)
==1
k=0
1
-laaaabal Mocal
Igloobal
C
desk sk VoLc
(69)
K-1
17
Σ
desk sk VoLo
(70)
C
where,
7 Tglobal7/ocal-
Similar to Section D.I
we can bound T1, T2 and T3 in 31 as:
D.2.1 Bounding Tj
(VoL(0L-1)., -
(71)
K-1
11 (VoL(0)-1), 0
Σ
Σ
desk sk V (P
(72)
C
K-1
= --
-
(RV
(73)
a=l K=D
= =-nK(RV (l(0 1) RV (oL(O--1))
(74)
RV ((
[(:oo_(0 - RV [V
(75)
0
$ -
- Vo + Mocal
(c.0²+
CHO Vm +
(1 (1+
(76)
where,
[Ve(((-111cc)) VeC(0--10,0)]
CHQ
110-1,, 1,-1,
(77)
vm
k-1
co2² +
CHQI
1
Mocal V VoL(
(78)
Vm
j=0
CHCI
< 1Mocal
coe² +
KG
(79)
mm
D.2.2 Bounding T2
C K-1
(0) 0, - 02-12=
Vof(O--1;)), -
11
2
Σ
desk sk VoL,
(80)
1/2
C K-
2
= C
Σ
(p desk V VoL,
(81)
a=0
k=0
< +2) "2 K2
(82)
21

---

for some suitable choice of E2 and b which we set later. Thus,
1 Tee
-1 (v - -000))00)))000100)))-)----)))) - 12 co (1+ G²
(83)
20
2
c=1
t=1
D.2.3 Bounding T3
Using eigen-decoopposition,
(0: - V - -0-1)- = Σ
Afaac(r - 10-1,,5,ac
(84)
j=1
where Ajeess Viã for -., .pl are the eigen-values and eigen-vectors of life e ((:(((-1) respectively.
C
K-1
(0, - 0 Vy,.c)
12(
desk VoLc(Oc(-11)), Vys,c)2
(85)
C K
C K (oo(( 1
S k=
1
0
Σ
Vol L (0.
1,,x
k=0
(86)
C K
26 C
VoLc(Oee-IA)11
(87)
K=0
C K-1
V Fo (0 Vise
+ + K2e)
(88)
k=U
for suitable choice of E4 and b.
A (0c -1-1,, -0-11, vascd)
(89)
2C
Tee
The
C K-1
V VoLcOootLLk) Vohee )
2
Asai
+ (67 +6
20
The i=1
j=1
k=
(90)
2 QCH K²G² 2
1 (2e4 + 2
QRCH K² Vm.
(91)
V
2
D.2.4 Combining Ti, T2,T,TT:
Combining Tı, T2 and T3, we get.
[(0.) < ((0-)))))) -K(1)))))---)) +1 Toocl
ccee² +
CHCi vm
(1 (1+62)K²02
1/2
2 Ca (1 13) k2 + 1 quc K² + 112 [2e4 + +
QREH Vm
K²G²
(92)
mm
< L(f--1) - -nK(1 Vo
+ 17 Mocal c= +
CHQI
(1 (1+ + nc co (1 + Ex)2 0² +
QCH 1) vm
(2ea + +
QRECH
K²G²
vm
2
vm
(93)
22

---

Using PL condition and iterating over this recursion we have:
[((r) ((0)))((0)) < (1 -2mn =((0)) + C 1- (1 21 E1))²
(1 2107 - E1))
(94)
- - m,
(95)
2/1
where 1) I 2uR(1-c)- For simplicity, let Nglobal
i.e. 1j Mocal and E1 E2
E3 =E And,
C,(,, = -
1
[(
CH9 vm
1 2 (1 +e)²² e² (1 (1+
qCH
1
(28 + e²) qRCR
2/m
2
m
(96)
= K2(
CH
ERCH
(97)
Vm
Vm
D.2.5
Choice of sketching dimension b:
Using A.I for each step & [0, ......T), -T) union bounding over Ti, T2 and T3. Union bounding all T time-steps, we get
(
b = Ω E2
( NTTCCI/)))
(98)
where N = Ee
21=1 Dh is the number of training samples.
E Communication Efficiency
Theorem E.1. If Theorem 4.2 holds, then with Ô (C max
1, Ca(r.m,s)
2p-(1-eJKe log [((6))-(('))
bits of communication, with w.p at least 1 5 Algorithm: satisfying:
outputs an e-optimal solution Or
[(Or) ((0)) Sc
(99)
Proof.
Setting 7) min( [ux(1-2)) C)( and T
[1na[1-e] log ((C(())-<(0'))
have:
(CT) 2((-)) S
(100)
nC21E, m, KG
S
(101)
2/
(C(0.) - ( < (()) -(())))))(()))) ) C/22E, m, KG2
S (102)
2/u
The total number of communication bits are given
CFT = - =( {1, C2(e, m,c, 21 z)e
log (((b) -((8))
(103)
where, Ô hides poly-log dependence.
F Additional Experimental Details & Results
We provide the overview of the local Top-r baseline in Algorithm
23

---

Algorithm 2 Local Top-r-based Distributed Learning.
Hyperparameters: server learning rate Tiglobal, local learning rate Moxal-
Inputs: local datasets 'De of size ne for clients c = 1,..., C, number of communication rounds T. Output: final model By-
Broadcast random SEED the clients.
2: for = .,..., do
3: On Client Nodes:
4: for -e. C do
5: if then
6:
Receive the random SEED from the server. Initialize the local model 0c,1 E RP using the random SEED.
7:
else
Receive Á.- from the server.
8:
9:
Update the model parameters + 00-1 Š.-1- Assign the local model's parameters De,t +- O2 to be updated locally.
10:
end if
12:
for k = 1....,KK K do
13:
0ctt f- 0c, Mecal VoL.(00,1)
14:
end for
15:
Dr,t f- Oc.t Og
16:
Send Top-r sparsified updates Åe, + Top-r(Ao,) to the server.
end for
18:
On the Server Node:
20:
Receive Top-r sparsified updates Å., from clients C.
Aggregate them Å, +- Tiglobal Ee Âe,t
21:
Broadcast Å to the clients.
22:
23:
end for
G
Additional Results on Spectral Density of Predictor Hessian
10000
11000
10000
8000
9000
8000
7000
6000
4000
5000
2000
4000
3000
0
500
750
1000
1250
1500
1750
2000
250
2002
3250
1750
2000
Epoch
Epoch
(a) for BCE Loss
(b) A for MSE Loss
Figure 2: Estimate R of EELIAA over training iterations.
This section presents empirical results verifying our assumption on the sum of absolute eigenvalues of the predictor Hessian (Assumption 4.2p. Since it is infeasible to compute all eigenvalues for deep models like ResNet-18, we rely numerical approximations introduced in several prior works 24 78). SLQ is a technique The complete algorithm can be found in Algorithm 3 We use PyHessian [24],
numerical linear algebra to estimate the ESD of large matrices.
open-source implementation of SLQ our experiments. We keep the default parameters ni, = 10, m = 100 for the plots they have been shown to be of high accuracy. We train ResNet-13 18 from scratch model
24

---

10
S00
Gzech
Epoch-1500
100)
10"1
-300 Ligunvalus
Eigerraatus
Eigarnaaus
(a) Epoch 500
(b) Epoch 1000
(c) Epoch 1500
Figure 3: Spectral Density of Predictor Hessian : (0, x.) (2, x;) for a fixed training input across training epochs. Dataset: 1000 samples from CIFAR- 10 dataset. Model: ResNet-18 The model has 1.1 x 10? parameters. Loss function: Binary Cross Entropy(BCE) Loss.
Algorithm 3 Stochastic Lanczos Quadrature for ESD Computation (Ghorbani et al. 22
1: Input: Parameter: 0, degree m, and n..
2: Compute the gradient of 0 by backpropagation, i.e., compute go 4f(0:,.) d8
3: for i= 1,2,....7 Rh do
4: (Different Seeds)
5: Draw a random vector from (o,, 1) and normalize (same dimension as 0).
6: Get the tridiagonal matrix T through Lanczos algorithm.
7: Compute T (6) and 2 from T.
8: t
Ei= The ()u;tt))
9:
end for
Return d(t) En T (((
10:
to classify 1000 samples from two classes of CIFAR-10. We use a learning rate of 1e - 3, SGD as the optimizer and and perform GD. We use fix training sample plot the ESD of predictor Hessian H (0;x;) = é, VJJ f (0; x:) for two choices of loss functions BCE (Binary Cross Entropy) and MSE (Mean Squared Loss). We compute the spectrum by backpropagating through the output layer instead of the loss. From the results in Figure [] we note that bulk of eigenvalues
close zero, and the spectral density decays quickly far from zero, i.e., most of the eigenvalues are tiny. Given the approximate density as normalized histogram, we compute the empirical average and estimate & an approximation to Σ [A:l/Amae. Our results for & are presented in Figure 2
25

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: All of our claims supported by our theoretical and empirical results throughout the paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss limitations throughout the paper and also in Section 6
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
26

---

Justification:
Proof sketched are provided right after theorem statements and the full proofs are given in the Appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification:
We provide experimental details in Section 5 and Appendix F
Guidelines:
The answer NA means that the paper does include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
27

---

Answer: [No]
Justification: Due to confidentiality constraints, we are unable to share the code at this time, but we provide sufficient details reproduce the results.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips... cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run reproduce the results. See the NeurlPS code and data submission guidelines (https /nips.cc/public/giee guwde//ddee/// Cot///oodeeeeeseeeeeeee wwwpibmmsssiooooiooooooo
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We provide implementation details in Section 5 and Appendix F
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Computing spectral density for predictor Hessian is computationally de- manding. Therefore, we repeated the experiments only once with randomly selected homogeneous subset of the dataset.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
28

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative emrr rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide details in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification:
We followed the Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer:
[Yes]
Justification: provide a discussion on broader impact in Section
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
29

---

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We properly cited the models and datasets in Sections and 5 and in the appendix.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the Datasets.papersstt paperswi thcode. cow/ datasets package should be provided. For popular datasets,
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
30

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
31


Reviews:
[{"review_id": "bjpFaQzX4M", "reviewer": "Reviewer_wSrY", "summary": "This paper presents a novel analysis of sketching-based distributed learning algorithms for deep neural networks. The authors provide a tighter, dimension-independent convergence analysis by leveraging the restricted strong smoothness (RSS) property of deep learning loss functions. This work addresses a gap between existing theoretical analyses (which suggest dimension dependence) and empirical results (which show competitive performance without such dependence). The paper provides theoretical guarantees, improved communication complexity bounds, and empirical validation of their approach.", "questions": "1. Could you provide insights on how your approach might perform on other types of deep learning tasks beyond image classification (e.g., natural language processing or reinforcement learning)?\n2. What are the main challenges in implementing your approach in practical, large-scale distributed learning systems?", "limitations": "yes", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper presents the first dimension-independent convergence analysis for sketching in distributed deep learning without requiring restrictive assumptions like heavy-tailed distributions or top-r sparsification.\n2. The work bridges a gap between theory and practice, explaining why sketching-based methods perform well empirically despite previous pessimistic theoretical bounds. The authors present experimental results that support their theoretical findings, demonstrating the effectiveness of their approach on real-world datasets.\n3. The paper covers both single-step and multi-step local update scenarios, providing a thorough theoretical treatment.", "weaknesses": "1. While the paper removes some restrictive assumptions from previous work, it still relies on certain assumptions about the loss function (e.g., PL condition) and the eigenvalue distribution of the Hessian. The robustness of the results to violations of these assumptions could be discussed further.\n2. While the paper provides a thorough theoretical analysis, more discussion on the practical aspects of implementing the proposed approach in large-scale distributed learning systems could be beneficial.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are thankful to the reviewer for the suggestions and comments. Below we provide answers to each of the concerns:\n\n**Q1. Robustness to violation of assumption on eigenvalues:**\n 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\n    +\\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.\n\n---\n\n**Q2. Robustness to violation of assumption on PL condition:**\n 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}$.\n\n---\n\n | Loss function ($\\mathcal{L}$)  | Analysis | Iteration Complexity | Assumption  |\n|-------------------------------------|-----------|--------------------------|-------|\nPL  | Ours (Theorem 1) | $\\mathcal{O}\\left(\\frac{(1+\\varepsilon)^2c_s\\varrho^2 + \\frac{c_l c_H}{\\sqrt{m}}\\left[1+ \\kappa(\\varepsilon^2 + 2\\varepsilon)\\right]}{\\mu(1-\\varepsilon)^2}\\log\\left(\\frac{\\mathcal{L}(\\theta_0)-\\mathcal{L}(\\theta_{*})}{\\epsilon}\\right)\\right)$ | RSS ($\\kappa = o(1)$) |\nSC | [Song et al. 2023](Theorem E.1) | $\\mathcal{O}\\left(\\frac{\\left(c_s\\varrho^2 + \\frac{c_l c_H}{\\sqrt{m}}\\right)p(1+\\varepsilon)^2}{\\mu(1-\\varepsilon)^2}\\log\\left(\\frac{\\mathcal{L}(\\theta_0)-\\mathcal{L}(\\theta_{*})}{\\epsilon}\\right)\\right)$ | $\\beta$-smoothness|\nNon-convex | Ours (Theorem 2)| $\\mathcal{O}\\left(\\frac{(1+\\varepsilon)^2c_s\\varrho^2 + \\frac{c_l c_H}{\\sqrt{m}}\\left[1+ \\kappa(\\varepsilon^2 + 2\\varepsilon)\\right]}{(1-\\varepsilon)^2}\\left(\\frac{\\mathcal{L}(\\theta_0)-\\mathcal{L}(\\theta_{*})}{\\epsilon}\\right)\\right)$ | RSS ($\\kappa = o(1))$|\nNon-convex|[Song et al. 2023] (Theorem E.3) | $\\mathcal{O}\\left(\\frac{\\left(c_s\\varrho^2 + \\frac{c_l c_H}{\\sqrt{m}}\\right)p(1+\\varepsilon)^2}{(1-\\varepsilon)^2}\\left(\\frac{\\mathcal{L}(\\theta_0)-\\mathcal{L}(\\theta_{*})}{\\epsilon}\\right)\\right)$ | $\\beta$-smoothness|\n\nTable 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.\n\n---\n**Q3. Performance on tasks beyond image classification and practical challenges:**\n\nWhile 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.\n\n\n\n\n\n[Song et al. 2023] . Song, Y. Wang, Z. Yu, and L. Zhang. Sketching for first order method: Efficient algorithm for low-\nbandwidth channel and vulnerability. In International Conference on Machine Learning, pages 32365–\n32417. PMLR, 2023.\n\n[Rothchild et al. 2020] D. Rothchild, A. Panda, E. Ullah, N. Ivkin, I. Stoica, V. Braverman, J. Gonzalez, and R. Arora.\nFetchSGD: Communication-efficient federated learning with sketching. In H. D. III and A. Singh,\neditors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of\nMachine Learning Research, pages 8253–8265. PMLR, 13–18 Jul 2020\n\n[Melas-Kyriazi et al. 2022] L. Melas-Kyriazi and F. Wang. Intrinsic gradient compression for scalable and efficient federated\nlearning. In B. Y. Lin, C. He, C. Xie, F. Mireshghallah, N. Mehrabi, T. Li, M. Soltanolkotabi, and\nX. Ren, editors, Proceedings of the First Workshop on Federated Learning for Natural Language Processing\n(FL4NLP 2022), pages 27–41, Dublin, Ireland, May 2022. Association for Computational Linguistics.\ndoi: 10.18653/v1/2022.fl4nlp-1.4.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UXp12NopdN", "reviewer": "Reviewer_E7W7", "summary": "This paper provided ambient dimension-independent convergence and communication complexity analysis for sketching-based gradient descent algorithms using the second-order geometry of loss Hession.", "questions": "1. Why does this paper not study the more widely used SGD algorithm? This paper studies GD but compares it with FetchSGD in theory and experiment, which is not quite reasonable.\n\n2. The theoretical result of the multiple-local step in Theorem 2 does not converge to 0 with the number of iterations, so why do the authors say that the result is similar to the sublinear convergence rate of SGD? Also, the authors should provide more explanatory notes on why the GD algorithm using multiple-local step is similar to SGD.\n\n3. What is the notation $\\\\mathbf{v}\\_{0}$ in Assumption 3.3? It is not defined in the text nor used elsewhere. Is the notation taken from the literature [28]?\n\n4. As the authors state, this paper does not make innovations to the algorithm. Therefore, the authors should focus on describing the theoretical contributions of this paper and should not demonstrate the superiority of the algorithm by unfair comparison with flawed algorithms (biased compression algorithms without error feedback).\n\n5. The authors did not provide code, nor did they describe the experimental setup in detail. The reviewer wonders if it is feasible to simulate 100 clients for full-batch gradient descent (i.e., a batch size of 500 per client) using a single NVIDIA Titan X GPU. Also, it is not reasonable to run the mini-batch based FetchSGD algorithm using the full batch setting.\n\n6. The symbolic notation and written presentation in lines 137 to 153 of this paper are very similar to that of literature [28, Sections 3 and 4].\n\n7. The references of the paper are badly formatted, and many of them do not contain information on publication conferences or journals, e.g. [57, 58, 59, 60, 61, 62, 63, 64, 65, 67, 70, 71, 76].", "limitations": "Please refer to Weaknesses and Questions.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "* This paper proves that the dimensional dependence comes from the global smoothness assumption and provides a convergence analysis of the sketching method independent of the ambient dimensions.\n\n* The analysis of this paper does not require the heavy hitter assumption and the Top-$r$ sparsification operation.", "weaknesses": "* This paper studies gradient descent (GD) and uses additional assumptions, limiting the application of the analytical results.\n\n* Comparison experiments are not reasonable, and the experimental results are roughly plotted.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for carefully reading our paper and providing constructive feedback. Below we provide answers to each of the concerns:\n\n\n**Q1: SGD v/s GD**:\nFirstly 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.\n\n---\n\n**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. \n\n---\n\n  **Q3: Similarity to SGD:**\n    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.\n    \n   ---\n    \n **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. \n    \n ---\n \n**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. \n___\n\n   **Q6.Symbolic notation and citation errors:**\nWe 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.\n\n\n[Karimi et al] Linear convergence of gradient and proximal-gradient methods under the polyak-ojasiewicz condition, 2020. URL  https://arxiv.org/abs/1608.04636\n\n[Banerjee et al.] Restricted strong convexity of deep learning  models with smooth activations.  In  The Eleventh International Conference on Learning Representations, 2023. \n\n[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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9pVZCQZ1wR", "reviewer": "Reviewer_LRSh", "summary": "This paper considers the sketch-DL framework for distributed / federated learning studied by prior works such as Rothchild et al., Song et al. The authors identify that for these works, their convergence either has a dependence on the dimension (under rather minimal assumptions) which is unfavorable, or assumes something stronger such as the sketched gradient has heavy-hitter coordinates and it is enough to recover using Top-$r$. The model studied in this paper is a deep, feedforward neural net with a 1-Lipschitz activation (say ReLU). They assume the loss function is PL and the eigenvalues of the predicator Hessian are dominated by a few top eigenvalues. Under these assumptions, they manage to remove the linear dependence on the dimension for convergence from Song et al. They also perform experiments to confirm that using sketching matrices such as CountSketch gives better performance than local top-$k$ or FetchSGD due to Rothchild et al. without error feedback.", "questions": "A few typos:\n\n* In References, citation [56] should be Sketching as a tool for numerical linear algebra instead of Computational advertising: Techniques for targeting relevant ads.\n\n* On line 592, end of page 15, what is the matrix $\\\\|\\cdot \\\\|_2$ norm? Is that the spectral norm?\n\nA few questions:\n\n* Can you extend the analysis to non-PL losses (say the loss is only convex, or even non-convex) similar to Song et al.?\n\n* I think it will be greatly beneficial to have a table comparing the set of assumptions posed in this paper and Song et al., together with different convergence results.", "limitations": "Yes.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This paper resolves a big issue from Song et al., where the convergence depends linearly on the dimension. They are able to provide a tighter analysis and a sharper bound by examining a different set of assumptions that utilize information of the Hessian, going beyond first-order information. This is quite important as in some sense, this paper obtains the best of both worlds: the algorithm it examines is the simplest form of applying linear sketching due to Song et al., where one could easily integrate extra mechanism such as DP into it effortlessly. The new set of assumptions imposed in this paper are quite reasonable, and it also performs experiments to confirm simply applying sketches are enough to guarantee fast convergence.", "weaknesses": "There are two main weaknesses of this paper.\n\n* Fair comparison with the analysis of Song et al. This paper makes 3 assumptions: 1. the loss function is PL, 2. the eigenvalues of Hessian are dominated by a few top eigenvalues and 3. a uniform upper bound on the $\\ell_2$ norm of the gradient. For fair comparison, the corresponding results in Song et al. are assuming 1. the loss function is strongly-convex and 2. the loss function is smooth. Note that Song et al. does not need to assume a uniform upper bound on the size of all gradients, which is commonly criticized as an unrealistic assumption. Moreover, Song et al. provides a slew of results for convex and non-convex loss function. In my opinion, this paper should provide a more comprehensive comparison with Song et al. For example, what is the crux to remove the linear dependence on the dimension? Is it due to PL, restricted strong smoothness (RSS) or uniform bound on the gradient? I would assume it's due to the RSS assumptions, but the authors should explicitly spell it out. Also, is it possible to remove the gradient upper bound assumption?\n\n* Experimental comparison with Rothchild et al. Figure 1 presented in the main body of the paper only compares the sketching algorithm with FetchSGD without error feedback, and hence obtains better dimension reduction ration v.s. accuracy, which is not surprising, as FetchSGD *requires* error feedback correction to make the algorithm correct. In Figure 2 on page 23, authors did compare with FetchSGD with error feedback, and the \"correct\" FetchSGD does have the superior performance. I do understand authors intention here as the vanilla FetchSGD is nonlinear, hence very hard to integrate any other pieces such as DP into it. From this perspective, the sketching algorithm in Song et al. slightly trades the performance for a simpler and more general algorithm. Still, I don't think it's fair to put a figure in the main body comparing the sketching algorithm with a \"wrong\" version of FetchSGD. Authors could probably put Figure 2 in the main body, and provide comprehensive justifications on why one would expect FetchSGD to perform better and why sketching algorithm could potentially yield more use cases in practice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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:\n\n**Q1. Comparison with [Song et al. 2023]:**\nBefore comparing, we briefly summarise the difference in our setup, assumptions and results:\n1.  **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.\n\n3. **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]. \n 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.\n\nUnder 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. \n\n| Loss function ($\\mathcal{L}$) | Analysis | Iteration Complexity | Assumption |\n|-------------------------------------|-----------|--------------------------|-------|\nPL | Ours (Theorem 1) | $\\mathcal{O}\\left(\\frac{(1+\\varepsilon)^2c_s\\varrho^2 + \\frac{c_l c_H}{\\sqrt{m}}\\left[1+ \\kappa(\\varepsilon^2 + 2\\varepsilon)\\right]}{\\mu(1-\\varepsilon)^2}\\log\\left(\\frac{\\mathcal{L}(\\theta_0)-\\mathcal{L}(\\theta_{*})}{\\epsilon}\\right)\\right)$ | RSS ($\\kappa = o(1)$) |\nSC | [Song et al. 2023](Theorem E.1) | $\\mathcal{O}\\left(\\frac{\\left(c_s\\varrho^2 + \\frac{c_l c_H}{\\sqrt{m}}\\right)p(1+\\varepsilon)^2}{\\mu(1-\\varepsilon)^2}\\log\\left(\\frac{\\mathcal{L}(\\theta_0)-\\mathcal{L}(\\theta_{*})}{\\epsilon}\\right)\\right)$ | $\\beta$-smoothness|\nNon-convex | Ours (Theorem 2)| $\\mathcal{O}\\left(\\frac{(1+\\varepsilon)^2c_s\\varrho^2 + \\frac{c_l c_H}{\\sqrt{m}}\\left[1+ \\kappa(\\varepsilon^2 + 2\\varepsilon)\\right]}{(1-\\varepsilon)^2}\\left(\\frac{\\mathcal{L}(\\theta_0)-\\mathcal{L}(\\theta_{*})}{\\epsilon}\\right)\\right)$ | RSS ($\\kappa = o(1))$|\nNon-convex|[Song et al. 2023] (Theorem E.3) | $\\mathcal{O}\\left(\\frac{\\left(c_s\\varrho^2 + \\frac{c_l c_H}{\\sqrt{m}}\\right)p(1+\\varepsilon)^2}{(1-\\varepsilon)^2}\\left(\\frac{\\mathcal{L}(\\theta_0)-\\mathcal{L}(\\theta_{*})}{\\epsilon}\\right)\\right)$ | $\\beta$-smoothness|\n\nTable 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.\n\n| Loss function ($\\mathcal{L}$) | Analysis | Iteration Complexity | Assumption |\n|-------------------------------------|-----------|--------------------------|-------|\nPL | Ours (Theorem 2) | $\\mathcal{O}\\left(\\frac{\\varrho^2 + \\frac{c_H}{\\sqrt{m}}+\\frac{\\varepsilon\\kappa c_H}{\\sqrt{m}}}{\\mu^2(1-\\varepsilon)^2\\epsilon}\\log\\left(\\frac{\\mathcal{L}(\\theta_0)-\\mathcal{L}(\\theta^{*}}{\\epsilon}\\right)\\right)$ | RSS, Bounded Gradient |\nSC | [Song et al. 2023](Theorem F.10 ($\\sigma$=0)) | $\\mathcal{O}\\left(\\frac{p\\beta K}{\\mu}\\log\\left(\\frac{\\beta\\mathbb{E}[\\vert\\vert\\theta_0-\\theta^{*}\\vert\\vert^2_2]}{\\epsilon}\\right)\\right)$ | $\\beta$-smoothness |\n\nTable 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. \n\n---\n\n**Q2.  Regarding Figure 2**\nThank 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. \n\n---\n**Q3.Typos**: \nOn 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.\n\n[Zhou et al.] Characterization of gradient dominance and regularity conditions for neural  \nnetworks, 2017. URL  https://arxiv.org/abs/1710.06910. \n[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", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
**General Response**:

We thank the reviewers for their positive comments about our work and their efforts in providing meaningful feedback. In particular, we thank Reviewer `LRSh` for stating that our results are **quite important** ,**resolve big issues from previous work** under **reasonable assumptions**, and have meaningful implications for downstream works in Differential Privacy; Reviewer `wSrY` for finding that our work is **thorough** and **addresses a gap between theory and practice** and lastly, Reviewer `E7W7` for pointing out that our work **gets rid of unrealistic assumptions**.

We have addressed the questions and comments of reviewers in individual responses. Below we provide a brief description of the contributions and the novel aspects of our work:

1. We identify a gap in theory and practice of sketching methods for distributed learning -- sketching methods incur a dependence on ambient dimension in the convergence rate [Song et al. 2023] while in practice they are observed to be at par with their uncompressed counterpart ( Appendix B.3 of [Rothchild et al. 2020])

2. We identify and demonstrate the widely used global smoothness assumption of loss functions as the root cause of the dimension dependence. To alleviate this dimensional dependence, [Rothchild et al. 2020] make assumptions on the heavy-tailed nature of gradient distribution and perform additional Top-r sparsification resulting in a biased algorithm.

3. We get rid of these dimensional dependence while avoiding these assumptions by exploiting the second-order structure of Loss Hessians and the fast decay of Hessian eigenvalues [Ghorbani et al.],[Zhang et al.],[Xie et al.].  Our results instead depend on the sum of the absolute values of the eigenvalues of the Hessian which is potentially much smaller. We also provide empirical evidence verifying this claim.


[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

[Ghorbani et al.] An investigation into neural net optimization via hessian eigenvalue density. In  International Conference on Machine Learning, pages 2232–2241. PMLR, 2019.
[Zhang et al.] Y. Zhang, C. Chen, T. Ding, Z. Li, R. Sun, and Z.-Q. Luo. Why transformers need adam: A hessian  perspective.  arXiv preprint arXiv:2402.16788, 2024.
[Xie et al.] . Xie, Q.-Y. Tang, Y. Cai, M. Sun, and P. Li. On the power-law hessian spectrums in deep learning,  2022. URL  https://arxiv.org/abs/2201.13011

Meta-review: The paper considers distributed learning with compressed communication of updates. It gives novel theoretical insights if sketching is used as the compression mechanism, under realistic assumptions, in particular studying the dimension dependence. Some concerns remained that the baseline algorithms in the experiments were weak (should be changed to include methods using error-feedback to ensure convergence). Nevertheless, overall positive assessment remained, and we hope the authors will incorporate the mentioned feedback by the reviewers.
Decision: Accept (poster)