Title: Sketching for Distributed Deep Learning: A Sharper Analysis
Venue: NeurIPS 2024
Accepted: yes
Keywords: sketching

distributed learning

federated learning

optimization
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

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 [62] or need additional restrictive assumptions such as heavy-hitters [59]. 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.

1 Introduction

Distributed learning is a 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 a 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 a result of the system setup, distributed learning generally requires a high communication overhead due to the frequent communication between the server and the clients [42]. Federated learning, a popular distributed learning setup, further extends the generic setting to 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 in compressing

Berivan Isik ∗

Sanmi Koyejo

the model updates [35, 43] through sparsification [44, 73, 7, 30, 31], quantization [4, 74, 10], lowrank projection [71, 67, 68, 8], and sketching [59, 62, 50, 33, 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, 9] as opposed to the popular Top-r 2 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.e., in contrast to the lower sketching dimension) which may be extremely large for modern overparameterized deep models [39, 20, 11]. This dimension dependence has been a 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, 59]. 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 a 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 eliminate 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 [6, 72, 22, 53, 52, 15], and our results are 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 to Table 1 for an overview of our results. We do not restrict our analysis to any specific sketching matrix, and our results hold for any symmetric sub-Gaussian sketching matrix. For instance, for computational benefits, one can use the popular Count-Sketch [14] or Hadamard sketch [2] approaches, which are both examples of sub-Gaussian sketching. Our contributions can be summarized as follows:

Weidentify and demonstrate the widely used global smoothness assumption of loss functions as the root cause of the dimension dependence of sketching methods.

Weprovide a novel analysis for sketching in distributed and federated learning that eliminates the ambient dimension dependence in the convergence error.

We are the first to break this dimension dependence without making restrictive assumptions such as the heavy hitter assumption by Rothchild et al. [59].

We are again the first to do this without a Top-r sparsification step in the framework which would have required additional measures to eliminate the bias.

We are the first to use 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.

2 Related Work

Communication-Efficient 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 to 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, 71], sketching [59, 62, 34], and sparse subnetwork training techniques [31, 40, 41, 47, 32]. While some of these compression methods are already unbiased [4], many are biased and have to be combined with other mechanisms to reduce the bias for better convergence [44, 59]. Linearity is another desired feature, as it simplifies the implementation of distributed learning with security-promoting techniques like secure aggregation (in the compressed dimension). Among the general compression approaches mentioned, sketching stands out as a simple linear and unbiased operation, allowing for computations in the reduced dimension before desketching. We do not propose a new compression method in this work but instead provide a substantially improved convergence analysis for sketching-based frameworks that breaks the dimension dependence-which explains why sketching would not explode the converge error in large models.

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 communication efficiency [34, 33, 59, 62, 27]. The linearity of these sketching-based frameworks has led to their successful integration with secure aggregation and differential privacy as well [16, 66, 62, 17]. Despite the empirical success of these sketching-based applications in distributed and federated learning, the existing upper bounds on the convergence errors have a 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 a

dimension-independent convergence bound for their sketching algorithm, called FetchSGD (which is a combination of Count-Sketch 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 FetchSGD 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.

Notation. For a quantity x , we use the notation x t to refer to the global variable at round t shared by all the clients and the server, and x c,t to refer to the local variable for client c at round t . For a positive integer n , [ n ] = { 1 , · · · , n } . We use I p × p as the p × p identity matrix. We use E [ · ] for expectation. For a vector x , we use ∥ x ∥ or ∥ x ∥ 2 to denote its L 2 norm. For a matrix A , we use ∥ A ∥ or ∥ A ∥ 2 for the spectral norm of A . We use ∇ θ L ( θ ′ ) = ∂ L ∂θ | θ = θ ′ . For a random vector y , ∥ y ∥ ψ 2 denotes its sub-Gaussian norm [70]. The notations ˜ O ( t ) , ˜ Ω( t ) , ˜ Θ( t ) are the same as the common O ( t ) , Ω( t ) , Θ( t ) but they hide the dependence on logarithmic terms. polylog( n ) denotes O (log k ( n )) for some k .

3 Sketching for Distributed Learning

We consider a distributed learning framework with C clients, each client c = 1 , . . . , C having a local dataset D c = { x i,c , y i,c } n c i =1 of size n c and a local loss L c : R p → R defined as L c ( θ ) = 1 n c ∑ n c i =1 ℓ ( y i,c , ˆ y i,c ) , where θ ∈ R p is the parameter vector, ˆ y i,c := f ( θ ; x i,c ) is the prediction of the model for input x i,c , and ℓ : R × R → R is a loss function that measures the error between ˆ y i,c and the groundtruth y i,c . Our goal is to minimize the empirical loss L ( θ ) :

Unlike much of the existing literature [59, 33, 62], our analysis does not ignore the fact that the predictor f is a deep learning model - i.e., the loss L c is not just an arbitrary smooth loss. The motivation behind this focus is that many of the modern models being used in distributed learning are 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.

Following standard literature [22, 54, 45, 6], we consider a fully-connected feed-forward neural network f of depth L , with widths m and activations α ( l ) for each layer l ∈ [ L ] := { 1 , . . . , L } , described as:

where W ( l ) t ∈ R m l × m l -1 is the layer-wise weight matrix for layer l ∈ [ L ] and v t ∈ R m L is the last layer vector at iteration t , ϕ ( · ) is the smooth (pointwise) activation function, and m 0 = dim( x ) = d . We denote the total set of parameters as

For simplicity, we will assume that the width of all the layers is the same, i.e., m l = m for all l ∈ [ L ] , and thus p = ( L -1) m 2 + md + m . We consider deep models with only one output, i.e., f ( θ ; x ) ∈ R , but our results can be extended to multi-dimensional outputs.

As prevalent in literature [22, 6], we make the following assumptions regarding the activation function ϕ , loss function ℓ , and the initialization parameter θ 0 which hold true for the commonly used activation functions, loss functions, and initializations used in practice.

Assumption 3.1 ( Activation function ) . The activation ϕ is 1-Lipschitz and β ϕ -smooth,i.e. | ϕ ′ | ≤ 1 and | ϕ ′′ | ≤ β ϕ .

Assumption 3.2 ( Initialization ) . We initialize θ 0 with w ( l ) 0 ,ij ∼ N (0 , σ 2 0 ) for l ∈ [ L ] where σ 0 = σ 1 2 ( 1+ √ log m √ 2 m ) , σ 1 > 0 , and v 0 is a random unit vector with ∥ v 0 ∥ 2 = 1 .

Assumption 3.3 ( Loss function ) . The loss ℓ i = ℓ ( y i , ˆ y i ) with ℓ ′ i = dℓ i d ˆ y i , ℓ ′′ i = d 2 ℓ i d ˆ y i is (i) Lipschitz, i.e., | ℓ ′ i | ≤ c l , and (ii) ℓ ′′ i and smooth ℓ ′′ i ≤ c s for some c l , c s > 0 .

3.1 Sketching-Based Distributed Learning

Random sketching [75, 48] 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 ∈ R b × p where typically b ≪ p . Examples of sketching matrices include Count-Sketch [14], Subsampled Randomized Hadamard Transforms ( SRHT ) [13], and sparse Johnson-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 as 'desketching'.

We outline the sketching-based distributed learning framework in Algorithm 1. Each client receives a random seed from the server to initialize the local parameters θ c, 1 , and generate a sketching matrix R . At each local step k ∈ [1 , · · · , K ] , each client performs local gradient descent (GD) over their local dataset D c . At each communication round, the client accumulates the changes over K -local 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:

While we use the same sketching matrix across communication rounds t = 1 , . . . , T , in general, using different matrices for each round does not affect the analysis.

Choice of sketching matrix: Weuse a (1 / √ b ) -sub-Gaussian matrix as the choice of sketching matrix. We say R ∈ R b × p is a (1 / √ b ) -sub-Gaussian matrix [69] if each row R i is an independent mean-zero, sub-Gaussian isotropic random-vector such that ∥ R i ∥ ψ 2 ≤ 1 / √ b . We assume E [ R ⊤ R ] = I p × p . From the above definition, we can see that for g 1 , g 2 ∈ R p

3.2 Limitations of the Existing Analyses

When analyzing the convergence rates of the sketching-based distributed learning frameworks, previous works [62, 59] assume that the loss function L is β -smooth, i.e,

In the model updates based on sketching followed by desketching, the term ∥ θ ′ -θ ∥ 2 effectively yields a term of the form ∥ R ⊤ Rg ∥ 2 2 , where g ∈ R p stands for a suitable gradient on the full model. While E [ R ⊤ Rg ] = g , i.e., the desk -sk operation is unbiased, we unfortunately have ∥ R ⊤ Rg ∥ 2 2 = Θ( p b ∥ g ∥ 2 2 )

Algorithm 1 Sketching-Based Distributed Learning.

Hyperparameters: server learning rate η global, local learning rate η local .

Inputs: local datasets D c of size n c for clients c = 1 , . . . , C , number of communication rounds T . Output: final model θ T .

1: Broadcast a random SEED to the clients.

2: for t = 1 , . . . , T do

3: On Client Nodes:

4: for c = 1 , . . . , C do

5: if t = 1 then

6: Receive the random SEED from the server. Initialize the local model θ c, 1 ∈ R p and generate the sketching matrix R ∈ R b × p (hence sk , desk ) using the random SEED .

7: else

8: Receive sk( ¯ ∆ t -1 ) from the server.

9:

Desketch and update the model parameters θ t ← θ t -1 +desk(sk( ¯ ∆ t -1 )) .

10:

Assign the local model's parameters θ c,t ← θ t to be updated locally.

11: end if

12: for k = 1 , . . . , K do 13: θ c,t ← θ c,t -η local · ∇ θ L c ( θ c,t )

14: end for

15: ∆ c,t ← θ c,t -θ t

16: Send sketched updates sk(∆ c,t ) to the server.

17: end for

18:

19: On the Server Node:

20: Receive sketched updates sk(∆ c,t ) from clients c = 1 , . . . , C .

21: Aggregate: sk( ¯ ∆ t ) ← η global · 1 C ∑ C c =1 sk(∆ c,t )

22: Broadcast sk( ¯ ∆ t ) to the clients.

23: end for

with high probability (see Lemma B.2 in Appendix). Thus, such an analysis picks up a dimension dependence Θ( p ) which can be neutralized only if the sketching dimension is b = Ω( p ) . However, such a high-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 in 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 of the gradient vector, with an analysis having to rely on heavy-hitter assumptions. Interestingly, Rothchild et al. [59] noted that the sketching-based distributed deep learning approach seemed to work fine empirically without getting the adverse effect of dimension dependence despite what their theoretical results (based on smoothness) suggest. Our work sheds light on this discrepancy and proves (see Section 4) why the sketching-based distributed learning approach in its simplest form (see Algorithm 1) does not pick up the dimension dependence.

3.3 Restricted Strong Smoothness (RSS)

In this section, we describe the RSS property, an interesting property of deep learning losses that we use to derive dimension-independent convergence guarantees for sketch-DL. Using Taylor's expansion, the loss at θ = θ ′ can be written as:

where ˜ θ = ϵθ ′ +(1 -ϵ ) θ t , ϵ ∈ [0 , 1] and ∇ 2 θ L ( ˜ θ ) is 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 ):

where ∇ f i,c = ∇ θ f ( θ ; x i,c ) and H i,c = ℓ ′ i,c ∇ 2 f i,c = ℓ ′ i,c ∇ 2 θ f ( θ ; x i,c ) respectively. Recent works [45, 6, 22] have derived sharp upper bounds on the spectral norm of the predictor Hessian: ∥∇ 2 f ∥ = Λ max ( ∇ 2 f ) = O (1 / √ m ) . This leads to the following restrcited smoothness property of deep learning losses:

In contrast to the β -smoothness assumption which is common in optimization literature, this new property shows that deep learning losses exhibit strong smoothness in a restricted set of directions. The effect of strong smoothness suitably manifests for θ ′ such that | cos( θ ′ -θ t , ∇ f i,c ) | ≥ κ > 0 for a restricted set of directions. In other directions, the strong smoothness constant is O (1 / √ m ) , i.e., a tiny value. For our analysis, we use an even sharper perspective on RSS based on the eigenvalues of predictor Hessian, H i,c = ℓ ′ i ∇ 2 f i,c . While the smoothness perspective based on the spectral norm of O (1 / √ m ) is promising, it implicitly assumes all directions have this level of smoothness which would impact the analysis since there are p = O ( Lm 2 ) directions. Based on empirical evidence [55, 26], most directions have smoothness much smaller than 1 / √ m , and an analysis based on the eigenvalues of the H i,c captures this sharper perspective as opposed to picking up a dependence on 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 1, using the RSS property of deep learning losses and provide novel dimension-independent convergence results. First, we state our key assumptions in Section 4.1 and explain why they are well-supported by recent work. Then, we provide our novel analysis that 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 to the multiple-local step ( K > 1 ) case. We conclude in Section 4.4 by deriving the communication cost of Algorithm 1 under our novel convergence analysis.

4.1 Assumptions

Before we introduce the assumptions, we would like to recall a key property.

Definition 1 (PL condition) . Consider a loss function L : R p → R and the solution set X ∗ = { θ ′ : θ ′ ∈ argmin θ L ( θ ) } and we use L ∗ to denote the corresponding minimum value. Then, L is said to satisfy the Polyak-Łojasiewicz (PL) condition with constant µ if

Remark 4.1. PL condition [57] can be used to establish linear convergence of gradient descent while still being weaker than strong convexity. One can show that if a function f is µ -strongly convex then it also satisfies PL condition with the same constant µ . 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 in 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 H i,c as follows:

Assumption 4.2 (Predictor Hessian Eigenspectrum) . Let Λ i,c,t, 1 , Λ i,c,t, 2 , · · · , Λ i,c,t,p be the eigenvalues of the predictor Hessian H i,c,t = ℓ ′ i,c ∇ 2 f ( θ t ; x i,c ) for t ∈ [ T ] and Λ max = max j | Λ i,c,t,j | then there exists κ = O (1) such that ∑ p j =1 | Λ i,c,t,j | ≤ κ Λ max.

Remark 4.2. Several works [24, 49, 76] have shown that the spectrum of loss Hessian follows a bulk and outliers structure where the bulk can be attributed to H and follows power law trend [Figure 3(b), Figure 4 in Papyan [56]], motivating our assumption. Further empirical evidence is presented in Appendix G (refer to 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 1) can be written as θ t ← θ t -1 -η desk (sk( ∇ θ L ( θ t -1 ))) . We follow a similar analysis to Song et 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.1.

Theorem 4.1 (Informal version of Theorem D.1) . Set K=1 in Algorithm 1. Under Assumptions 4.1, and 4.2, for suitable constants ε, δ < 1 , learning rate η = η global · η local independent of p , and b = Ω ( 1 ε 2 polylog( p 2 NT δ ) ) , with probability at least 1 -δ , we have:

where θ ∗ is a minimizer of (1).

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 as Top-r sparsification (as done by Rothchild et al. [59]), which also introduces a bias that needs to be suitably corrected.

4.3 Multiple-Local Step ( K > 1)

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 to guarantee convergence. For the purpose of our analysis, we assume that the gradient norms are bounded, i.e., ∥∇ θ ℓ ( θ ) ∥ ≤ G . We refer the reader to Theorem D.2 in the Appendix for a more formal statement. We now present an informal version of our theorem for the convergence of the multiple-local step case:

Theorem 4.2 (Informal version of Theorem D.2) . Let ∥∇ θ ℓ ( θ ) ∥ ≤ G . and 4.2, for a suitable constants ϱ, c and ε < 1 , learning rate η = η η < 1 , and

Under Assumptions 4.1 H global local 2 µK (1 -ε ) , we have:

where θ ∗ is a minimizer of the problem 1, N is the number of training samples, and

We refer the reader to Appendix D.2 for the full proof.

4.4 Communication Efficiency

As a direct consequence of our analysis, we now provide an improved communication complexity for Algorithm 1 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 a substantial improvement over the analysis of Song et al. [62]. Unlike Rothchild et al. [59], we break this dependence without requiring a Top-r sparsification step or a heavy-hitter assumption. In the theorem below, we state the required number of communication bits to achieve an ϵ -optimal solution based on our sharper convergence analysis in the previous sections.

Theorem 4.3 (Informal version of Theorem E.1) . Under Assumptions 4.1 and 4.2, Algorithm 1 obtains an optimal solution satisfying the error

We provide the full proof in the Appendix E. We state the result by Song et al. [62] for comparison:

Song et al. [62] assumes that the loss function L is β -smooth and the expectation is over the randomness of sketching. Note that our bound depends on κ instead of the ambient dimension p , presenting a sharper bound. This shows the efficacy of our approach in improving the overall communication complexity over prior works.

5 Experimental Results

In this section, we provide a comparison of the sketching approach in Algorithm 1 with other common approaches such as local Top-r [44] and FetchSGD [59]. We note that both Local Top-r (as outlined in Algorithm 2 in 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 in Algorithm 1, we use Count-Sketch . This means the only difference between FetchSGD without error feedback and Count-Sketch is the extra global Top-r step at the server. We conducted our experiments on NVIDIA Titan X GPUs on an internal cluster server, using 1 GPU per one run.

We train ResNet-18 [28] on CIFAR-10 dataset [38] that is i.i.d. distributed to 100 clients. Each client performs 5 local gradient descent iterations (i.e., using full-batch of size 500 ) at every round. Figure 1 shows that Count-Sketch -based distributed learning approach in Algorithm 1 performs competitively with FetchSGD . This result highlights the potential of sketching alone, without additional modifications as in FetchSGD , to maintain competitive accuracy. Additionally, the error-feedback free approach enables compatibility with Differential Privacy(DP) techniques which we leave as future work.

0.94

0.92

0.90

LocalTop-rwitherrorfeedback

0.88

LocalTop-rwithouterrorfeedback

0.86

FetchSGDwitherrorfeedback

0.84

FetchSGDwithouterrorfeedback

0.82

Count-Sketch

0.80

0.78

5

10

15

20

25

30

35

DimensionReduction

We note that we do not claim novelty of any of the methods discussed in this section as they have been introduced [44, 34], improved [59], and analyzed [62] extensively in prior work. We present these empirical comparisons for completeness and to provide support for why a 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 et al. [59].

6 Discussion and Conclusion

We provide a significantly improved convergence analysis for sketching-based 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 in 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.

Limitations. We provide a sharper analysis for sketching in distributed learning. In future work, we plan to extend our analysis to federated learning by allowing client dropout.

Broader Impact. This paper provides a tighter analysis for sketching-based distributed learning and federated learning frameworks. We expect this work to be helpful for the community as it contributes to the efforts in making machine learning models more decentralized, accessible, and trustworthy.

References

K. J. Ahn, S. Guha, and A. McGregor. Graph sketches: sparsification, spanners, and subgraphs. In Proceedings of the 31st ACM SIGMOD-SIGACT-SIGAI symposium on Principles of Database Systems , pages 5-14, 2012.

N. Ailon and B. Chazelle. Approximate nearest neighbors and the fast johnson-lindenstrauss transform. In Proceedings of the thirty-eighth annual ACM symposium on Theory of computing , pages 557-563, 2006.

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.

D. Alistarh, D. Grubic, J. Li, R. Tomioka, and M. Vojnovic. Qsgd: Communication-efficient sgd via gradient quantization and encoding. Advances in neural information processing systems , 30, 2017.

A. Banerjee, S. Chen, F. Fazayeli, and V. Sivakumar. Estimation with norm regularization. Advances in neural information processing systems , 27, 2014.

A. Banerjee, P. 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.

L. 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.

R. B. Basat, S. Vargaftik, A. Portnoy, G. Einziger, Y . Ben-Itzhak, and M. Mitzenmacher. QUICKFL: Quick unbiased compression for federated learning. arXiv preprint arXiv:2205.13341 , 2022.

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.

J. Bernstein, Y.-X. Wang, K. Azizzadenesheli, and A. Anandkumar. signsgd: Compressed optimisation for non-convex problems. In International Conference on Machine Learning , pages 560-569. PMLR, 2018.

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 preprint arXiv:2108.07258 , 2021.

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 arXiv:1611.04482 , 2016.

C. Boutsidis and A. Gittens. Improved matrix algorithms via the subsampled randomized hadamard transform. SIAM Journal on Matrix Analysis and Applications , 34(3):1301-1340, 2013. doi: 10.1137/120874540. URL https://doi.org/10.1137/120874540 .

M. Charikar, K. Chen, and M. Farach-Colton. Finding frequent items in data streams. In International Colloquium on Automata, Languages, and Programming , pages 693-703. Springer, 2002.

S. Chen and A. Banerjee. Structured estimation with atomic norms: General bounds and applications. Advances in Neural Information Processing Systems , 28, 2015.

W.-N. Chen, C. 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.

W.-N. Chen, B. Isik, P. Kairouz, A. No, S. Oh, and Z. Xu. Improved communication-privacy trade-offs in $l_2$ mean estimation under streaming differential privacy. In Forty-first International Conference on Machine Learning , 2024. URL https://openreview.net/forum? id=x1G7ieRgRd .

G. Cormode and S. Muthukrishnan. An improved data stream summary: the count-min sketch and its applications. Journal of Algorithms , 55(1):58-75, 2005.

A. Dasgupta, R. Kumar, and T. Sarlos. A sparse johnson: Lindenstrauss transform. In Proceedings of the Forty-Second ACM Symposium on Theory of Computing , STOC '10, page 341-350, New York, NY, USA, 2010. Association for Computing Machinery. ISBN 9781450300506. doi: 10.1145/1806689.1806737. URL https://doi.org/10.1145/1806689.1806737 .

J. Dean, G. Corrado, R. Monga, K. Chen, M. Devin, M. Mao, M. Ranzato, A. Senior, P. Tucker, K. Yang, et al. Large scale distributed deep networks. Advances in neural information processing systems , 25, 2012.

E. Dobriban and S. Liu. Asymptotics for sketching in least squares regression. Advances in Neural Information Processing Systems , 32, 2019.

S. Du, J. Lee, H. Li, L. Wang, and X. Zhai. Gradient descent finds global minima of deep neural networks. In International conference on machine learning , pages 1675-1685. PMLR, 2019.

C. Dwork, A. Roth, et al. The algorithmic foundations of differential privacy. Foundations and Trends® in Theoretical Computer Science , 9(3-4):211-407, 2014.

B. Ghorbani, S. Krishnan, and Y. Xiao. An investigation into neural net optimization via hessian eigenvalue density. In International Conference on Machine Learning , pages 2232-2241. PMLR, 2019.

M. Greenwald and S. Khanna. Space-efficient online computation of quantile summaries. ACM SIGMOD Record , 30(2):58-66, 2001.

G. Gur-Ari, D. A. Roberts, and E. Dyer. Gradient descent happens in a tiny subspace, 2018.

F. Haddadpour, B. Karimi, P. Li, and X. Li. Fedsketch: Communication-efficient and private federated learning via sketching. arXiv preprint arXiv:2008.04975 , 2020.

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 770778, 2016.

I. C. F. Ipsen and A. K. Saibaba. Stable rank and intrinsic dimension of real and complex matrices, 2024. URL https://arxiv.org/abs/2407.21594 .

B. Isik, T. Weissman, and A. No. An information-theoretic justification for model pruning. In International Conference on Artificial Intelligence and Statistics , pages 3821-3846. PMLR, 2022.

B. Isik, F. Pase, D. Gunduz, T. Weissman, and Z. Michele. Sparse random networks for communication-efficient federated learning. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=k1FHgri5y3-.

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.

N. Ivkin, D. Rothchild, E. Ullah, I. Stoica, R. Arora, et al. Communication-efficient distributed sgd with sketching. Advances in Neural Information Processing Systems , 32, 2019.

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.

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 in federated learning. Foundations and Trends® in Machine Learning , 14(1-2):1-210, 2021.

D. M. Kane and J. Nelson. Sparser johnson-lindenstrauss transforms. Journal of the ACM (JACM) , 61(1):1-23, 2014.

H. Karimi, J. Nutini, and M. Schmidt. Linear convergence of gradient and proximal-gradient methods under the polyak-Łojasiewicz condition, 2020.

A. Krizhevsky, G. Hinton, et al. Learning multiple layers of features from tiny images.(2009), 2009.

Y. LeCun, Y. Bengio, and G. Hinton. Deep learning. nature , 521(7553):436-444, 2015.

A. Li, J. Sun, B. Wang, L. Duan, S. Li, Y. Chen, and H. Li. Lotteryfl: Personalized and communication-efficient federated learning with lottery ticket hypothesis on non-iid datasets. arXiv preprint arXiv:2008.03371 , 2020.

A. Li, J. Sun, X. Zeng, M. Zhang, H. Li, and Y. Chen. Fedmask: Joint computation and communication-efficient personalized federated learning via heterogeneous masking. In Proceedings of the 19th ACM Conference on Embedded Networked Sensor Systems , pages 42-55, 2021.

M. Li, D. G. Andersen, J. W. Park, A. J. Smola, A. Ahmed, V. Josifovski, J. Long, E. J. Shekita, and B.-Y. Su. Scaling distributed machine learning with the parameter server. In 11th USENIX Symposium on operating systems design and implementation (OSDI 14) , pages 583-598, 2014.

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.

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.

C. Liu, L. Zhu, and M. Belkin. On the linearity of large non-linear models: when and why the tangent kernel is constant. In Advances in Neural Information Processing Systems , 2020.

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:1595415964, 2020.

Y. Liu, Y. Zhao, G. Zhou, and K. Xu. Fedprune: Personalized and communication-efficient federated learning on non-iid data. In International Conference on Neural Information Processing , pages 430-437. Springer, 2021.

M. W. Mahoney. Randomized algorithms for matrices and data. Foundations and Trends® in Machine Learning , 3(2):123-224, 2011. ISSN 1935-8237. doi: 10.1561/2200000035. URL http://dx.doi.org/10.1561/2200000035 .

M. W. Mahoney et al. Randomized algorithms for matrices and data. Foundations and Trends® in Machine Learning , 3(2):123-224, 2011.

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.

A. Mohtashami, M. Jaggi, and S. Stich. Masked training of neural networks with partial gradients. In International Conference on Artificial Intelligence and Statistics , pages 58765890. PMLR, 2022.

S. Negahban and M. J. Wainwright. Restricted strong convexity and weighted matrix completion: Optimal bounds with noise. The Journal of Machine Learning Research , 13(1):1665-1697, 2012.

S. Negahban, B. Yu, M. J. Wainwright, and P. Ravikumar. A unified framework for highdimensional analysis of m -estimators with decomposable regularizers. Advances in neural information processing systems , 22, 2009.

Q. Nguyen. On the proof of global convergence of gradient descent for deep relu networks with linear widths, 2021.

V. Papyan. Measurements of three-level hierarchical structure in the outliers in the spectrum of deepnet hessians, 2019.

V. Papyan. The full spectrum of deepnet hessians at scale: Dynamics with sgd training and sample size, 2019.

B. Polyak. Gradient methods for the minimisation of functionals. USSR Computational Mathematics and Mathematical Physics , 3(4):864-878, 1963. ISSN 0041-5553. doi: https://doi. org/10.1016/0041-5553(63)90382-3. URL https://www.sciencedirect.com/science/ article/pii/0041555363903823 .

D. Povey, X. Zhang, and S. Khudanpur. Parallel training of deep neural networks with natural gradient and parameter averaging. In ICLR (Workshop) , 2015.

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 International Conference on Machine Learning , pages 8253-8265. PMLR, 2020.

L. Sagun, L. Bottou, and Y. LeCun. Eigenvalues of the hessian in deep learning: Singularity and beyond. arXiv preprint arXiv:1611.07476 , 2016.

L. Sagun, U. Evci, V. U. Guney, Y. Dauphin, and L. Bottou. Empirical analysis of the hessian of over-parametrized neural networks, 2018.

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 , pages 32365-32417. PMLR, 2023.

S. U. Stich, J.-B. Cordonnier, and M. Jaggi. Sparsified sgd with memory. Advances in Neural Information Processing Systems , 31, 2018.

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.

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-1485, 2017.

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-34708. PMLR, 2023.

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.

S. Vargaftik, R. B. Basat, A. Portnoy, G. Mendelson, Y. B. Itzhak, and M. Mitzenmacher. Eden: Communication-efficient and robust distributed mean estimation for federated learning. In International Conference on Machine Learning , pages 21984-22014. PMLR, 2022.

R. Vershynin. Introduction to the non-asymptotic analysis of random matrices , page 210-268. Cambridge University Press, 2012.

R. Vershynin. High-Dimensional Probability: An Introduction with Applications in Data Science . Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 2018.

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.

M. J. Wainwright. High-dimensional statistics: A non-asymptotic viewpoint , volume 48. Cambridge university press, 2019.

H. Wang, S. Sievert, S. Liu, Z. Charles, D. Papailiopoulos, and S. Wright. Atomo: Communication-efficient learning via atomic sparsification. Advances in neural information processing systems , 31, 2018.

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.

D. P. Woodruff. Computational advertising: Techniques for targeting relevant ads. Foundations and Trends® in Theoretical Computer Science , 10(1-2):1-157, 2014. ISSN 1551-3068. doi: 10.1561/0400000060. URL http://dx.doi.org/10.1561/0400000060 .

Z. Xie, Q.-Y. Tang, Y. Cai, M. Sun, and P. Li. On the power-law hessian spectrums in deep learning, 2022.

C. Xing, D. Arpit, C. Tsirigotis, and Y. Bengio. A walk with sgd. arXiv preprint arXiv:1802.08770 , 2018.

X. Zhang, Z. Bu, Z. S. Wu, and M. Hong. Differentially private sgd without clipping bias: An error-feedback approach, 2024.

A Sketching Guarantee

Lemma A.1. For a randomised matrix R ∈ R b × p with i.i.d. subgaussian rows R i , i.e. ∥ R i ∥ ψ 2 ≤ 1 / √ b and E [ R i R T i ] = I p × p and for U, V ⊆ R p , such that b = Ω( log 3 ( p | U || V | /δ ) ε 2 ) with probability atleast 1 -δ , we have:

Proof. We use the following lemma from Song et al. [62]:

Lemma A.2 (Lemma D.24 from [62]) . Let R ∈ R b × p denote a random Gaussian matrix . Then for any fixed vector h ∈ R p and any fixed vector g ∈ R p , the following property holds:

Following a similar proof as Song et al. [62], we can get a sub-Gaussian version of this lemma:

Lemma A.3. Let R ∈ R b × p be a random sub-Gaussian matrix, with ψ 2 norm of each entry bounded by 1 √ b . Then we have:

̸

Proof. Note for i = j , R ∗ ,i , R ∗ ,j are two independent sub-Gaussian vectors. Let z k = R k,i R k,j and z = ⟨ R ∗ ,i , R ∗ ,j ⟩ . Then according to the definition of sub-Gaussian random variables, z k ∈ SE ( c 2 1 b 2 , c 1 b ) is a sub-exponential random variable with an absolute constant c 1 . Thus, we have z = ∑ b k =1 z k ∈ SE ( c 2 1 b , c 1 b ) , by sub-exponential concentration Lemma B.7 in Song et al. [62] we have:

Picking t = c 3 √ log ( p 2 /δ ) /b , we have:

̸

Taking the union bound over all ( i, j ) ∈ [ p ] × [ p ] and i = j , we complete the proof.

Then following the same proof with the only difference to apply Lemma A.3 instead of Lemma D.18 in Song et al. [62], we can get the following sub-Gaussian version of Lemma D.24 in Song et al. [62].

Lemma A.4. Let R ∈ R b × p denote a random sub-Gaussian matrix, with ψ 2 norm of each row bounded by 1 √ b . Then for any fixed vector h ∈ R p and any fixed vector g ∈ R p , the following properties hold:

Based on Lemma A.4, taking union bound over u ∈ U, v ∈ V and setting b = Ω ( C log 3 ( p | U || V | /δ ) ε 2 ) , we get the claim.

̸

B Dimension dependence

We make use of the following lemma from Banerjee et al. [5] to show dimensional dependence in ∥ R ⊤ Rg ∥ 2 2 :

Lemma B.1 (Theorem 5 in Banerjee et al. [5]) . Let A ⊆ S p -1 and let X ∈ R n × p be a design matrix with independent isotropic sub-Gaussian rows, i.e., ∥ X i ∥ ψ 2 ≤ κ and E [ X i X T i ] = I p × p . Then, for absolute constants η, c > 0 , with probability at least 1 -2 exp( -ηw 2 ( A )) , we have:

where w ( A ) denotes the Gaussian width of any set A given by : w ( A ) = E g [sup a ∈ A ⟨ a, g ⟩ ] , where g is an isotropic Gaussian random vector, i.e., g ∼ N (0 , I p × p ) .

Remark B.1. By choosing X to be a 1 √ n sub-Gaussian matrix, we get the form :

which we will use next.

Lemma B.2 (Dimension Dependence in sub-Gaussian sketching) . Fix a vector g ∈ R p . For a sketching matrix R defined earlier, we have for some constants η 1 , η 2 with probability at least 1 -2 exp( -η 1 ϵ 2 b ) -2 exp( -η 2 ϵ 2 p ) ,

Proof. Let G ⊆ R p . Gaussian width of G , w ( A ) = O ( √ log |G| ) . For g ∈ G using Lemma B.1,

Similarly, we can define B = { h : h = Rg ∀ g ∈ G} , Gaussian width w ( B ) = O ( ∥ R ∥ √ log |G| ) . Now setting the design matrix X for Lemma B.1 to be √ b/p · R ⊤ :

Taking ϵ = max { c 1 w ( A ) √ b , c 2 w ( B ) √ p } ,

Succintly, we can write with probability 1 -2 exp( -η 1 ϵ 2 b ) -2 exp( -η 2 ϵ 2 p ) ,

C Background

We state some useful results from Banerjee et al. [6]:

Lemma C.1 (Predictor Hessian Spectral Norm Bound, Theorem 4.1 and Lemma 4.1 in Banerjee et al. [6] ) . Under Assumptions 3.1 and 3.2, for any x ∈ X , θ ∈ B Spec ρ,ρ 1 ( θ 0 ) , with probability at least (1 -2( L +1) m ) , we have

where,

LemmaC.2 ( Local Smoothness , Theorem 5.2 in Banerjee et al. [6]) . Under Assumptions 3.2 and 3.3, with probability at least (1 -2( L +1) m ) , ∀ θ, θ ′ ∈ B Frob ρ,ρ 1 ( θ 0 ) ,

D Analysis of Convergence

We will bound T 1 , T 2 and T 3 , which we introduce below:

where ℓ ′ i , ℓ ′′ i are the first and second order derivatives of the point-wise losses. For an iterate θ t -1 , by the second-order Taylor expansion of the empirical loss around θ t -1 we have (ignoring h.o.t.),

Using 29,

where, H i,c = H i,c ( θ t -1 ) for brevity.

D.1 Single-step scheme (K=1)

Theorem D.1. Under assumptions 4.1 and 4.2, set K = 1 in Algorithm 1. For a suitable constant ε < 1 , η = η global η local = 2(1 -ε ) ((1 + ε ) 2 c s ϱ 2 + c H c l / √ m (1 + κ (2 ε + ε 2 ))) and b = Ω ( 1 ε 2 log 3 ( pT/δ ) ) with probability at least 1 -δ , we have:

where θ ∗ is the minimizer of the problem 1.

Proof.

where η = η global η local for notation simplicity.

D.1.1 Bounding T 1

where, a follows from Lemma A.1 for some for suitable choice of ε 1 .

D.1.2 Bounding T 2

Thus,

for some for suitable choice of ε 2 .

D.1.3 Bounding T 3

Using eigen-decomposition of H i,c ,

where Λ j,i,c , v j,i,c for j ∈ [1 , · · · , p ] are the eigen-values and eigen-vectors of H i,c respectively. Now,

where, Λ max ,i,c = ∥ H i,c ∥ 2 = | ℓ ′ i,c | · ∥∇ 2 f ∥ 2 = c l c H √ m .

Thus,

for some for suitable choice of ε 3 .

D.1.4 Combining T 1 , T 2 , T 3

Combining T 1 , T 2 , T 3 , we get:

Using PL condition,

where,

Setting η < 2(1 -ε 1 ) (1 + ε 2 ) 2 c s ϱ 2 + c H c l √ m (1 + κ (2 ε 3 + ε 2 3 )) gives a contraction map. If we iterate this recursion, we get:

D.1.5 Choice of sketching dimension b :

Using A.1, for each step t ∈ [0 , · · · T ] , union bounding over T 1 , T 2 and T 3 and over all T time-steps, Union bounding over all T time-steps, we get

where N = ∑ C i =1 n c is the number of training samples.

D.2 Multi-step scheme

Theorem D.2. Let ∥∇ θ ℓ ( θ ) ∥ ≤ G . Under Assumption 4.1, for a suitable constant ε < 1 , η = η global η local < 1 2 µK (1 -ε ) and b = Ω ( 1 ε 2 log 3 ( p 2 NT | C || K | /δ ) ) , with probability at least 1 -δ , we have:

where θ ∗ is the minimizer of the problem 1. And,

where, ϱ and c H are defined in Lemma C.1.

Set we get

Proof. We assume that the gradients at each time step t are bounded i.e ∥∇ θ L ( θ ) ∥ 2 ≤ G for some suitable constant G.

First, we introduce some notation: In the K-step case since the parameters are shared only in the sync step, we define θ c,t,k ∈ R p as the local parameters for the client c at time step t and step k ∈ [1 , 2 , · · · , K ] . Based on the sync step in Alg: 1, we can see that θ c,t, 0 is the same for all clients c ∈ C and thus, we denote it as θ t , i,e. θ c,t, 0 = θ t ∀ c ∈ C . Thus, we can write the update in the sync step as:

where,

Similar to Section D.1, we can bound T 1 , T 2 and T 3 in 31 as:

D.2.1 Bounding T 1

where,

D.2.2 Bounding T 2

for some suitable choice of ε 2 and b which we set later. Thus,

D.2.3 Bounding T 3

Using eigen-decomposition,

where Λ j,i,c , v j,i,c for j ∈ [1 , · · · , p ] are the eigen-values and eigen-vectors of H i,c = ℓ ′ i ∇ 2 θ f ( θ t -1 ; x i ) respectively.

for suitable choice of ε 4 and b .

D.2.4 Combining T 1 , T 2 , T 3 :

Combining T 1 , T 2 and T 3 , we get,

Using PL condition and iterating over this recursion we have:

where η ≤ 1 2 µK (1 -ε ) . For simplicity, let η global = 1 i.e. η = η local and ε 1 = ε 2 = ε 3 = ε . And,

D.2.5 Choice of sketching dimension b :

Using A.1, for each step t ∈ [0 , · · · T ] , union bounding over T 1 , T 2 and T 3 , Union bounding over all T time-steps, we get

where N = ∑ C i =1 n c is the number of training samples.

E Communication Efficiency

Theorem E.1. If Theorem 4.2 holds, then with ˜ O ( C max { 1 , C 2 ( ε,m,κ ) G 2 2 µ 2 (1 -ε ) Kϵ } log ( L ( θ 0 ) -L ( θ ∗ ) ϵ )) bits of communication, with w.p at least 1 -δ Algorithm: 1 outputs an ϵ -optimal solution θ T satisfying:

Proof. Setting η = min { 1 2 µK (1 -ε ) , ϵµ C 2 ( ε,m,κ ) KG 2 } and T = 1 2 µηK (1 -ε ) log ( 2( L ( θ 0 ) -L ( θ ∗ )) ϵ ) , we have:

The total number of communication bits are given as:

where, ˜ O hides poly-log dependence.

F Additional Experimental Details & Results

We provide the overview of the local Top-r baseline in Algorithm 2.

Algorithm 2 Local Top-r -based Distributed Learning.

Hyperparameters: server learning rate η global, local learning rate η local .

Inputs: local datasets D c of size n c for clients c = 1 , . . . , C , number of communication rounds T . Output: final model θ T .

1: Broadcast a random SEED to the clients.

2: for t = 1 , . . . , T do

3: On Client Nodes:

4: for c = 1 , . . . , C do

5: if t = 1 then

6: Receive the random SEED from the server. Initialize the local model θ c, 1 ∈ R p using the random SEED .

7: else

8: Receive ¯ ˆ ∆ t -1 from the server.

9: Update the model parameters θ t ← θ t -1 + ¯ ˆ ∆ t -1 .

10:

Assign the local model's parameters

11: end if

12: for k = 1 , . . . , K do

13: θ c,t ← θ c,t -η local · ∇ θ L c ( θ c,t )

14: end for

15: ∆ c,t ← θ c,t -θ t

16: Send Top-r sparsified updates ˆ ∆ c,t ← Top-r (∆ c,t ) to the server.

17: end for

18:

19: On the Server Node:

21: Aggregate them ¯ ˆ ∆ t ← η global · 1 C ∑ C c =1 ˆ ∆ c,t

20: Receive Top-r sparsified updates ˆ ∆ c,t from clients c = 1 , . . . , C .

22: Broadcast ¯ ˆ ∆ t to the clients.

23: end for

G Additional Results on Spectral Density of Predictor Hessian

10000

8000

6000

<Y

4000

2000

250

500

750

1000

1250

1500

1750

2000

Epoch

250

500

750

1000

1250

Epoch

This section presents empirical results verifying our assumption on the sum of absolute eigenvalues of the predictor Hessian (Assumption 4.2). Since it is infeasible to compute all eigenvalues for deep models like ResNet-18, we rely on numerical approximations introduced in several prior works [24, 78]. SLQ is a technique in numerical linear algebra to estimate the ESD of large matrices. The complete algorithm can be found in Algorithm 3. We use PyHessian [24], an open-source implementation of SLQ for our experiments. We keep the default parameters n v = 10 , m = 100 for the plots as they have been shown to be of high accuracy. We train a ResNet-18 from scratch model

11000

10000

9000

8000

7000

6000

5000

4000

3000

0

θ

c,t

θ

t

to be updated locally.

K

1500

1750

2000

(a) Epoch 500

(b) Epoch 1000

(c) Epoch 1500

Epoch:500

Epoch:1000

Epoch:1500

100

100

100

Scale)

10-1

Scale)

Scale)

10-2

10-2

10-2

(Log

10-3

Density(Log

Density (Log

Density

10-4

10-4

10-4

10-5

10-6

10-6

10-6

10-7

-300

-200-100

0

100

200

300

-100

-50

0

50

100

-300-200-100

0

100

200

300

Eigenvalue

Eigenvalue

Eigenvalue

Algorithm 3 Stochastic Lanczos Quadrature for ESD Computation (Ghorbani et al. [24])

1: Input: Parameter: θ , degree m , and n v .

2: Compute the gradient of θ by backpropagation, i.e., compute g θ = df ( θ ; x i ) dθ .

3: for i = 1 , 2 , . . . , n v do

4: {Different Seeds}

5: Draw a random vector v from N (0 , 1) and normalize it (same dimension as θ ).

6: Get the tridiagonal matrix T through Lanczos algorithm.

7: Compute τ ( i ) k and ˜ λ ( i ) k from T .

9: end for

8: ϕ z i σ = ∑ q k =1 τ k f ( ˜ λ k ; t, σ )

10: Return ϕ ( t ) = 1 n v ∑ n v l =1 ( ∑ q i =1 τ ( l ) i f ( ˜ λ ( l ) i ; t, σ ) )

to classify 1000 samples from two classes of CIFAR-10. We use a learning rate of 1 e -3 , SGD as the optimizer and and perform GD. We use a fix training sample to plot the ESD of predictor Hessian H ( θ ; x i ) = ℓ ′ i ∇ 2 θ f ( θ ; x i ) 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 3, we note that a bulk of eigenvalues are close to zero, and the spectral density decays quickly far from zero, i.e., most of the eigenvalues are tiny. Given the approximate density as a normalized histogram, we compute the empirical average and estimate ˆ κ an approximation to κ = ∑ | Λ i | / Λ max . Our results for ˆ κ are presented in Figure 2.

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 are 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 in 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.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of 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 in 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 on 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 a few datasets or with a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs 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 to 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 be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide experimental details in Section 5 and Appendix F.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a 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 a 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 a 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 it 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 to 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 a 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 submissions 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 to reproduce that algorithm.

(b) If the contribution is primarily a new 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 a way to access this model for reproducing the results or a 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 to describe 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., to 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [No]

Justification: Due to confidentiality constraints, we are unable to share the code at this time, but we provide sufficient details to reproduce the results.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so '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 to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) 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. If only a 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) is 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, hyperparameters, 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 not include experiments.

The experimental setting should be presented in the core of the paper to a 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, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: Computing spectral density for predictor Hessian is computationally demanding. Therefore, we repeated the experiments only once with a 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 are accompanied by error bars, confidence 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.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or 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 a 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 error 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

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide details in Appendix F.

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 as well as 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 it 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 https://neurips.cc/public/EthicsGuidelines ?

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 a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in 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: We provide a discussion on broader impact in Section 6.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or 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.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a 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 is 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 to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: 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)?

Answer: [NA]

Justification: [NA]

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a 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 a 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 1 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 or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.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 package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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.

If this information is not available online, the authors are encouraged to reach out to the asset'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/model 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 as details about compensation (if any)?

Answer: [NA]

Justification: [NA]

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, 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 not involve crowdsourcing nor 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 NeurIPS 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\nThere 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.\n\nQUESTIONS:\nA 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.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-20"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper provided ambient dimension-independent convergence and communication complexity analysis for sketching-based gradient descent algorithms using the second-order geometry of loss Hession.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\n1. 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].\n\nLIMITATIONS:\nPlease refer to Weaknesses and Questions.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}