Sketched Adaptive Distributed Deep Learning: A Sharp Convergence Analysis
Authors: Zhijie Chen
Qiaobo Li
Arindam Banerjee
Venue: NeurIPS 2025 poster ()
Keywords: Optimization
Distributed Learning
Communication Efficiency
Primary area: optimization
TL;DR: We provide sharp convergence analysis on sketched adaptive distributed deep learning that only depends on the intrinsic dimension of loss Hessian, instead of the full dimensionality..

Abstract: Combining gradient compression with adaptive optimizers is a highly desirable goal in distributed learning, with potential benefits in both fewer communication rounds and less per-round communication. In spite of preliminary empirical promise, certain major challenges in the convergence analysis of such methods have stayed open: handling compression based approximation of both first and second moments (pre-conditioner) which appear as a ratio; avoiding dependence on the number of parameters, which is extremely large in modern deep models; and providing high-probability guarantees instead of in-expectation, which can hide high variance behavior.

In this work, we introduce a family of Sketched Adaptive Distributed Learning (SADL) algorithms which can use suitable unbiased gradient sketching for compression with suitable adaptive optimization algorithms. As our main contribution, we provide theoretical convergence guarantees of SADL algorithms which addresses all of the existing challenges. In particular, our guarantees hold with high probability, picks up only a logarithmic dependence on the number of parameters, and the first and second moment approximation is handled precisely yielding a dependence on the intrinsic dimension of the loss Hessian, which is significantly smaller than the full dimensionality of deep learning models.  Empirically, the SADL algorithms are shown to be competitive with and often outperform baselines on both vision and language tasks, in both supervised fine-tuning and training-from-scratch regimes. Further, the SADL algorithms are also competitive with the state-of-the-art communication-efficient distributed learning algorithms based on error feedback.

Sketched Adaptive Distributed Deep Learning: A Sharp Convergence Analysis
Zhijie Chen
Qiaobo Li
Siebel School of Computing and Data Science University of Illinois Urbana- Chaamaaiaa Jucmon@illinoie illiiiiin.eeeeee edu
Siebel School of Computing and Data Science University Illinois Urbana- Champaign qiaobol20:
Arindam Banerjee
Siebel School Computing and Data Science
University of Illinois Oabana-Champaigee
arindamboilling edu
Abstract
Combining gradient compression with adaptive optimizers is a highly desirable goal in distributed learning. with potential benefits in both fewer communica- tion rounds and less per-round communication. In spite of preliminary empirical promise, certain major challenges in the convergence analysis of such methods have stayed open: handling compression based approximation of both first and second moments (pre-conditioner) which appear a ratio; avoiding dependence on the number of parameters, which is extremely large modern deep models; and providing high-probability guarantees instead of in-expectation, which can hide high variance behavior. this work, we introduce family Sketched Adaptive suitable unbiased gradient
Distributed Learning (SADL) algorithms which sketching for compression with suitable adaptive optimization algorithms. As our main contribution, provide theoretical convergence guarantees of SADL algo- rithms which addresses all of the existing challenges. In particular, our guarantees hold with high probability, picks up only logarithmic dependence of parameters, and the first and second moment approximation is handled precisely yielding dependence on the intrinsic dimension of the loss Hessian, which is significantly smaller than the full dimensionality of deep learning models. Empiri- cally, the SADL algorithms baselines on both vision and language tasks, in both supervised fine-tuning and training- teaning-from--ee from-scratch regimes. Further, the SADL algorithms are also competitive with the state-of-the-as based on error feedback.
the number
shown be competitive with and often outperform
coaimmniiattio--fffciii distributed learning algorithms
Introduction
Despite the recent progress in distributed deep learning (Liu et al. 2022). the cost of communication arguably remains the main challenge. [Wag et al. (2023) showed that 20 Ghps network bandwidth necessary bring the communication overhead to a suitable scale for finetuning GPT-J-6B, which is unrealistic in distributed settings. Even with good network conditions, reduction of communication complexity means one can train much larger models given the same communication budget.
The communication cost of vanilla distributed learning can be represented as OddT), where d the ambient dimension of the parameter space, i.e., the number of parameters, and T is the number of communication rounds for convergence. Various methods have been proposed to minimize T, e.g..
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

local training (Stich, 2018), large batch training (Xu et al. 2023), etc. Folklores in centralized training regimes suggest that T heavily relies on the choice of optimizers, where adaptive methods (such as Adam, AMSGrad, and variants) usually demonstrate faster convergence and better generalization performance, especially in transformer-basee machine learning models (Reddi et al. 2019).
Alternatively, communication costs can be reduced by being more thrifty on the communication bits in each round, reducing the ((d) factor to ((b), b < d. In modern deep learning models, the (dd) term is the dominant factor in the communication complexity since d >> 'I' 'n modern models. Considerable efforts have been devoted to design efficient gradient compression methods, which
compress (gradient) vector of dimension d to an effective size b. Popular gradient compression
methods include quantization (Alistarh al., 2017: Chen et al.. [023 Reisizadeh et al. [2020] Liu
et al., 2023a), sparsification (Alistarh et al. 2018: Wu et al., 2018 Rothchild et al.
20200 and
sketching (Spring et al.] 2019) Jiang et al. 2024 Song et al.] 2023).
While progress has been made along both of these lines (see Section[A] there are certain truly unique challenges in provably doing both gradient compression (for communication efficiency) and adaptive optimization (for faster convergence) simultaneously First, adaptive methods (such as Adam (Kingma and Ba, 2014) and AMSGrad (Reddi al., (019)) work with ratio of the first moment and second moment, and gradient compression will approximate both the numerator and denominator of such a ratio, making both numerical and statistical analysis of the ratio quite complex. Note that this major challenge is entirely non non-existent in non- non-adaptive methods, which form bulk of the literature in distributed and federated learning (Rothchild et al. 2020; Shrivastava et al. 2024). Second, there are two primary sources of noise in any such analyses: due to mini-batching in client updates and due to stochasticity in gradient compression; and a third secondary but more complex source of noise due to sequential dependencies in adaptive updates over iterations. While an "in- expectation" analysis can mute some of these noises (Gorbunov et al. (02I), the algorithm may have high variance which will lead to unstable training in practice. To avoid such issues, our focus is to provide high--roobbbility convergence bounds, while taking into account all three of the above sources of noise. Third, almost all existing analysis based on gradient compression picks up an undesirable ambient dimension d dependence in the optimization analysis. Biased compressors generally need to handle errors from first-order terms because of the distortion in the gradient direction whereas unbiased compressors typically have high variance which affects the second-order terms. Since modern deep learning models have ambient dimension in hundreds of billions even trillions, it is important to avoid such a dependence. The existing analyses on adaptive methods with gradient compression are quite alarming, which show that the iterations T needed for convergence can be inversely proportional to the compression rate (Chen et al. 2022; Song et al. (023). For constant per-round communication bits, the bounds indicate the iteration complexity T to scale Od), i.e., linearly with the ambient dimensionality, which is prohibitively large for modern deep learning models. Appendix A jincuudes more challenges in the algorithmic and theoretical aspects.
In this work, we introduce a family of Sketched Adaptive Distributed Learning (SADL) algorithms which have the flexibility of using suitable unbiased gradient sketching compressors with suitable adaptive optimization algorithms. At a high level, SADL algorithms are analogous to previous attempts (Tang et al.) 202T] Chen et al., 2022 Wang et al. 2022a), which showed preliminary empirical success of applying gradient compression with adaptive optimizers, and (Song et al.] 2023) which introduced the -esseech framework distributed learning, but comes with a sharp convergence analysis which addresses all three of the unique challenges.
A key technical contribution of our current work is that the convergence rate of the proposed SADL algorithms depends only logarithmically (instead of linearly) on the ambient dimension d. The central technical challenge in addressing the dimensional dependence is to handle the entry-wise sketching noise in both the the first moments and the second moments (preconditioners) of the adaptive optimizers, which has been acknowledged to be non-trivial (Tang al., 2021 Wang et al., 2022a). Our sharper analysis yields bounds based on the intrinsic dimension (instead of the ambient dimension d) of the loss Hessian in deep learning, i.e., the ratio of sum of absolute eigenvalues over the largest eigenvalue. Recent observations on the Hessian spectrum of deep learning models have demonstrated that the intrinsic dimension is significantly smaller than the ambient dimension, by showing the eigenvalues decay sharply, with most eigenvalues being close to zero (Li et al., 2020; Liao and Mahoney (20TT] Liu et al. 2023b), and even arguably conforming with pawer-law decay (Xie et al. 2022) Zhang et al.) 2024). This specific eigenspectrum structure provides significant advantages in the sharp analysis of sketching noise in adaptive methods. The SADL algorithms

---

Algorithm Sketched Adaptive Distributed Deep Learning (SADL)
Input:
Learning rate 7), initial parameters To, adaptive optimizer OPP
Output: Updated parameters ZT
Initialize server moments: my 0, vo 0. ög O, client initial parameters: z6.0 20, client
moments: m 0, v/ 0, 66 0, vcc [C];
for -= ,, do
Client Updates:
for c c= 1,2..... C do
Client model synchronization: 2i,0, mb, of, of ADA_OPT(sf mi- of-1, 08-1, me)
for k= 1, ..,,.... K do
Compute stochastic gradient -i.k1 with respect the parameters ii,k-ii
Perform gradient step: zf,5 Fi, x-
end for
Sketch (compress) the parameter updates and send m, to server: m, = =k(rfoo xf,x))
end for
Server Updates:
Average sketched client updates and send my back to clients: my 2 EE= mi;
Update parameters and moments: 25, my, Ut, de = ADA_0PT((==1, m, m,-1, Vt- 0z- mc).
end for
do not involve computing the Hessian eigenspectrum, which is only used for the convergence analysis. Further, our results do not follow from the standard toolbox for sketching methods, such as (Kootufe and Sriperumbudur, (020)), which ensures nearby vectors stay close after sketching, but has no direct implications an optimization setting, especially with adaptive optimizers which will sketch both the first and second moment, and work with their ratio. Based on such analysis, our work has the following main contributions:
(I) We introduce SADL algorithms which combine random sketching and adaptive methods, While the preconditoner in adaptive methods morphs the shape of sketching noise, posing challenges in leveraging the fast-decaying Hessian eigenstructure, we prove that the proposed sketching effectively balances iteration complexity and sketching dimension b. We derive a high probability bound showing that sketch dimension of Ollog d) suffices to achieve an (II/VT) convergence rate depending only on the intrinsic dimension in non-convex deep learning settings.
(2) Unlike existing works (Reddi et al. 2020 Xie al., 2020), we provide a general convergence analysis without assuming uniformly bounded gradient norms on either the server or client side. We prove that SADL automatically generates bounded gradients along the entire optimization trajectory, where the involvement of mini-batch stochasticity and multiple local training steps calls for careful analysis on connecting the noisy local training steps with the global loss.
(3) We validate our thcoretical claims with empirical evidence on deep learning models from vision (ResNet, Vision Transformer) and language (BERT) tasks. We cover both fine-tuning and training- from-scratch regimes. SADL in general achieves comparable performance with adaptive methods without any gradient compression, using only 0.1% of the full dimensioanlity. Further, SADL is competitive with the state-of-the-ee algorithms based on error feedback and adaptive methods.
2 Sketched Adaptive Distributed Deep Learning
canonical distributed learning setting involves C' clients, each associated with local data dis- tribution D... The goal is to minimize the averaged empirical risk: C(x) = J/2 Σ Eę=D. Uxx E), where is the loss function, x E IRd is the parameter vector, and E is the data sample. We denote (z)x) lEg-D_ {), c E [C] the client loss computed over the local distribution. We denote gh,, as the mini-batch gradient over (x) at global step t and local step A.
adaptive methods, which calls adaptive optimizers as subroutines. We denote T as the total training rounds. At each round, after K local SGD steps, client e sends to the server the sketched local model updates with a sketching operator sk: d + g². If < d without deteriorating the performance too much, the communication
Algorithm I presents generic framework of communication-effici
3

---

Algorithm 2 OPT (AMSGrad)
Inputsiterate I;-1, moments m,-1, 1t-1, 6--1 sketched updates m,
Parameters:Leaaming rate Hi 31, 32, Small constant c, Sketch size
Output:Updated parameters 2t, and moments me, ta. 0,
Update my = 31 14-1 + (1 - 31) desk((m))) (14)):
Update vi = 32 11 + (1 - B2) desk (b/20)/)
Update ô, = max(0;;111 (1
Update 2(+1 2t Votte my Zt RV,," 1/2 m..
cost per round can be reduced from O(d) to O(6). Algorithm 2 projects the compressed updates and second moments back to the ambient dimension using desketching operator desk: R² - R and implements a single-step adaptive optimization. The server and clients call Algorithm 2 at every epoch, i.e. communication round, to update the global model and synchronize local models. Notice that in Algorithm 1 both the clients and the server transmit sketched vectors which results in an overall (bTT) communication cost.
Remark 1. (Sketching Randomness) At each single round, the sketching operators sk's are shared among clients, via the same random seed. On the other hand, fresh sk's at different rounds to get statistical independence which helps the analysis.
We list some desired properties for sk, desk and discuss sketching methods which satisfy them. Property 1. (Linearity). The compression operators are linear w.Ft the input vectors, i.e. sk n i= v2) = -. sk(1;) and desk().. Th 0: = = desk(T:), V(u,, 0; E R
Property 2. (Unbiased Estimation). For any vector e E IR ((desk(sk(e)) =
Property 3. (Bounded Vector Products). For any fixed vector v, E TRd, ((((desk(sk(tt),))
(v, 1) 2 (d/8)) |||]|||Aa||> S (()).
) vb
We denote R E 3R as the sketching operator, and sk(c) Re and (T) = R Different instantiations of R constitute rich family of sketching operators, including i.i.d. isotropic Gaus- sian (Song et al. 2023), Subsampled Randomized Hadamard Transform (SRHT) (Lu et al. 2013).
and Count-Sketch (Charikar al., 2002), among others. The specific bounds for these special cases
can be found in Appendices B.1
B.2. and B.3 respectively.
2.1 Convergence Analysis
We first state set of standard assumptions commonly used in the literature of first-order stochastic methods. We will use || - || to denote L2-norm throughout the work.
Assumption 1. (Bounded Global Gradients). Square norm of the gradient is uniformly bounded, i.e., [|V[(=)]| < Gz
Assumption 2. (Bounded Client Gradients). For every client, there exists constant Ge 2 0, such that |VLL (z))l < GZ e C [C].
For simplicity, in this section we define G = max max ( Gg). In Section we show that Assumptions I and are not necessary to derive the convergence bound we present. We assume the local stochastic mini-batch noise is aub-Gaussian, which is widely adopted in first-order optimization (Harvey et al. 22199] Mou et al. 2020).
Assumption 3. SSub-Gaussian Noise). The stochastic noise V (x) (x)))) at each client is nuu-Gaassiian random variable, iee. P(IIVL(<)) - (=))) 2 S 2exp( (12 / (o), for all t.
Besides, we present assumptions on the Hessian eigenspectrum v.). of the loss function L. Assumption 4. (Hessian Matrix Eigenspectrum) The smoothness of the client loss function Ci, i.e. the largest eigenvalue of the loss Hessian Hc, is bounded by L.
The local smoothness assumption commonly used in distributed learning settings (Safaryan et al. 2021 Fatkhullin et al. 2024) and holds for general deep learning losses. It can be directly derived from Assumption that the global loss L a DE Lo is smooth.

---

Definition 2.1. (Intrinsic Dimension) Let { be the eigenspectrun of the loss Hessian He. The intrinsic dimension is defined I = = ii Ael/ max, AAil-
The definition of intrinsic dimension is analogous to that in the classical literature (lpsen and Saibaba 2024). but we take the absolute values of eigenvalues. Intuitively, the Hessian matrix possesses an anisotropic structure in different directions, whereas the conventional smoothness is pessimistic estimation of the loss curvature. the intrinsic dimension of the Hessian in deep learning losses is significantly smaller than the ambient dimensionality d. (Ghorbani al., 2019 Li et al., 2020 Liu et al., 2023b) show the eigenspectrum enjoys sharp decay in magnitude. (Sagun et al. 2016 Liao and Mahoney| 2021 show the eigenspectrum have bulk parts concentrate zero. (Xie et al., 2022) Zhang et al. 2024
large volume of recent empirical literature has identified that
Liao
further show the eigenvalues satisfy power-law distribution, and in this case the intrinsic dimension a constant independent of d. We quote their plots Appendix EE] for completeness. Our empirical verification under the setting of distributed learning can also be found in Fig. 5 in Appendix E
Remark 2. (Three types of noises in Algorithm 7). One of the key technical contributions this work is to theoretically balance the noises from different sources. The noise in the training process stems from the local mini-batch training, the compression error due to sketching, and the aggregate noise over the training horizon. The mini-batch stochastic error is o-sub-Gaussian by Assumption 3 The sketching error depends on the specific choice of sketching methods, but is always subject to the bounded property on vector products (Property B We will denote the 'bad probability' from mini-batch and sketch respectively as 6g and 6, which are usually viewed as a tiny value (10 (10-5) in high probability bounds. These two kinds of noise are unbiased, but leads to sequential dependencies. In the analysis (Appendix(C) C we will show that the aggregated noise due to the sequential dependencies form a martingale, which will be used to derive gegh-probability concentration bounds. We denote as the scale of the 1/2-norm (Vershynin 2018) in the martingale.
L
Now we characterize the convergence of Algorithm in Theorem 2.2 All technical proofs for this section are in Appendix LoanndixCandd and we provide an outline of the proof techniques in Section 2.2
Theorem 2.2. (Informal version of Theorem C.1)
Suppose the sequence of iterates {T:)T_ is
generated by Algorithm 7 with constant learning rate Ti- 17. Under Assumptions 1-4, for any T and e 0, with probability 1 (8)).
kij TKC/J(b,d d) +E
Σ [V(())))1 sc(x)) + - TOT + J(b,d) + TTO +g)
+ I
where J(b,d) = "ICK""B) and ž1
vb
Remark 3. (Nature of High Probability) While our convergence result holds with high probability,
the informal version in Theorem 2.2 hides the dependencies on the three types of noise (Remark 2
The dependencies follow exponential concentrations, and are shown in Corollaries [ [[nd]] 2
Remark 4. (Rate of Convergence) The rate of convergence may be hard to read off from Theorem 2.2 At a high level, we show in Corollaries and [ that the rate is (1/TT) when T < 0(1/²²) and is
O(I/VT) when 2 2(1/22). The specific results provide additional details.
A non-asymptotic convergence bound of training with practical decaying learning rates can be found in Theorem C.4]in appendix. Given that we only introduce logarithmic factors on d in the iteration complexity and the per-round communication b is constant, the total communication bits of training a deep model till convergence is also logarithmic w.r.t d.
To better understand Theorem 2.2 can investigate different regimes based the training stages. For large T, where T > ((1/e)), we can achieve O(L/VT) convergence rate in Corollary Corollary 1. (Asymptotic Regime of Theorem 2.2g With the same condition as in Theorem 2.2 and
a constant learning rate = KVT for sufficiently large T 2 G² with probability 1 ()) O(exp(-S((v2)) - -5g
DIv VT
4 (1 (1+J))
4KIL G 20(z))
1-33)² +
LG 2 VT
2T +v (12022 + + logg 3,
))
KVT VT
where ó, 8g and v correspond to the three noise types (Remark 2 and J2 = log "CKK"(" 16)
vb
5

---

More interestingly, when T' is relatively small, i.c. T < ((1/e)), we can observe that the coefficient of || [x+)ll on the left hand side Theorem 2.2 and C.4 is approximately a constant, given that is relatively small. Therefore, SADI. can achieve an O (T/T) convergence for small T, which yields faster convergence rate than non-adaptive methods.
Corollary 2. Meear-mitializzttio Regime of Theorem 2.2 With the same condition as in Thereom 2.2 and a constant learning rate tp KVT set 2 log (CK&TT2//) and constant J3 > 22G, then for any T 2 La-y2e with probability 1 - (6)) O(exp((ee(x")))
6g.
J&T ΣIve(E)) S f(z))e
1 T
2T )) I.G v 10² log (3,
4xII+ G G2 (1-8811 c
RT
T
T
where ö, 8g and u correspond to the three noise types (Remark 2
2.2 Technical Details and Proof Sketch
In this section, we provide proof sketch behind the main results. We focus the proof of Theorem 2.2 and the proof of Theorem C.4]shares shares the main structure. The proof contains several critical components, which are unique adaptive methods. We adopt AMSGrad (Alg. 2 as the server optimizer and it would be straightforward to extend the analysis to other adaptive methods.
First, we introduce the descent lemma for AMSGrad. For conciseness, denote the precondtioner Bn matrix diag((Vét + <)²) as Vi. Define auxiliary variable 2; Iz 1-81 (xa 21-1). The trajectory of L over { )?=1 can be tracked by the following lemma. Lemma 2.3. (Informal version of Lemma C.2 For any step t [T'.
((2+1)) (f(2))
V[(a)) V. R. Regi + (z+ - x.) Hc(2))(z441 - ze),
C
where H H(())) is the loss Hessian at some 2g within the element-wise interval of (Ia, zi), and S omits
the less important terms.
Our objective henceforth is to bound the first-order descent term and the second-order quadratic term on the right hand side respectively.
Second-Order Quadratic Term. Denote {)y, 1)) as the eigen-pairs of Hz(2)). The quadratic term can be written as (z: - =x)" Hc(&))(<+11 2)) = ビ X, (z++1 2,, uy{ - xe, 22)) The inner product term is a projection of the updates onto anisotropic bases. Since 2++1 2t and Iy can be expressed by r++1 3e and Jy. ft-1, we can bound the quadratic term using the following lemma.
Lemma 2.4. For any & (T). [0 It-1, mz)l kn
log (CKtt/()))
KG with probability 8.
Jb
proof of generalized version of this statement deferred the appendix. Induction method is used to address the temporal dependence introduced by the momentum factor in AMSGrad. Combining Lemma 2.4 with Assumption 4 yields dimension-free bound on the second-order quadratic term.
Remark 5. A straightforward application of smoothness to the second-order term yields a quadratic
term R Rgl.? which is linearly proportional to d in scale (Rothchild al. 2020: Song al.
2023). We avoid this dimension dependence by combining Property of sketching and the intrinsic dimension of the deep learning Hessian matrix.
First-Order Descent Term. The first-order term in the descent lemma can be decomposed into three components, which we will handle separately:
V((x)) (V R. Regia v((x)) V/"/VVC (wx) VL(xx) V (RT Regis (cfxx)
D
v((x)) V (VL (z.)..
First, DE can be reduced to a second-order term by Taylor expansion on VL. Since this term does not involve any stochasticity from random sketching, we can directly upper bound Dj by Cauchy- Schwartz. Next, since á _ (x+) = VL(x.), Di can viewed as a scaled squared gradient norm. Applying wiement-wise high probability bound on random sketching yields lower bound.

---

= 10 Bylestpoc1
1,6 10P Byteveepcch
RACSA
90%
200
J20
104.7
kpack
Figure 1: Model performance on CIFAR-10 with ResNet of 42M parameters. The plot starts from the
10th epoch for better demonstration; First (Second): Validation errors under compression rate 0.1% (1%). Third:
Validation error on SADL with different communication costs. The legend 1.6 x 10 represents training in the ambient dimension without sketching. Fourth: Training error SADL with different communication costs. Higher compression rate improves the convergence rate in training and the peak validation error is achieved when the compression rate is 0.1%.
Lemma 2.5. For V.-1 generated by Algorithm (SADL), with probability 1 8.
VL(x)) V (((x) 2 M V
where M
Ing (CKKte))))
nKG+E
Vb
Martingale for zero-centered noise. Dí contains a zero-centered noise term R Reg (xf where the randomness is over R, and the mini-batch noise at round f. Although xi has temporal dependence, the fresh noise due mini-batching and sketching-desketchingg at round t is independent of the randomness in the previous iterations. Therefore, the random process defined by the aggregation of D2 forms martingale. The martingale difference can be bounded with high probability under proposed sketching method. By adapting Azuma's 's inequality on a Gau-Gausiinn martingale, we have
Lemma 2.6. With probability O(exp(-S((2)) - 5 6g.
Σ Σ C
VV V((x)) (1
Resi, (z) S SUKVT log (CKTd/0) G
+ log
27
)).
Vb
8g
Finally, applying union bounds to these parts and telescoping the descent lemma leads to Theorem 2.2
While the intrinsic dimension plays key role in avoiding the ambient dimension dependence in our convergence bound, it alone is insufficient to address all the technical challenges in sketched adaptive distributed learning. In fact, intrinsic dimension only addresses the second-orde term in optimization. The analysis of sufficient descent (Lemma [.3 and 2.5) and bounded aggregate sketching noise (Lemmaa2. 2.6) manages to address the noise, the norm of which is in-expectation linear to d, in the first-order term.
Remark 6.
3 Bounded Gradient Norm Along Optimization Trajectory
Although the gradient norm assumptions (Assumption and are standard in adaptive optimiza- tion (Reddi et al. 2020) and distributed learning research (asu et al., 2019 Xie et al. 20200), in this
section, we show that the two assumptions are not necessary to derive the convergence bound.
Our main idea is to show the gradient norm is bounded over the entire optimization path with high probability. We rely the following lemma to demonstrate the boundedness.
Lemma 3.1. For any L-smooth function C(x) with optimal value L 2 0, [ < 2LC(x).
As stated in Lemma 3.1 for any smooth function, the gradient norm can be bounded by the function value the specific iterate. That being said, we can derive upper bound the gradient norm along the optimization trajectory via bounding the function values over the iterates. However, the technical difficulty of the analysis lies in the involvement of the local training steps, which might be noisy and the relation which with the global iterate is unclear.
Our analysis can be divided into two steps: 1) We first relate the averaged local gradient norm to the global function value based on the local smoothness. Notice that this step does not require any additional assumptions, such the deviation between local and global function values; 2) We apply
7

---

3 xx10 Bytes/epoch
3.0 Bytes/epoch
x Byleslepoch
MARL
7%
Neud500 NAXIN
595
Veldstten Erei
2%
FastE50
NARUMA
Hat NA
Coodall
Coccal
67
(61 AA
2%
176
2%
1%
2%
1.56
10
15
20
25
30
35
40
10
1.5
20
25
30
35
10
1.5
25
40
Epoch
Epoch
Epoch
Figure 2: Validation Error on CIFAR-10. We finetune a ViT-base model (with 86M parameters). 1 bit- Adt-Aamm has comparable communication cost with 3 x 10th Bytes/epoch. SADL shows competitive performance under all communication budgets.
the induction method to show the global loss is contained in the neighborhood of the function value initialization, and the bound of the gradient norm follows immediately by applying Lemma 3.1
The following lemma shows how the local gradient norm can be related to the global loss at iterate Te, which will be a key component in the following analysis.
Lemma 3.2. Under Assumption Let rj 21/KK The local gradients as of K can be bounded
with probability 1 - CK3 CK exp(- 20pp-22///P/ /²/o²)bb by
C (x
22VI + 2 2 In + A
Next, we show the global loss contained in the neighhorhood of the function value at initialization.
With a slight abuse notation, we define a constant G == max /2 2 M} where the full form can be found in 77 in Appendix D and the following lemma holds. Lemma 3.3. We can derive an upper bound on ((=T) with high probability,
2CK de )2 (L)/VV++
[(2r) S KIVTMIG + wnI/TMMVV + Ms + Emm + xay MsG + MeVC
where [Mi]T are constants independent of K, and G, defined in Appendix D
With the closeness of žy and Ty, we can show C(xr) < 2r. under appropriate choice of к and and it is sufficient to ensure the gradient norm is bounded over the entire optimization path, which directly leads to the following convergence guarantee. The full proof can be found in Appendix D
./ subject to yo I min + log log* "CKKTd /0) 10))-1 ST and 2LVR
Theorem 3.4. Let
vb
K = va Then under Assumption and 4 with probability Texp( ((o(-(((2))- - TCK5,
TCK exp(-22/00) T5, the gradient on the iterates vt generated by Algorithm| 1 are bounded by G, i.e. [|V[(x))] < G. Consequently, the averaged gradient converges with rate O(1/VT) by
G
[V
LT (nKMs +6),
where Ms
log (CKPTT)/8 Vb
(v2G + 22(11 + 22 2)
4 Empirical Studies
In this section, we instantiate the algorithm framework of SADL to demonstrate the effect of sketching in common distributed deep learning settings.
Experimental Configurations. We adopt three experimental settings, from vision to language tasks. For the vision task, we train ResNer101 (Wu and He 2018) with total of 42M parameters from scratch and finetune a ViT-Base (Dosovitskiy al. 2020) with 86M parameters on CIFAR- 10 (Krizhevsky et al. 2009). For the language task, we adopt SST2, text classification task, from
8

---

0'010' Bytes/epoch
x 10th Bytes/epoch
10%
SADI
50%
5/31
30%
8
C2Adam
CDAccm
8 10
Emmr
25%
MARRIA
Woldatne Ereer
40%
MAR No
Wildaato Eiee
25%
Cockoa t5
20%
15%
20%
1.5
10%
10%
5
10
1.5
20
25
20
10
15
10
1.5
20
30
Epoch
Epech
Epoch
Figure 3: Validation Error on SST2 (GLUE) with BERT of 100M parameters. Left: compression rate 0.2%; Middle: 2%; Right: SADL with communication costs 18 x 10 8 x 10 10 x 10 x 10º} Bytes/epoch. The legend 4 x 10t represents training in the ambient dimension without sketching. Higher compression rate improves the convergence rate and all compression rates achieve comparable test errors at the end of training.
the GLUE benchmark (Wang et al. parameters. For all experiments we split the training dataset uniformly over 5 clients. Our baselines include FetchSGD (Rothchild et al. 2020), MARINA (Gorbunov et al. (021). CocktailSGD (Wang)
2018). We train a BERT model (Devlin 2018 with 100M
et al. 2023), CDAdam (Wang et al., 2022a), bit-Adam (Tang et al. 20. FedCAMS (Wang et al. 2022b) and FedPAQ (Reisizadeh et al. 2020). A comparison of the theoretical guarantees of the baselines can found in Table 12 in the Appendix. We define the compression rate as the ratio between the transmission size and the full model size (both in bytes). Specifically, for SADL, the compression rate is b/d, i.c. the ratio between sketch size and full dimensionality.
In each set of experiments, we compare SADL and baseline algorithms under compression rates of 0.1% and 1% respectively. For fair comparison, All algorithms, including SADL, use Adam as the server optimizer when applicable. SADL adopts SRHT as the sketching algorithm. Notice that most existing works only demonstrate the performance under up 1% compression rate, while our work further pushes the limit of compression to 0.1%. We will show SADL's highly competitive performance in this extreme case, where some baselines fall short.
Sharp-Decaying Hessian Eigenspectrum. Our theoretical result builds upon the notion of intrinsic dimension. While existing research has repeatedly shown supporting evidence on the sharp-decaying cigenspectrum, we also provide a verification in distributed deep learning in Fig. 5 in the Appendix.
We present the empirical performance of SADL under three different tasks and report the mean values across 3 independent runs, with sources of randomness including client data partitioning, data shuffling, and sketching. The statistical significance analysis and the hyper- parameter budgets can be found in Appendix E Fig. T shows the error curves on the validation of CIFAR-10 when training ResNet(40M) with communication budget in (1.6 10 1.6 x 100} bytes/epoch, representing 0.1% and 1% compression rate respectively. For SADL, the two compression rates are achieved with sketch size b (4 x 10a 4 10 The compression rate of Ibit-Adam and FedPAQ is 3%, and are compared with other basclines with compression rate 1%. We can observe that under both communication budgets, SADL outperforms other optimizers by significant margin. Further, we compare the model performance of SADL with different sketch sizes and find that in this experimental setting, the validation error is not monotonic with the sketch size and reaches the peak value when b = 4 x 10 (i.e. compression rate 0.1 ) On the other hand, the training error, which better reflects the convergence speed, is strictly monotonic with sketch sizes- larger sketch size leads to faster convergence and agrees with our theory. The discrepancy between the two rates suggests that sketching methods may implicitly improve the model generalization ability.
Convergence Results.
Similar phenomenon is observed in the language task. Fig- shows the validation errors of training SST2 with BERT (100M parameters). The compression rates are selected from (0.2%, 2%). Fig. 3 includes baseline methods which achieve comparable performance. The full plot including all baselines presented in Fig. in the appendix. We observe SADL converges faster and achieves slightly better validation performance across communication budgets. In Fig. 3 we further compare SADL with different compression rates from (0.02%, 0.2%, 2%) and distributed adaptive method without sketching. Remarkably, 1) SADL converges faster with higher compression rates; 2) Under compression rate 0.2%, SADL achieves almost the same performance as the original unsketched version; 3) Given that the ambient dimension is 100M, it is thrilling to see under an extremely low compression rate (0.02%). SADI. still achieves comparable performance as trained in the ambient dimension. We also present results from finetuning a ViT-Base model (80M parameters) in Fig. 2

---

Table 1: Ablation Study of SADL on the ResNet experiment with communication costs f1.6 x 10º, 1.6 x } Bytes/epoch. Validation errors of the last communication round are displayed. The comparison is across different combinations of server optimizers (Adam. AMSGrad) and sketching algorithm (SRHT, CountSketch). Adam SRHT is the recommended combination in practice.
1.6 x 10s Bytes/epoch
1.6 10 Bytes/epoch
SRHT
CountSketch
SRHT CountSketch
Adam
17.0%
23.8%
16.7%
23.6%
AMSGrad
18.1%
25.1%
17.0%
24.7%
The compression rates are selected from {0.1%, 1%, 10%). The full plot can be found in Fig. o in the appendix. We observe that SADL advantageous especially under lower compression rates. Remarkably, CocktailSGD achieves the best validation error under compression rates 1% and 10%, while falling short in 0.1%. SADL has a consistent performance across all communication budgets.
Performance under Extremely Low Communication. We experiment with extremely low com- be empirically grounded. The
pression rates to show the theoretical logarithmic dependence experiments are conducted under the same setting as Figure | and respectively. We adopt com- pression rates from 0.001 % 1% where the lowest amounts to b 400 in the ResNet experiment and 6 = 2000 in the SST-2 task. We present the validation errors along the training process in Figure 8 in Appendix E We observe that although the validation accuracies converge distinct values, SADL converges under all sketch sizes. More interestingly, different sketch sizes achieve similar convergence speed. Even under extremely tiny sketch sizes, SADL converges in the first 100 (25 resp.) epochs in CIFAR- 10 (SST2 resp.) task. This observation aligns with our theoretical results on the logarithmic dependence on d the convergence rate.
Ablation Studies. SADL naturally supports a custom choice of server optimizers and sketching algorithms in practice. In the previous experiments, choose Adam the server optimizer and SRHT as the sketching algorithm. Since the mainstream adaptive optimizers share similar designs in utilizing momentum and preconditioners, the theory developed can be easily extended to Adam. The theoretical guarantees can also be seamlessly adapted to other sketching algorithms. We extend our experimental results the combination of [Adam, AMSGrad] [SRHT, CountSketch]. Table presents the validation errors during the last communication round.
Regarding the choice of server optimizer, Adam achieves slightly better performance than AMSGrad. Regarding the choice of sketching methods, SRHT consistently outperforms CountSketch, while the latter is still competitive with other baseline methods in Fig. 1. This observation aligns with the error bound of CountSketch (Lemma B.3 which differs by 1 /b factor with SRHT (Lemma B.1
5 Conclusion
In this paper, we investigated sketched adaptive methods for distributed learning. While the motivation behind combining sketching and adaptive methods is clear, there is limited understanding on its empirical success. We show an exponential improvement in the communication cost from (d) to Ollog d), i.e., logarithmic in the number of parameters. The improvement especially important for modern deep learning models with large number of parameters. Our analysis introduces a novel technique that handles three sources of noise and adaptive optimization, both of which pose key challenges in applying sketching methods to communication-effise learning. We empirically show that the proposed SADI. consistently outperforms the existing baselines, meaning our method is practical, while having strong theoretical guarantees.
adaptive distributed deep
Acknowledgment
The work was supported by the National Science Foundation (NSF) through awards IIS 21-31335, OAC 21 31-38835, DBI 20-21898, well as a C3.ai research award. Compute support for the work was provided by the National Center for Supercomputing Applications (NCSA) and the Illinois Campus Cluster Program (ICCP).
10

---

References
Alistarh, D., Grubic, D., Li, J., Tomioka, R., and Vojnovic, M. (2017). Qsgd: Communication- efficient sgd via gradient quantization and encoding. Advances in neural information processing systems, 30.
Alistarh, D., Hoefler, T., Johansson, M., Konstantinov, N., Khirirat, S., and Renggli, C. (2018). The convergence of sparsified gradient methods. Advances Neural Information Processing Systems, 31.
Basu, D., Data, D., Karakus, C., and Diggavi, S. (2019). Qsparse-local-sgd:: Distributed sgd with quantization, sparsification and local computations. Advances in Neural Information Processing Systems, 32.
Beznosikov, A., Horvath, S., Richtarik, P., and Safaryan, M. (2023). On biased compression for distributed learning, Journal of Machine Learning Research, 24(276)11-50.
Charikar, M., Chen, K., and Farach-Colton, M. (2002). Finding frequent items in data streams. In International Colloquium on Automata, Languages, and Programming, pages 693-703, Springer.
Chen, C., Shen, L., Liu, W., and Luo, Z.-Q. (2022). Efficient-adam: Communicaiion-eficient distributed adam with complexity analysis. urXiv preprint arXiv:2205. 14473.
Chen, G., Xie, K., Tu, Y.. Song, T., Xu, Y., Hu, J., and Xin, L. (2023). Nqft: Nonuniform quantization for communication efficient federated learning. IEEE Communications Letlers.
Devlin, J. (2018). Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.0480)
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al. (2020), An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010.1 1929.
Fatkhullin, I., Tyurin, A., and Richtárik, (2024). Momentum provably improves error feedback! Advances in Neural Information Processing Systems, 36.
Ghorbani, B., Krishnan, S., and Xiao, Y. (2019). An investigation into neural net optimization via hessian eigenvalue density. In International Conference on Machine Learning, pages 2232-2241.
PMLR.
Gorbunov, E., Burlachenko, K. P., Li, Z., and Richtárik, P. (2021). Marina: Faster non-convex
distributed learning with compression. In International Conference on Machine Learning, pages 3788-3798. PMLR.
Harvey, J., J., Liaw, C., Plan, Y., and Randhawa, S. (2019). Tight analyses for non-smooth stochastic gradient descent. In Conference on Learning Theory, pages 1579-1613. PMLR.
Ipsen, I. C. and Saibaba, A. K. (2024). Stable rank and intrinsic dimension of real and complex matrices. arXiv preprint arX1v:440..21594.
Jiang, S., Sharma, P., and Joshi, G. (2024). Correlation aware sparsified mean estimation using random projection. Advances in Neural Information Processing Systems, 36.
Kingma, D. P. and Ba, J. (2014). Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6960
Kpotufe, S. and Sriperumbudur, B. (2020). Gaussian sketching yields a ji lemma in rkhs. In International Conference on Artificial Intelligence and Statistics, pages 3928-3937. PMLR.
Krizhevsky, A., Hinton, G., et al. (2009). Learning multiple layers of features from tiny images.
Li, X., Gu, Q. Zhou, Y., Chen, T., and Banerjee, A. (2020). Hessian based analysis of sgd for deep nets: Dynamics and generalization. In Proceedings of the 2020 SIAM International Conference on Data Mining, pages 190-198. SIAM.

---

Liao, Z. and Mahoney, M. W. (2021). Hessian eigenspectra of more realistic nonlinear models. Advances in Neural Information Processing Systems, 34:20104-20117.
Liu, H., He, F., and Cao, G. (2023a). Communication--ffiicen federated learning for heterogeneous edge devices based on adaptive gradient quantization. In JEEE INFOCOM 2023-IEEE Conference on Computer Communications, pages 1-10. IEEE.
Liu, H., Li, Z., Hall, D., Liang, P., and Ma, T. (2023b), Sophia: A scalable stochastic second-order optimizer for language model pre-training, arXiv preprint arXiv:2305. 14342.
Liu, J., Huang, J., Zhou, Y., Li, X., Ji, S., Xiong, H., and Dou, D. (2022). From distributed machine learning to federated learning: A survey, Knowledge and Information Systems, 64(4):885-917.
Lu, Y., Dhillon, P., Foster, D. P., and Ungar, I.. (2013). Faster ridge regression via the subsampled randomized hadamard transform. Advances in neural information processing systems, 26.
Mishchenko, K., Gorbunov, E., Takac, M., and Richtarik, P. (2024). Distributed learning with compressed gradient differences. Optimization Methods and Software, pages 1-16.
Mou, W., Li, C. J., Wainwright, M. J., Bartlett, P. L., and Jordan, M. I. (2020). On linear stochastic approximation: Fine-grained polyak-ruppert and non-asymptotic concentration. In Conference on learning theory, pages 2947-2997, PMLR.
Reddi, S., Charles, Z., Zaheer, M., Garrett, Z., Rush, K., Konecny, J., Kumar, S., and McMahan, H. B. (2020). Adaptive federated optimization. arXiv preprint arXiv:2003.002: 00295..
Reddi, S. J., Kale, S., and Kumar, S. (2019). On the convergence of adam and beyond. arXiv preprint vrXiv:1904.09237
Reisizadeh, A., Mokhtari, A., Hassani, H., Jadbabaie, A., and Pedarsani, R. (2020). Fedpaq: A communication-effiiiee federated learning method with periodic averaging and quantization. In International conference artificial intelligence and statistics, pages 2021-2031. PMLR.
Richtárik, P., Sokolov, I., and Fatkhullin, (2021). Ef21: new, simpler, theoretically better, and practically faster error feedback. Advances in Neural Information Processing Systems, 34:4384- 4396.
Richtárik, P., Sokolov. I., Gasanov, E., Fatkhullin, I., Li, Z., and Gorbunov, E. (2022). 3pc: Three point compressors for communication aggregation. In International Conference on Machine Learning, pages 18596-18648. PMLR.
distributed training and a better theory for lazy
Rothchild, D., Panda, A., Ullah, F., Ivkin, N., Stoica, I., Braverman, V., Gonzalez, J., and Arora, R. (2020). Fetchsgd: Communication--ffice federated learning with sketching. In International Conference on Machine Learning, pages 8253-8265. PMLR.
Safaryan, M., Islamov, R., Qian, X., and Richtárik, P. (2021). Fednl: Making newton-type methods applicable to federated learning. arXiv preprint UrXiv:2106.022
Sagun, L.,, Bottou, L.,, and LeCun, Y. (2016). Eigenvalues the hessian in deep learning: Singularity and beyond. arXiv preprint
Seide, F., Fu, H., Droppo, J., Li, G., and Yu, D. (2014). 1-bit stochastic gradient descent and its application to data- parallel distributed training of speech dnns. In Interspeech, volume 2014, pages 1058-1066 Singapore.
Shrivastava, M., Isik, B., Li, Q., Koyejo, S., and Banerjee, A. (2024). Sketching for distributed deep learning: A sharper analysis. Advances in Neural Information Processing Systems, 37:6417-66 -447.
Song, Z., Wang, Y., Yu, Z., and Zhang, L. (2023). Sketching for first order method: efficient algorithm for low-bandwidth channel and vulnerability, In Internationaal Conference on Machine Learning, pages 32365-32417. PMLR.
Spring, R., Kyrillidis, A., Mohan, V., and Shrivastava, A. (2019). Compressing gradient optimizers via count-sketches. In International Conference Machine Learning, pages 5946-5955. PMLR.
12

---

Stich, S. U. (2018), Local sgd converges fast and communicates little. (at.::00800977
arXiv preprint
Szlendak, R., Tyurin, A., and Richtárik, P. (2021). Permutation compressors for provably faster distributed nonconvex optimization. arXiv preprint arXiv:2110.0333
Tang, H., Gan, S., Awan, A. A., Rajbhandari, S., Li, C., Lian, X., Liu, J., Zhang, C., and He, Y. (2021). 1-bit adam: Communication efficient large-scale training with adam's convergence speed. In International Conference on Machine Learning, pages 10118-10122 PMLR.
Vershynin, R. (2018). High-dimensional probability: An introduction with applications in data science, volume 47. Cambridge university press.
Wang. A., Singh, A., Michael, J., Hill, F., Levy, O., and Bowman, S. R. (2018). Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv: (80.000004774..
Wang. J., Lu, Y.. Yuan, B., Chen, B., Liang, P., De Sa, C.,, Re, C., and Zhang, C. (2023). Cocktailsgd: Fine-tuning foundation models over 500mbps networks. In International Conference on Machine Learning, pages 36058-36076 PMLR.
Wang, Y., Lin, L., and Chen, J. (2022a). Communication-compressss adaptive gradient method for distributed nonconvex optimization. In International Conference on Artificial Intelligence and Statistics, pages 6292-6320. PMLR.
Wang, Y., Lin, L., and Chen, J. (2022b). Communication-effcce adaptive federated learning. In International conference on machine learning, pages 22802-2222 22838. PMLR.
Wu, J., Huang, W., Huang, J., and Zhang. T. (2018). Error compensated quantized sgd and its applications to large-scale distributed optimization. International Conference on Machine Learning, pages 5325-5333. PMLR.
Wu, Y. and He, K. (2018). Group normalization. In Proceedings of the European conference computer vision (ECCV), pages 3-19.
Xie, C., Zheng, S., Koyejo, S., Gupta, I., Li, M., and Lin, H. (2020). Cser: Communication-efficieent sgd with enror reset. Advances in Neural Information Processing Systems,
Xie, Z., Tang, Q.-Y., Cai, Y., Sun, M., and Li, P. (2022). On the power-law hessian spectrums in deep learning, arXiv preprint 2rXiv:2201.
Xu, H., Zhang, W., Fei, J., Wu, Y., Xie, T., Huang, J., Xie, Y., Elhoseiny, M., and Kalnis, P. (2023). Slamb: accelerated large batch training with sparse communication. International Conference Machine Learning, pages 38801-38825, PMLR.
Yao, Z.,, Gholami, A., Keutzer, K., and Mahoney, M. W. (2020). Pyhessian: Neural networks through the lens of the hessian. 2020 IEEE international conference on big data (Big data), pages 581-590. IEEE.
Zhang, Y., Chen, C., Ding, T., Li. Z., Sun, R., and Luo, Z.-Q. (2024). Why transformers need adam: A hessian perspective. arXiv preprint 17878.2402.16788
13

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justificcton right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No]", ", it is perfectly acceptable to answer "[No] provided proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: We have clearly stated the contributions and scopes of this work in the abstract and introduction.
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 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]
14

---

We point out that our main results depends on the intrinsic dimension in T2
Justification:
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point any strong assumptions and how robust the results 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 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 or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails 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, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant 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:
All proof are provided 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 referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide 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 supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: dlyeer-prameed disclosed in the appendix.
Guidelines:
15

---

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 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 specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, 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 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 submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily 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 way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome deseribe the particular way they provide for reproducibilityy 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?
Answer: [Yes]
Justification:
The code is public.
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/
CorooodeeSuumisiiooooooooiyyyy 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 to reproduce the results. See the NeurlPS code and data submission guidelines (https
//ipss.cc/publi//guue
Codeeuumiiii wettsbbmmssiiioi ssw.ssolicyyyyiyi 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 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).
16

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
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?
Question:
Answer: [Yes]
Justification: Hyperparameters described in the appendix.
Guidelines:
The answer NA means that the paper does not 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
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: The analysis of statistical significance can be found in Appendix
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.)
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.
OK to report 1-sigma error bars, but one should state 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 out of range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables 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 to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Compute resource information is disclosed in appendix.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
17

---

The paper should provide the amount of compute required for each of the individual experimental runs as 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 https://neurips. http://wwww.....ppaaoccccccccaaa
Answer: [Yes]
Justification: All ethic codes observed.
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 to preserve anonymity (e.g.. if there special consid- eration 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: The broader impact is stated in the appendix.
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 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 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 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
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, scraped datasets)?
Answer: [NA]
Justification: There is no model of high risk.
Guidelines:
The answer NA means that the paper poses such risks.
18

---

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 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
Are the creators 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?
Question:
Answer: [Yes]
Justification: All assets are cited.
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 used and, if possible, include 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, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- 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 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
The code is well-documented.
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 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: No crowdsourcing used.
19

---

Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information 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 the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, 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 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: Not Applicable.
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 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.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods this research? Note that if the L.M used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality the research, declaration is not required.
Answer: [NA]
Justification: No LLM is involved.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy /tte::////errrpsccccooofereeeeee for what should or should not be described.
20

---

A Related Work and Key Challenges
In this section, we briefly discuss the related work especially in the context of the key challenges in combining gradient sketching with adaptive methods, At high level, analysis of adaptive methods with gradient compression has to manage the eπror or variance due to compression for both the first moment and second moment terms, which appear as a ratio in the algorithm. There are two broad categories of gradient compressors: biased and unbiased. Biased compressors generally need to handle errors from first-order terms because of the distortion in the gradient direction whereas unbiased compressors typically have high variance which affects the second-order terms.
Unbiased Compression. Denote as the compression operator over (gradient) vector x. The compression error w can be characterized by B[((()) - x|| < w||=[|| The convergence rates of such compressed gradient methods heavily depend on w. For the family of unbiased compressors. can have linear dependence on d. For instance. L2 -z-quantization and unbiased RandK sparsi- d
fier (Beznosikov et al. 2023) achieves w & - 1. While recent works (Szlendak et al. 2021 show that the convergence rate depends on the number of clients C, arguably in many practical settings with modern deep learning models, the number of parameters d' (hundreds of billions or more) is much larger than the number of active clients C (millions). The usage of such unbiased compressors effectively leads to dimension-ddependent convergence rate in compressed gradient based distributed learning methods such as MARINA (Gorbunov et al. 2021) and FedPAQ (Reisizadeh et al. 2020). (Song et al. 2023) is an inspiring work that introduces the sketch-and aehb-andddeskette framework to dis- tributed learning. However, the convergence bound derived in (Song et al.) 2023) limited since their communication cost scales linearly with the number of model parameters d, i.c., O(d), which is also there in prior results using other compression algorithms.
Biased Compression. Biased gradient compressors achieve a lower variance in the compression error. TopK and biased RandK, which are commonly-used contractive compressors, achieve w < (1 4). The issue of the biased methods leading to optimization divergence even in simple cases (Beznosikov cet (al.) can be mitigated by introducing error feedback (EF) mechanisms (Seide et 2014). However, the ofate-of-the-er aut eπrоr feedback EF21 (Richtárik et al. 2021) still suffer from distortion error which is proportional to x- The dimensional dependence is inherited to the convergence rate of CocktailSGD (Wang al., 2023) and 3PC (Richtárik al., (022) that employ biased gradient compressions. Further, most of the EF developments do not explicitly show compatibility with adaptive methods, which involve anisotropic and nonlinear updates (Tang al. (002T)
Communication. The use of gradient compression calls for designing adequate transmission mechanisms. For sparsifying compressions, such TopK and RandK, the average of sparse client gradients is possibly dense, which increases the downlink (server-to-client transmission overhead. In the worst case, plain average of the client gradients in MARINA (Gorbunov al. 202 leads to &C in the number of non-zer zero bits. FetchSGD (Rothchild et al. 2020) mitigates the problem by adopting an extra call of topK compressor on the server side at additional compression costs.
Additional Desiderata and Challenges. Instead of an in-expectation analysis, which can hide high variance behavior, we want a high probability bound, which has to consider two inherent sources of noise such settings: stochastic mini-batch and compression, as well as the third derived source of noise, due to the sequential dependencies of stochastic adaptive iterative updates. Further, while assuming gradient norms to be uniformly bounded simplifies analyses, we would like the guarantees to continue to hold without such an assumption. Finally, want the analysis to hold when the local clients are doing multiple steps of updates, unlike the single update common in traditional distributed learning (Mishchenko et al. 2024: Tang et al. [021] Wang et al. 2023).
B Lemma for Random Sketching
For completeness, we provide the following lemmas that give high probability bounds on the inner products.
Lemma B.I. (SaHe)/Same as Lemma D.23 Song et al. (2023) Let R E Rubxd denote a subsample randomized Hadamard transform or AMS sketching matrix. Then for any fixed vector h E R and any
21

---

fixed vector g E IR the following properties hold:
log log2.3 (d/5) (dd/))
K Kg R Rhh
Th
||||||10||2 s (
vb
Lemma B.2. (Gaussian) Samee Lemma D.24 Song et al. (202)) Let R. E. I denote random
Gaussian matrix. Then any fixed vector h E R2 and any fixed vector E the following properties hold:
log 1.5 (d/5)
P Kg T R Rh 99
1/112
|
1 O()).
vb
Lemma B.3. (Count-Sketch)f as Lemma D.25 Song et al. (2023)] Let E bR denote a count-sketch matrix. Then for any fixed vector h R and any fixed vector g R the following properties hold:
P [[g R Rh g¹ hh| log(1/6)1l9112 [1]2]] s (()).
C Proof of Main Results
Theorem C.1. Suppose sequence of iterates {2x)T= is generated by Algorithm (SDDL) with a constant learning rate = 1/- Under Assumptions 1-4, for any T and > 0, with probability 1 ((6) ((xx((-((v2)) 89.
knJiK T
[Vxx)))]1 <((x1) = =kn= LK* + VT(
log 1.5 (CKTd/8) C
log 1 27
vb
8g
+1/xI(1+ log (CKdT" 10)
SKILK? + 2G G² (1-3:)²
vb
where Jg. and v the randomness of sketching, sub sub-Gaussian noise, and martingales respectively.
and J1 = (V1+ ICK' nK
VB.
Remark 7. (Dependence on K) The convergence bound in Theorem C.I has a dependence on K. The primary focus of this work is reduce the communication cost in FL. algorithms, where the cost only depends on T and sketching dimension b. Therefore, we view K as a constant throughout the work. As we will show in Corollaries 1 and if we set step-size y as O(1KVTT), which is the same Reddi al. 2020), the dependence on K in the bound can be eliminated.
C.1
Proof of Lemma 2
Let
B1 (x) xt-1)
1
B1
+ 1 81
T1-1-
11
1
11
22

---

Then, the update on 2 can be expressed as
B1
2(+1 - Zt
(I++1
(Tt 21-1)
1
B1
1-
31
1
=V,
-1/2
mm +
31
xVii-1
-1/2
-m4-1
31
1-
1
RVe
-1/2
(31m4-1 + +1 (1-B1)
RR m+) +
31
1/2
m
31
1- B1
C
81
1/2
RVV. -1/2)
m4-1
18 C
V1
112
R
1-
31
31
(KVV 1/2 - xV, -1/2
m4-1
-.V. V1
1/2
R
R4 (2 xg,K)
I- 81
C K
B1
(KVi-1/2 nVV. 11/2)
Ki V/
1/2
2
Σ
R Rigi.
m
1 1- 31
C
By Taylor expansion, we have
1 ((2++1) = [(z)) + V((2)) [(0+1 =1) -2))+(2++1 21) iic(2t+1 -x1)
=[(x) + V((x)) ( - 2) + +VV((z) - v=(xx)) (2441 - 21) + + - Hc(E441 21).
+2)++
(1)
Bounding the first-order term
(2411) - ze)
(31 31 (1- (NViT1/² xV, 1/2) m4-1
C K
-V((x))
KJ 112
Σ Σ
RT Regix
k=1
C K
B1
(nt/=1/2 v T
RV/i - 112) m4-1- C
"
(no, - 1/2
R Rigi
-1 B1
1=
K (c(x)) v/ 1/2 C
R Rigi
For the difference term, applying Lemma B.2 yields
K
" ((()) (KVi 1/2
VV.
1/2
R
Righ
nK C (1+
log (CKTd/0)
1/2
/
112
Σ
C K
19
Vb
c= k=1
Denote [:] the i-th element of a vector. The 12-norm
11V 1/2
v max
max D 0.-1.
Vo--i++ Vor, +E ( + E) 0e, +c)
tr Ōz-1,i
max
( (or-1 (vv + al/00rs + (
23

---

By definition, d4 = max(0(-1, If 02,s = - the RHS is O. Otherwise, θz,e = 01,i.
1/2
1
Vt 12-
max
(V00-15 + +
01 ( + v vor
(1 [2) 0(-1,1)
< max
22 2/LLL 82)
vI- 32
max
v Drii
6²
V1 32
C
K
R Rugf
2
1=1
21 (1 82)
log 1.5 (CKId) 10)
G.
2
vb
The first inequality is from 6 1. The second inequality comes from dz,i 2 Uz >
(1 - --,) The last inequality follows from applying Lemma B.2 to each dimension of g. Plugging into the bound for the difference term
1122 -1121/1 !/ ((xx)) (NV)
~V ΣΣ
R Regi
C
k=1 (CK)tP))))))))))))))))))))))))) (0) 3012 10)
2n nv (1 (1-32) 32)
log
2 (1+
vb
The quadratic terms can be written as
(VC(2)) V((z)) - 2)) (2e xT) Hcl THc(
1
( 22)
B1
(xr x(-1)),
B1
31
where He is second-order Taylor remainder. So the quadratic term can be further
quadratic form over 2++1 24 and 2t It, denote as Q(2++1 2t, - 2²). For the same reason, the term (Ee41 - 2²) Hz(=111 21) can also be written into a quadratic form Q(2+11 - 22, 2+41 2)). Putting the two terms together yields quadratic form of @(2+41 24, 2 21).
Overall, the descent lemma can be written as
C(2+11
<C(z)) + 2//=v1- 82 (1 -B1)c2 (1
log 1 (CKIP/0) 33/2 ) [ 1/2 1/2
R Rigix
Vb
+ (VL(=)) - vl(xz)) (z411
(((x)))((<4111 - 2) + 2 ( - T -o(==+1 =z).
C.2 Proof of Lemma
CR
(Generalized version of Lemma 2.4)
Proof. We can prove by induction. For / = 0, since mo = O, the inequality holds. Suppose we have for l E Idd s.t. |||]| < H. with probability 1 - @(( - 1)6),
log 1-5 (CKd/0) G
|m| Al < (1
Vb
24

---

Then by the update rule,
K Σ
mm h/] = 1(9 m4-1 + (1 (1)))
7
Σ
R2 Regi,)) M
C
c=k=
(1 B1)n
K
I 1/1
(R, Rigi, h)l
C
Cm km1
log
1.5 (CKd/0) )n
3/1mm/11 (1-
|196,112111111
vb
log log (CKd/8)))
$ (1+
vb mKGHH w.p. - O(t8).
Let 1/21/20
1/2 vx. Then ||A||2 s 1/e.. We have
log 1.5 (CKd//) GKG/
mo) S (1+
vb
C.3 Proof of Lemma 2.5
We first prove the element-wise lower bound of the diagonal matrix V/ Denote (V;=1/2), as the i-th element on the diagonal of v/1//2 By the update rule,
log 1.5 (CKIdd/ (M1))) (max( +c) (11+
nKG
w.p. (
Vb
where the last inequality follows by letting h as a one-hot vector h4 in Lemma B.I observing that the elements can be transformed to inner product form ve,i = v2 h4. Then the scaled gradient norm can be lower bounded as
d min(V 1/21; Σ [(())
[l 2 min(V,
2 log" (CKtd/6)
WKG - [KC+))"/[V(aaa) w.p. O(dd)
vb
which completes the proof by applying union bounded on the dimension d.
C.4 Proof of Lemma 2.6
Since the noise is zero-centered, we view the random process of
(Y Σ k=1
vL(x,) VV. R. 2
as martingale. The difference of YY+1 - Yel is bounded with high probability
Σ VV((x) (V= 1/ Regi sials log 1.5 1og". (CKd/8
K
[1+1 Y,
vb
w.p.1-
Then by Azuma's inequality,
log 1.5, (CKd/0)
P( 2 V
Σ
1/2]V,11AVL((x))))) ((cxp(-(((2)) T8
vb
25

---

Note that the original Azuma's is conditioned on a uniform bound of the difference term, while our bound here is of high probability. Hence, we need another union bound. similar bound can be achieved for the sub-Gaussian noise in stochastic gradient. Let
k
Z = Σ T=1
(L(x,) V; (o)(),
c= k=1
Then
T
2T }) O(exp(-(((2)) + óg 8g
P(IZT) 2 vk
log
Combining the two bounds by union bound completes the proof.
C.5 Proof of Theorem C
We first introduce the lemma
Lemma C.2. For any round E [7].
C K
((2+1) <f(ze)
ΣΣ
VL(x.)
V. R Regi, (2+11 He(()(Eett 21)
C
c=1 1=1
27/²x
log (CKidd /8) 31200
+ (1 -1- B1)e (VL(z)
V((x)) - z)) + (1+
Vb
where H2(2)) is the loss Hessian some 2, within the element-wise interval of [x,, 2n]
After applying Lemma C.2
The second order quadratic forms in the descent lemma can be written as
(VL(=)) v[(x))) (zt+1 21)
1
31
=(2 21)
Hcl
1 31 (I1+1 -22)
(It
xI-1))
1- 31
81
(V,-1/ m)-1) Acc
1
(-xV,-1/2 "ma) 112m)
31 (-xV,111² m(-1))
= ạ -1-8 B1
1 31
- 31
=x²
31
(1-11/2
1) T Hcc(V, mi) x²
18
(V/11/2
11-1/2mm-1 m Mc (V m--1),
(1 3,y²
(1 B,
and
(#t+1 - Ti1((244111 - 21)
31
(xx xx-1)) Mcl -xccc) 1
81
(Ie+1 -xt) 1 31
B1 (x1+1 xzz)
(I+ xz-1))
Bi
1 31
1
23g
(1 3:)3 (Ie+1 - xzt) [ - -xt) - (1- 31 (#e+1 - xza) [0] - It-1)
B7
+
(x: - x1-1) Hel I(-1);
(2)
(1 - 3,
1/2 which is essentially quadratic form defined on 1/2 m4 and 1 Hence, provide generalized version of Lemma 2.4 as follows.
Lemma C.3. With probability 1 (t6), for eigenvector i; the (1 log à (CKd/))
the Hessian matrix, Tual s
)KGG/e
Vb
26

---

Note that o can be any basis and is constant throughout the training process. Then the sum of quadratic forms written as
(V((z)
V[(x=)) (2411 - -2)
Sh²
81
(Vi-11² m4-1) -A,(VV HV -nk2 -
(V=11/
1/2 m4-1) (V m)-1);
(1 -3112
(1-
81
d
17
d
=
Σ
>,(VV- 1//-1/24441 1+-1) (u;e,) VV, 1/2 m/m - x² (1 B1)²
Σ
>i(V 1-1
1 1) (mru) \V, /V m
(1 - 31)
1=1
i=1
81
87
Sh²
1/2 ox
mel + R²2
-
m)-1) nil
31) 2
(1 B1)
i=1
1=1
2
log" 1.5 log-. (CKd/00)
Sh²
IL(1+
(3)
(1 - 31
vb
where the last inequality is by B1 S 1 and Lemma. C.3
First-Order Descent Term, The first-order term the descent lemma can decomposed into three components, which we will handle separately.
VL(x,) VV Rigi, VL(xx)
-
Rigi, VL
D.
Da
(V Vc.
(VL
Ds
First, D3 can be reduced to second-order term by smoothness over C,
vL(x:) (VL" VV'))
(2 x4)
v((xz) T VV, 1/2 V
7=1
Σ
11, LKG
(4)
Note that this term does not involve any stochasticity with regard to random sketching, which means we can directly derive the upper bound by Cauchy-Schwartz. in the last inequality.
(x)) = vL(x)). D1 composes scaled squared gradient norm. Applying element-we wise high probability bound on random sketching yields the lower bound for the scale. By Lemma 2.5 we can derive the lower bound for D, Note that applying union bound to D, does not introduce another dependence, since é, is monotonically non-decreasing.
Next observing that 1
C -==1
Martingale for zero-centered noise. DD contains zero-centered noise term R, R.gi, VL (xf where the randomness is over R, and the mini-batch noise at round t. Despite 2f, has temporal dependence, the fresh noise at round / is independent of the randomness in the previous iterations. Hence, the random process defined by the aggregation of these norm terms over time forms martingale. By Lemma 2.6 can bound this term Dz.
27

---

Finally, putting these parts together by union bound over (T') and telescoping the descent lemma leads To show this, with probability 1 (6))
to Theorem C.I
((4+1)
(
C
K
(24)
VC(x))
VV. R2 Regi, + (2+11 Hc(2(((2 - 21)
C
=lk=1
21²F1
log 1. (CKtdP 10) 0012 + (VL(2)) V-(--)) (2441 21)
+ (1 - 1-))
(1+
vb
(10)
log 1.3 1.:1.,(CC (CKIP/0) 16)
[KG++) 2 T.K?
1
vb
C K
+ K3)
VL(x)) TV, 1/2112(() (R. Rigi - - (zf,))
C
k=
2
IL(1+ log (CKd/)) 132=220112
R²
(1- B1)²
Vb
2772K K
log (CKttt 10) 3312 + (a++1 - (z(())((+11111- - 2))
+ (1 BI)E2 (1+
vb
/1+ log 1.5 log!. (CKtdP//) 15)
((0)
I ((2)) k37(/
1+
/ 1 Ln
vb
+ K7 C K
VL(z.) (V, Rigi, (z
0
c=1 k=1
2 (1 31
IL(1+ log 1.5 (CKd/0) 2
(1 (1- -1)²
db
27/2k
log2 (CKUd( (C)
(1 - (1-2 (1 +
vb
where (1) follows from Lemma. C.2
(it) follows from 3 and Lemma 2.5 (iti) follows from
Summing up with E [T) and moving the squared gradient norm term to the left hand side yields that with probability - O ((exp(-((())) gg
log 1.5 1.5 (CKK12)0)
T
nKG - Σ
vb
<((2)) + - + TTpT + (1+ 31)²
log 1. (CKd/0)
IL(1+ +
2
(1- -,)
vb
+ 27
log 1.5 (CKttt/))) 10) 10))3/223
(1+
31
vb
T
+ ΣΣ ΣΣ (T....)VVVV Regi. - (Z((z)
1=1
(4)
$ S(2)) + -na I. + 2T 2+ (1+ 31)
log 1.5 (CKd/8) 22222/2
IT.(1 +
(1-
vb
kT
log 1.5 (CKitd 18) 13223
+ (1 31)e2 (1+
vb
+ uKn K VT(
log" 1.5 (CKTd/8) G²
σ log
2T
o, )
Vb
where (1) follows from Lemma.
2.6
We conclude the proof by realizing 31 S 1.
28

---

C.6 Proof of Corollary 1
1 log log' (CKKPT3/0) K G approaches
In the aysmptotic regime, with sufficiently large T, the term
6, so the denominator on the LHS can be replaced with 2c. Then the derivation is straightforward by just substituting VTK -TK into Theorem C.1
C.7 Proof of Corollary 2
We first develop the convergence bound in Theorem
C.1
under the condition b 2 log (CK
(V2nKG++)** kn
T
< ((z))
+ ~
1
1=1
G² + DRnk VT
27
32
ILK²G²
log²
2T
8g
(1
31
2
The condition on T < Ja-y2G equivalent to
V2nKG + nkK
< J3,
since 1)
Then scaling the coefficient the left hand side and substituting
for n,
VTK"
VTK
derive
|L(x))
((z))e
1 LG
+ logå
2T
+
32
ILG
KT
c T
og
T
31 c
C.8 A non-usymptotic bound on practical learning rates
We first state a convergence bound on using practical learning rates, which decays as the optimization procedure.
Theorem C.4. Suppose the sequence of iterates {xe)T=I is generated by Algorithm 1 with decaying
learning rate пu VIFT5K' where Tu [ 1- 1. Under Assumptions 1-4, for any T and > with probability - (6)) - O(exp( (((p()))(()))
6g.
log (CKKTT418) nJKG++
Σ
V & ((x)) + KLG 1L log
vb
Vw log T( log. (CKTd/O) G2
2
21 log T(1
+ x²
(CKdd (8)
ILG
Vb
8,
vb
(1-3,)
2
where 8, Sy. and u are the randomness from sketching, sub-Gaussen stochastic noise and martingales respectively, and J is constant defined in Lemma.
Alike the analysis in the constant learning rate case, we first define auxiliary variables 2t
B1
B1
+ 1 B1
(x+ xt-1)
31
1 11
Then, the update on z can be expressed as
31
2t+1
(xe+1 - xz.)
(It - I(-1)
B1
1 B1
1/2 K
B1
(NV/-1/2
11/2
-nV, -1/2
KTTt
V.
R Reg
m4-1
1- B
C
c=1 km
29

---

By Taylor expansion, we have
1
((z+41) = [(z)) + v2(22) (24+1 - 21) 2 =[(x) + V[(x)) (2441 2)) + (V[(2)) V((x))) (2+41 21)
(2(41 - 21) TT) Hc( - 21)
1
(2442 22(++222 21) TH)TT(cC Hc(214 21).
2
Bounding the first-order term
vL(x)) (c(x)))(Z++1+ =za)
( 81 (NV:-1²² -kV, 112) m4-1
C K
-V((x))
KT) V.
-1/2
Σ
R ReDi,
B1
C
c= c=1==1
B1 [V
112,1 ||1111 -1/2
11,1))
1 - 81
C
K
C K
th (xxs (KVV
11/2
-KV 1/2)
Σ
R4
Regi,
-
Hill V(Ix)) [V,
1/2
R Regi,
c= 1=
c= k=
81
K
R₂ Regi ) (I
|m 10000010001
Σ
||:1// 11,1)
1 31
C K
63)
RR. Regin
The quadratic terms can be written as
V[(()))(((())))
- 2) = -21-x)) TT)TTTAA B1
1
(11+1 x24)
B1
(xr: xt-1)),
31
where He is second-order Taylor remainder.
To bound the quadratic term, the counterpart of Lemma can be stated as
Lemma C.5. With learning rate = 0 v+16
where To [ 1- 1.
Denote J
1-81 VT++I
VbT ). Then with probability 1 - O
log 1.5 (CKd/0) JKGH
| M S (1+
Vb
Proof. For f. = 0, since my 0, the inequality holds. Suppose have for h Rd., s.t. ||A||| & H. with probability - (( -))0),
log 1.5 (CKd/8)
mm I hl S (1+
JKGH
Vb
By the update rule,
C K
|mm h- = 1(3 m + (1-)))
Σ
R Regt TM
C
t=k=
k=
(1 (1-B1 11)1
C
$ 3/1mmT,h1 +
(RRT Regi h)l
C
log 1.5 (CKd/0)
S Bi B.immT1 + (1-8:)(1
MM
l 1 |2211//22
vb
log (CKd//)
< (1+ +
JKGHH w.p. 1 - ((t)).
Vb
30

---

By exactly the same as in Sec.
C.3
we can lower bound the scaled gradient term by
V
2 min min(V, V ) Σ [VL())))
2 2V1 log 15 (CKId/6) n KKG+
w.p. O(d)).
vb
On the martingale of zero-centered noises, we can simply incorporate the learning rate 7) into the martingale. Define the random process of sketching noise
K
(v,
Y r=l
VL(x,)
( TV R.g Di
k=1
martingale. The difference of [Y, Y - Y--1
Y--1 1| is bounded with high probability
71
VL(x)) (R, Regi si,x)
c= k=1
1.5 log
(d/5)
T/ I, w.p. 1 1- 1
Vb
Then by Azuma's inequality,
log 1.5
(d/0) ni
P(Y
Σ
) T5 (5)
Vb
A similar bound can be achieved for the sub-Gaussian noise in stochastic gradient. Let
K
Z
(((x
TV; (of - (x)xx,
Then
27 )) ((ex((-(((2)) 6g
P( 2
Σ
2 log( 5g
Combining the two bounds by union bound completes the proof.
D Convergence Without Bounded Gradient Norm Assumption (Proof of Theorem 3.4)
We first prove the local client gradient L bounded. The client performs stochastic gradient descent L = Ie - 1) -f==1 gi,r. Let
VK
Lemma D.1. Under Assumption
Let S (-22/02) ) by
The local gradients of lo S K can be bounded
2LVK'
with probability K5. - exp(
2
|V
222 In
+1
2
4LL((x)) 22.
8.
de
Proof.
VL < (x) S (z()) + Σ (v))
=( + 2
| 11.x-111
2L
(xx) -
(z) (L + E + 2
k=1
K
(x)) +
|VC (2 Σ
VC
(zf
31

---

Take induction basis 5 I h 1. We have bounded gradient (-),)) < G with probability
1 -rè exp( A /////2 /0) The RHS can be bounded with probability 1 - kdo kexp(-A2/o²)
by
(x,) (1
+ Σ
||,,,-1 12)
(x))
To VR
2KGA
2
+
ThL
K(G+A9)
60
K
2
(x)) 1oU 2GA
LG + A²)
8c
(x)) 10 G A
L A2) S
G
2L
Let 110 < 1.
27. and G = (v 2222 2n
2
+ V
2 In
2 (LL((I)) + A2)? we have
+
be
5
RHS (x)) + 170 G 7
In
2
L A In
2
2
80
1
1
2
1
2
G
S (x:) +
G
+
A
+
GG + A In
4L
2L
8.
4L
5c
2L
Lemma 3.2. Under Assumption
3 Let S 2LVK
The local gradients as of l S K can be bounded
with probability CKo. CK (oop--A2/0") by
(xf $ 2VVV + 2 21 2 In + A.
C
Proof. Applying the fact that [(x")= = EC= (xxx) to Lemma D.I the averaged local gradient can be bounded by the global loss,
Σ IVC 2VVI/ 2
de + +
The averaged local gradient norm will be key component in the following analysis that focuses on the global gradients.
Consider the server optimizer
1
((2411) = [(zx) + VL(22) T (ze+1 -2) - 2,) He(2e+11 - 2x)
1
= =((z)) + VL(x)) ( - +
+ (VD(z) vl(x,)) ((++1 2x) + + - Ac(2 - 21).
vL(xx) (((+1) zx)
-V((z))
1- [ - xV." 1/3) M4-1
31 1/
KT
V.
1/2 C K
R
RediR
C
c=1
k=1
B1
vL(xe)
(Nt/-1/22 - -iVi -1/2
m
1- Bn
C K
K7)
TV.
-1/2
(z)
+ VL + t -gî, + R. Regi
C
k=
32

---

1
1
2
2 +1LL((11)
7
(x
||
Σ
2
+
2
In
+ A
C
C
de
dc
$ 0
Σ
2V
(Z) + +2
2A2In In
2
+
dc
2VI
2
C
(x1) + 1/1 2 In
+
A
C
8.
c=1
-2 -2VLV ((x+) + v
2A In
+
A
8
where the third inequality follows by Cauchy-Schwarz, On the server side, we consider the induction in
C basis II [x,,) < [(x)) S 2I w.p. 1 - -tck-KKKeeppp - holds
J
for 3 T - 1. The following inequality holds with probability 1 - K6g- CK exp((A22p1)ppp /
K ΣΣ (Z)
C
K v V/
1/2
+ (z,)
C K
k-1 Σ I
KJ C
IVL
1/2
c=
k=
I C
1/2
C
K
Gi-ll
Kn I C
C K
k-1
[V ||
ΣΣΣ
(xf
A
1=1===
Kn K VL VL(E)) + 2
2 21
+ +1)+
V [
5e
V2 ((xx)+ +
2n12 1
2 2
VL(x:)]AA(
8e
And for the difference term, applying Lemma B.2 yields
C K 1/2 -x
7 (LLx)) 1/2
Σ RT Regi,
C
log (CKTd/0)
C K
s n c (1+
)|
1/2 - -1-1/211
1|2
Σ
|
vb
c=1 k=
Denote [:] as the i-th element of a vector. The 12-norm
11, 1/2 -V,1//11 max
max D 0,-1, 1,
/ VoL, +E
( ( + E) 0e, +c)
Er Gz-1,i
max
( (or-i + Vvr. + D(v + (
33

---

By definition, 8/ =
If 02,5 = - the RHS is O. Otherwise, θz,e = 01,
[V, 1/2
//-1//2 max
0/ 1,5
(v (11-11 + +e((v + V + (
(1 (2r, 14-1,))
max
<² ev(1 -B2)04s
vT- 132
S max
Vz,i
2²
K
B2
112
max
R &
4V 2(1 82)
log
1. (CKtd 15) (V2G + 2 + 22 2CK 1). 8.
22
vb
The first inequality is from 62-1, 2 0t-1, The second inequality comes from θτ,i 2 4t 2 (1- 0r The last inequality follows from applying Lemma B.2 to each dimension of gi Plugging into the bound for the difference term
7 V((,) (NVV 1/2 xV
1/2 (CKid 1ol
R R.gi,
C
=1k=1
// - è (1 log 1.
2CK
+ + +1 2l 5.
))²
Vb
Consider the sketching noise term. Since the noise is zero-centered, we view the random process of
(Y,
1
VV
R.S DE
as martingale. The difference of \Y++ Yı bounded with high probability
C K
log 1.5 (CKd/o) K
M4+1 Vil
VL(x.) // R... Di S
C
vb
k=1
log 155 (CKd/8) K
2VLLC(1)) + +2V v /2ln d0
2CK
Vb
Then by Azuma's inequality, with probability at least 1 - Texp(- -S(22)) TT.
TCK exp( /2/22)
1/2
Yr Sv log 1.5 (CKTd/8)
(2V2LL(())) + 2IV + 2ln 8,
2CK
))²
vb
log 1.5
(CKTd/8) K VVTV22++2CCN + 22
2CK
<V
Vb
5,
where the second inequality follows from the induction basis.
We also consider the product term (V,-1/2 m
Lemma D.2. With probability 1 (tt) for eigenvector i the Hessian matrix, I(v,-1/3m)) Tval
(1 log (CK()))
KKG.e
Proof. We can prove by induction. For t= = O, since mo = 0, the inequality holds. By the induction basis, || has a uniform upper bound. Suppose we have for E IRd, s.t. |[b]| < H, with probability 1 @(( - 1)6),
log 1.5 (CKd/8)
2CK (nKH(V2G + 2 8, ))
|m_,/ hl < (1+
vb
34

---

Then by the update rule,
C
m h| 1(31 m4-1 (1-
1
Σ
Rd Regi..) Thl
k=1
(1 (1- Bm)m
C
K
S S/1lmi-/h +
[R Rigi h))
C
c1
log
C
(1--
1 l2211/|22
vb
c=1 k=1
$ (1+ log 1.5 (CKd/0)
2CK
JnKH(V2G + 20(1 +
2
}),
w.p. 1 - O(16).
vb
de
Let h V/ u. Then |1A||2 S 1/e. We have
(/2-1/2mm) T S (1
log 1.5 (CKd/0)
(nk(v2G + +20(1= + 2 6.
2CK
1)/6
vb
Then we consider the quadratic term, with probability - - ICK LCKKeCKKeppp expl-
(VL(z)) VL(x))) 21)
81
-xx .8
Sx²
(V, //
m--1)
(1-3,)3 m4-1) THo m-)))
(V.
1/2
(1 3132
81
\i 1/2 m.-1) >
(T (v 10, 11/2
me x²
8
>x 112 (/ (ora) ( 1//1//2 m
m--1)
(1
3
(1 (1-3))²
81
MM " 1/ + (1
3
MM Tm) m
(1
31)2
B1)
1=11
2R²
2
IL(1+
log log. (CKd
(, +2 +
22m 8.
2CK
1)
(1 31
Vb
8
log (CKK /5) 2,2 222 2 8.
2CK
<x²
IL(1+
(1 - 31)2
Vb
where the last but one inequality is by B1 < 1 and Lemma.
D.2
35

---

Putting all these things together, TCK exp( hap(-A2/00 /²)0²) T
with
probability
1 Texp(-S((2)) TC8.
(V mz) + R² B7 (Vi1//me) T (V
((TT) (2() +
1 1 3
2 (1 (1- (1-31)2
T
=f(=)) +
VL(x)) T (2++1- - -2) + + (V((z))
v((x:)) (21+1 - 2x) + 1/2 2((2+1 - 21) TT)TTAA Ac((241 - =n)
B1
(VL(2T), RV, 1/2 m2) /22) +
1/2
x² B
(V, 1/2
(V.11/3m m²) THC(V. 1/2 (H((V."(mmm)
+ #1-8 B1
2 (1 31 )
8n
T-1
2v 2Rn2
2 LGL
Ana 2M2KK A(1 2l
2CK de
(2( + (1 (1-B1)E
E
e
1=1
2n k / 2(11 - -22) -B2) (1+ log 1.5 (CKtdP/0)
/02 //2 + +20(1+ + 2ln
2CK ))
6²
Vb
8e
+ 2KTM log (CKTd10) K VIIVEG+2VVAL11 + 2ln
2CK
1) + 2KTTV /T KVG
Vb
8e
16
I1+ log 1.5 (CKd (G+ (++2A((++++++++++ / 2
2CK
)3//e2
(1 B1)²
Vb
5,
1 + $1 2
log log (CK& 18) )22 "22²'2K>"II III + (22 2
2CK
1 31 (1+
)3)e
Vb
8e
8n KIKKLIG 4V2G+
TR 2(11 - -3))
log
1.5 (CKtdd /0) /33 1//2166//2
22((21) + (1 - B1)C
22 (1+
vb
+ kmo log 1. (CKTd/8) K
KVG
4V2G 2xTOVU
vb
8 8 (1 + 31)² IL(1
log (CKPP)))))) PANK
log
+ 2
(1 (1-11)
Vb
8 nē 4V2G 1 2(11111 - B22
log
1.5
(CKt(ttt//// 18) /
<2((21)
(1 BI)E
62
vb
log 1.5 (CKTd/8) K
K
+ noV
4V2G + 21j0P
vb
8 8 (1 + B1)²
log 1.5 (CKT&/D) 24nkk
+
IL(++ +
(1 - 31
vb
where the second inequality holds by v2G 2 21(1+ + / 2 In
2CK ), the third inequality holds by h
N S v.. Let
No S E min
1 - 31
log 1.5 (CKT&18) )-1
H(1 +
3' 23, VL.
vb
2CK
K2
log 1 (CKTd/)) YoU (2())
2 max (2 + 21n
)2, 5,,1
de
c
vb
(6)
47/31
BbUKA A
8 (1 + 81)2
log 1.5 (CKTd)))) y3 2n2 KLLL
+ (1- B1
(1 +
(7)
(1 - -1)
vb
2
suffice to yield RHS < C
12.
36

---

127
10P
10P
100
201
20-
Eontralee
209
20P
20"
1071
107
10
00"1
107
10"1
10nh
107
30P
20"0
105
1,2
10P
(a) MNIST
(b) Fashion-MNIST
(c) CIFAR-IO
(d) CIFAR-100
Figure 4: The power-law structure of the Hessian spectrum LeNet. Quoted from Fig.1 (Xie et al. 2022).
Furthermore, the dropped positive terms regarding the gradient norm is
11+ log (CK nK(v2G + 2
2 +e
ZnT ZLTOVTL
))
[
Vb
8e
Rearranging the terms yields the convergence result.
Finally, we give the full forms of {M:)T
Mi == 2/2c log log..5 (CKTd(8) K
Vb
M2 2v log 1. (CKTd/6) K
2CK 2VGA(1+ 22
K
de ) 2v
A
Vb
81,31
Mg = 20(()) +
(1 B1)E
M4 4 2/2(1 - B2) (1
log 1.5
(CKtP) /6) /0)/3/2
2
vb
Ms 2
4 A(1
2CK
k 2(1 - -B2)
log 1 /0) 33/2 1n
(CKtdP//)
2CK
M
2
(1
50
2
vb
8,
My 8 8+ (1+ 81)
log
1.5
(CK& 18)
I(1
2
(1)
Vb
E Experimental Details and Additional Results
Aside from the experimental configurations described in the main paper, we provide additional details. We use Cross Entropy with label smoothing the loss function. The parameter for label smoothing 0.1. We use a cosine learning rate scheduler on the server optimizer, with the minimal learning rate le 5. Client batch size is 128, and weight decay le 4. We sweep wide range of server learning rates over the set [10-1, x 10-1 10-3 5 x 10-3 10-2, x 0.1.....55 1.0, 2.0]. The algorithms share all the other hyper-parameter within each of experiments. The only exception is CocktailSGD since it's unclear how to incorporate server-side momentum in their framework. To ensure fair comparison, we tune both the server learning rate and the client side momentum for CocktailSGD. The selected hyperparameters are displayed in Table 2 We select the hyperparameters balance between the training stability and model performance.
Our experiments were conducted on a computing cluster with AMD EPYC 7713 64-Core Processor and NVIDIA A 1100 Tensor Core GPU.
provide the full comparison of the baseline methods on the benchmarks Fig. and Fig. Although CocktailSGD and FedCAMS can achieve comparable performance under compression rate
37

---

Spochi
Epoch:10
Epoch: 15
Epach: 20
Epaci: 25
Epoct: 30
201
1/6 Siginsa
Fi 20 20
Fignsss
Esoch: 35
Epoch: 40
Epach 45
Epoch: 50
Epocts 35
Epects 50
2071
1041
100
Esoch: 65
Epoch: T0
Epoch 75
Epocis 80
Epocis 85
Epocis 30
2071
Figure 5: Eigenspectrum density every 5 epochs. The model is ViT-Small and trained CIFAR1O. The majority eigenvalues concentrates near O and the density enjoys super fast decay with the absolute values of eigenvalues, indicating summable eigenspectra.
Hyperparameter
SADL
CDAdam
FetchSGD
MARINA
Server learning rate
0.01
0.01
1.0
0.1
Server momentum
0.9
0.9
0.9
0.9
Server weight decay
0.0001
0.0001
0.0001
0.0001
Client learning rate
0.1
0.1
0.1
0.1
Client momentum
0.0
0.0
0.0
0.0
Table 2: Hyperparameters of SADL,, CDAdam, FetchSGD and MARINA in ResNet.
Hyperparameter
CocktailSGD
CAMS
bittAAam
PAQ
Server learning rate
0.04
0.01
0.0001
0.1
Server momentum
0.0
0.9
0.9
0.9
Server weight decay
0.0
0.0001
0.0001
0,0001
Client learning rate
0.1
0.1
0.1
0.1
Client momentum
0.9
0.0
0.0
0.0
Table 3: Hyperparameters of CocktailSGD, CAMS, Ibit-Adam and PAQ in ResNet.
Hyperparameter
SADL
CDAdam
FetchSGI
MARINA
Server learning rate
0.001
0.0001
0.5
1.0
Server momentum
0.9
0.9
0.9
0.9
Server weight decay
0.0001
0.0001
0.0001
0.0001
Client learning rate
0.001
0.001
0.001
0.001
Client momentum
0.0
0.0
0.0
0.0
Table 4: Hyperparameters of SADL, CDAdam, FetchSGD and MARINA SST.
Hyperparameter Server learning rate
CocktailSGD
CAMS
Ibit-Adam
PAQ
0.001
0.001
0.00001
0.001
Server momentum
0.0
0.9
0.9
0.9
Server weight decay
0.0
0.0001
0.0001
0.0001
Client learning rate
0.001
0.001
0.001
0.001
Client momentum
0.9
0.0
0.0
0.0
Table Hyperparameters of CocktaiISGD, CAMS, 1bit-Adam and PAQ in SST.
Hyperparameter
SADL
CDAdam
FetchSGD
MARINA
Server learning rate
0.005
0.005
0.5
Server momentum
0.9
0.9
0.9
0.9
Server weight decay
0.005
0.005
0.005
0.005
Client learning rate
0.01
0.01
0.01
0.01
Client momentum
0.9
0.9
0.9
0.0
Table 6: Hyperparameters of SADL, CDAdam, FetchSGD and MARINA in VIT.
38

---

3 10 Byleslepoch
3x Bytes/epach
3 10' Bytes/epoch
2%
sant
Valnaten Errr
5%
Valetatoon Emrr
2%
FastEG0
MAZE
Het NA
CBA.0m NicoooD
2%
Coccal
40%
Cocida
24
CAPE
2%
1N
15.
1%
10
20
30
10
1.5
20
25-
300
25
40
Epech
Epech
Epoch
Figure 6: Validation Error on CIFAR-10. We finetune ViT-base model (with 86M parameters)
from the pretrained backbone checkpoint (Dosovitskiy et al. 2020). bit-Adam has comparable communication cost with 10%. SADL shows competitive performance under all communication
budgets.
2e5 bits/epach
2eb bits/epoch
50%
5AD
CDAdam
50%
CDAcam
5/D.
30%
Valtatton Emor
MARRAA
Perc 50D
Mak IM
Wiidd tettoo Lnce
25%
30%
Cocktal 1500
CARS
30%
Cookia to
CAMS
20%
bhi-idem
20%
PAD
15%
30%
10%
Epoch
Figure 7: Validation Error on SST2 (GLUE) with BERT of 100M parameters. Left: compression rate 0.2%; Middle: 2%; Right: SADL with communication costs {8 10 x 10 10 4 x 10º} Bytes/epoch. The legend 4 10t represents training in the ambient dimension without sketching. Higher compression rate improves the convergence rate and all compression rates achieve comparable test errors at the end of training.
1%, their performance significantly degrades when the compression rate is lowered to 0.1%. On the other hand, our SADL framework has consistent performance across different communication budgets. Notably, we observe that all algorithms, including SADL and the baselines, exhibit stable behavior during training. We report the statistical significance level of the ResNet experiment in the following table. Each experiment is repeated three times under different random seeds, with sources of randomness including client data partitioning, data shuffling. and sketching. Table 8 presents the validation error during the last communication round, with the values after = indicating 3 standard deviation across the runs. Ibit-Adam and PAQ has fixed compression rate, which amounts to 1.6 10t Bytes/epoch.
ResNet
BERT
4e2
40%
2e3
80%
4e3
2e4
4e4
Test TestErro Error
70%
4e5
Test Test Error
30%
2e6
20%
50%
40%
10%
20
40
60
80
100
5
10
15
20
25
Epoch
Epoch
Figure 8: Comparing the performance of tiny sketch sizes on ResNet and BERT. The experiment settings are the same as in Fig. [ and Fig. The legends represent sketch sizes. principle, an extremely tiny sketch size (with 400 in vision tasks and 2000 in language tasks) still converges at similar rate but generates an unfavorable local minima that hardly generalizes.
39

---

1.6 Bytes/epoch
1.6 x 10" Bytes/epoch Sectch_500
50%
Sictch_560
SADL
SADL
Validation Error
Validation Error
40%
40%
30%
30%
20%
20%
0
100
200
300
400
500
100
200
300
400
500
Epoch
Epoch
Figure 9: Validation Error on CIFAR-10. Comparing Sketch-SGD and SADL. The experimental SADL achieves faster convergence and better model performance.
settings are the same as in Fig.
8 x 10 Bytes/epoch
10 Bytes/epoch
25%
Sketch_ @GD
SADL
SADL
Validation Error
Validation Error
20%
20%
15%
15%
10%
10%
30
10
Epach
Epach
Figure 10: Validation Error on SST. Comparing Sketch-SGD and SADL. The experimental settings are the same as in Fig. 3. SADL achieves faster convergence and better model performance.
Next, we present the empirical comparison on using different server optimizers. We present the comparison the three benchmarks in Fig, [9] and proposed in Song et al.] 2023); Shrivastava et al. (024)). We show that using Adam as the server optimizer consistently outperforms in all communication budgets. We also provide the experimental results of running AMSGrad on all three experimental settings the paper. Table 9 10 and [T present the validation errors during the last communication round, with the values after + indicating 3 standaadd deviation across three independent runs. We can observe that Adam and AMSGrad has comparable performance on the all experimental settings, indicating that SADL is insensitive the choice of optimizers.
and I1 The sketch-SGD algorithm is the same as
Additionally, include more challenging set of experiments that considers data heterogeneity and a larger number of clients in Fig. [2] We train a ViT-Base model on CIFAR 100 dataset. The FL system involves 30 clients (the number is maximized for frror-feedback based methods under hardware constraints). The data is distributed in a heterogeneous way under Dirichlet distribution with Dirichlet prior 0.1. We consider the compression rate 0.1 11 and 1%, in the main paper. The validation accuracy along the training process is reported in Fig. 12 It can be observed that SADL achieves the best validation accuracy and converges fast.
To verify Assumption we plot the full Hessian eigenspectrum throughout the training process in Fig. 5] We used stochastic lanczos algorithm implemented by the pyHessian library Yao et al. (2020) to approximate the distribution the full eigenspectrum. Our main claim in Assumption is that the
Hyperparameter
CocktailSGD
CAMS
Ibit-Adam
PAQ
Server learning rate
0.0001
0.001
0.0001
0.001
Server momentum
0.0
0.9
0.9
0.9
Server weight decay
0.0
0.005
0.005
0.005
Client learning rate
0.01
0.01
0.01
0.01
Client momentum
0.9
0.0
0.0
0.0
Table 7: Hyperparameters of CocktailSGDD, CAMS, Ibit-Adam and PAQ in VIT.
40

---

3 x 10 Bytes/epoch
10
Bytes/epach
4%
SADL
5ketch_5GD
2%
Sketch
SADL
Validalion Error
3%
Validation Error
2%
2%
2%
2%
2%
1%
10
20
25
35
40
10
15
20
25
30
35
Epoch
Epoch
Figure 11: Validation Error on fine-tuning VIT with CIFAR-10 dataset. Comparing Sketch-SGD and SADL. The experimental settings are the same as in Fig. 2 comparable performance.
SADL achieves faster convergence and
3e5 Bytes/epoch
3e6 Bytes/epoch
SADL
SADL
CDAcam
80%
CDAdam
80%
FetchSGD
Validation Error
MARINA
Validation Ertor
MARINA
Cockita
ISGD
60%
CocktaISGD
60%
CAMS
1bli-Adam
40%
40%
PAD
20%
20%
100 Figure 12: Validation Error of fine-tuning ViT-base (with 86M parameters) with CIFAR- 100 dataset trained among 30 clients with heterogeneous data distribution. Left: The communication budget 10 Bytes/epoch (0.1% compression rate). Right: The communication budget is x 10t Bytes/epoch (1% compression rate). SADL shows competitive performance under all communication budgets.
40
60
100
Hessian eigenspectrum iterate is summable and the sum is independent of the ambient dimension. which can be satisfied by common distributions, like power-laws. We run testing experiments on ViT-small and train CIFAR-10 dataset, with sketched Adam optimizer. In Fig. [5] we the majority of eigenvalues concentrates near The density enjoys super fast decay with the absolute values of eigenvalues. The decay also holds throughout the training process. This empirical evidence shows the validity of our assumption.
In the main body of the paper, we have achieved 99.9% compression rate and 99.98% compression rate for ResNet and BERT respectively. We further include the results on smaller in Fig. 8 In principle, extremely tiny sketch size (with 400 in vision tasks and 2000 in language tasks) still converges but generates an unfavorable local minima that hardly generalizes.
Bytes/epoch
1.6 10
1.6 105
SADL
(16.9 2.7.82..0)%
CDAdam
(27.7+3.1)%%
(.5.4+
CAMS
(41.2
(26.8
CocktailSGD
(47.4 2.3)%
(31.1
MARINA
(24.4 + 2.8)%
(2.88%4.1)%
FetchSGD
(36.2 2.4)%
(.1.2%
Ibit-Adam
N/A
( 16.8)%
PAQ
N/A
+ (
Table 8: Statistical Significance. The standard deviations computed from independent
the ResNet experiment.
41

---

Bytes/epoch
1.6 10
1.6 10"
Adam
17.8 2.0)%
16.9
AMSGrad
1.56%
(17.3 + 2.1)%
optimizer in SADL on the Table 9: Comparing validation errors of AMSGrad and Adam ResNet experiment.
Bytes/epoch
3 x 10
x 106
Adam
0.11%
(1.6 0.1)%
AMSGrad
0.3%%
(1.6 + 0.3)%
Table 10: Comparing validation errors of AMSGrad and Adam as server optimizer in SADL on the VIT experiment.
Additionally, summarize the theoretical guarantees of the existing approaches in Table 12 From the table, we can see all the comparisons made the main paper are fair.
Discussions
Limitations. The major focus of this work is distributed deep learning, which is subset of the broader distributed learning setting. The specific scenario is of particular interest, since the models involved in modern deep learning are extremely large (typically in billions of parameters), leading prohibitively high communication costs. Our theoretical analysis is based on the specific geometric structure in distributed deep learning. In extremely pessimistic cases conventional optimization problems, the benefits of combining sketching and adaptive optimization require further analysis.
Broader Impacts. This work will facilitate distributed learning in achieving higher communication efficiency. The proposed SADL framework is especially beneficial for resource-cosstrained scenarios and contributes more accessible distributed learning systems. However, the deployment of distributed training may also incur malicious attacks and data poisoning, which will require more attention in the future.
Bytes/epoch 8 10"
x 10"
Adam AMSGrad
0.27%
(8.1 :: 0.3)%
0.2)%
(8.7 :t 0.3)%
Table 11: Comparing validation errors of AMSGrad and Adam server optimizer in SADL on the
SST experiment.
42

---

Algorithms
Communication Bits
learning rate
Convergence Rate
FetchSGD
(
O(1/V)) +T + dd)) O(L/Vd)
O(/vT) (A)
CocktailSGD
0(1)
O(/vT de ((r)a/3) O(Va/vT) 0(/VV))
CD-Adam
0(1)
Onebit-Adam
((d)
0(1/V))
MARINA
0(1)
(1+ O(1/VT)
0( #0 1)/T) (B)
FedCAMS
0(1)
O(d/VT)
PAQ
((d)
O(1/VT)
0(1/VT)
Ours
0(1)
0(1/V))
0(//T)) (C)
Table 12: Comparison on Theoretical Guarantees. We only include the dependence d and T.
(A) Needs heavy-hitter assumption, otherwise deteriorated O(T2/3), (B) The rate is achieved either under deterministic case or use variance reduction methods. w is typically ((d/b) when the compressor is RandK or lz-quantization. (C) requires the assumption on the fast-decay Hessian eigenspectrum. Otherwise, the convergence rate can deteriorate to ((d/VT) under dimension- independent learning rate.
43


Reviews:
[{"review_id": "OusaGSuTSI", "reviewer": "Reviewer_Dwby", "summary": "This paper studies the problem of distributed optimization with adaptive optimizers applied over compressed local updates from local workers. There are multiple sources of errors that the authors take into account: stochasticity of the problem, extra noise from sketched updates, and the local nature of training. With all of them considered, the authors derive convergence guarantees in the nonconvex setting and study the obtained method empirically on small vision and language tasks.", "questions": "\"Therefore, SADL can achieve an $O(1/T)$ convergence for small $T$, which yields faster convergence rate than non-adaptive methods.\"  \nI don't understand this claim. The right-hand side includes terms that scale as $\\sqrt{T}$ and $T$, how exactly would this give $O(1/T)$ convergence rate. How do the authors obtain $O(1/T)$ result?", "limitations": "yes", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "The paper's main strength is the technical depth of the study. The authors took a hard problem, derived convergence guarantees, and performed numerical experiments to show how the method works in practice.  \nIt is also worth mentioning that the authors tried their best to provide the most rigorous theoretical results and explained well how one can avoid making any assumptions on boundedness of the gradients or iterates.\n\nThe main weakness of the work is, in my opinion, limited motivation for all the complexity. The adaptive methods are not well understood (e.g. no clean theory exists to explain what exactly makes rmsprop/adam so good) and their theory isn't particularly attractive for extending to more complicated settings. Most other components such as local updates, compression, and momentum are well understood, but it still remains unclear how much insight we get from combining them all together.  \nOverall, this paper seems to suggest that the combination makes sense to try to use and see if it works well, but this is often the case already with different optimizers: we need to try them in practice to know if they work even if there is theory supporting them. That could be studied by the authors, but their numerical evaluations are rather limited. A good example of how a study like this could be done is the recent work of Douillard et al. \"Streaming DiLoCo with overlapping communication: Towards a Distributed Free Lunch\", which the authors seem to now know, but it provides state of the art in distributed optimization.\n\nMinor things to fix:  \nTypo in line 78: \"in both the the first moments\"  \nLine 206: \"left hand side\" -> \"left-hand side\"", "was_revised": true, "final_justification": "I read the authors' feedback and the other review. As a primarily theoretical work, I think it suffers from several issues, especially in terms of complexity, which is why I remain on the reject side.", "rebuttal": "We sincerely thank the reviewer for taking the time to provide a detailed and thoughtful review of our paper. We appreciate all the positive comments regarding our contributions, and are especially grateful that the reviewer acknowledges the technical depth of our work. We would like to provide detailed responses regarding the weakness part and questions. \n\n- **Q1: \"limited motivation for all the complexity\"**\n\nWhile the SADL algorithm, as described in Alg. 1, is simple and concrete, the analysis is indeed involved since introducing sketching into adaptive distributed optimization might cause conflicts in various components of optimization, such as the preconditioning and the additional source of sketching noise. The **key insight** of our work is that both the convergence rate and the communication cost of SADL depends only **logarithmically** on the ambient dimension $d$, which has an exponential improvement over the existing bound which is **linear** in $d$. The result benefits from our novel analysis on the terms that could be directly affected by the use of sketching.\n\n1. The convergence rate is sensitive to the noise from every aspects of the optimization procedure. Reducing the communication cost unavoidably introduces **larger noise**, the norm of which is in-expectation **linear** to $d$. Our analysis manages to address the possible conflicts in the existing analysis of adaptive distributed deep learning by demonstrating that SADL can achieve both sufficient descent (Lemma 2.3 and 2.5) and that the aggregate sketching noise remains bounded (Lemma 2.6). \n2. Second, the use of adaptive optimizers poses unique challenges in leveraging the **anisotropic** structure (Definition 2.1) of deep learning optimization in the analysis. Adaptive optimizers naturally re-scales the update step to be **isotropic** over the ambient dimension via preconditioning. Our sharp analysis shows that the adaptive steps preserve the structure of the sketching noise so that the variance term can be bounded (Lemma 2.4).\n\nIn summary, our work provides the _first theoretical analysis_ that **sketching is compatible with all components of off-the-shelf adaptive distributed deep learning** methods. Each term in the convergence bound depends at most logarithmically on $d$ and we are the first to show this key insight. \n\n- **Q2: \"numerical evaluations are rather limited\"**\n\nWe would like to emphasize that our work primarily focuses on _distributed deep learning_ **theory**. The numerical experiments are included to illustrate and support our novel theoretical findings.\n\nWe appreciate the reviewer's pointer to [1]. Our main focus is on reducing the communication cost while preserving convergence guarantees. In comparison, [1] considers other aspects of distributed learning, such as synchronization mechanisms, which fall outside the scope of our paper. The experiment in [1] features simulations of compute utilization metric with varying bandwidths. However, such experimental setting is not suitable to demonstrate our theoretical findings. Since our algorithm and all baseline methods share the same synchronization mechanism, the effect of network bandwidth impacts all methods equally when the per-round communication cost is the same. \n\nPapers like [1] are highly sophisticated for real-world deployment and involve various building blocks, while we are more focused on the core distributed learning algorithm with reduced communication. We believe that our proposed **SADL can serve as a foundational component** for future work, including works such as [1].\n\n- **Q3: SADL's $O(1/T)$ convergence for small $T$.**\n\nThe statement is made from Corollary 2 for _small $T$_, i.e., $T \\le O(1/\\epsilon^2)$, and the right-hand-side of the bound only involves $O(1/T)$ terms.\n\nWe thank you again for providing the constructive reviews to our work. We hope our responses have provided clarity and addressed your concerns. We look forward to engaging with you in the discussion period. \n\n[1] Douillard, Arthur, et al. \"Streaming diloco with overlapping communication: Towards a distributed free lunch.\" arXiv preprint arXiv:2501.18512 (2025)."}, {"review_id": "uhCubURJCk", "reviewer": "Reviewer_FZNX", "summary": "The authors introduce Sketched Adaptive Distributed Learning (SADL) algorithms. SADL algorithms operate in a distributed learning setting with C clients, where AMSGrad is the outer optimizer, SGD is the inner optimizer, and a sketching operator is used to reduce the size of tensors communicated. The authors provide a new sharp analysis of convergence for SADL with high probability convergence guarantees. The authors find that a sketching dimension logarithmic in the number of model parameters yields a sub-linear convergence rate that only depends on the loss Hessian’s intrinsic dimension. The authors also show that their analysis holds without assuming that gradient norms are bounded, unlike previous work. Finally, the authors provide empirical experiments showing the performance of SADL when compared to other communication-efficient optimizers.", "questions": "- Since SOTA models are typically trained with gradient clipping, assuming that gradients are bounded seems realistic. Could you comment on this and explain why the results in section 3 are still important? \n- How many Hyperparameters were swept for each method presented in figures 1,2, and 3?\n- Are all methods communication-bandwidth matched in figures 1, 2, and 3? For example, sketching is ring all-reduce compatible, while top-k sparsification requires an all-gather.", "limitations": "Yes, for the theoretical portion of the paper, but there are some limitations that I highlight in the weaknesses that are not currently addressed in the paper.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths**:\n- The paper does a good job introducing the different theoretical results, placing them within the context of the existing literature, and explaining what is new about the theory shown.\n- The theoretical analysis is novel.\n- To the extend of my limited understanding of convergence analysis proofs (I am not intimately familiar with the proofs from existing work, but I did read through the math in the main text), the author’s theorems seem sound and to the best of my knowledge, they support what is mentioned in the implications sections/contributions.\n- The empirical evaluation covers both image and language tasks with transformers, a key setting to evaluate Adaptive optimizers.\n\n**Weaknesses**: \n- **Hyperparameters** While I was able to find the different hyperparameter values used in the appendix, I could not find the number of hyperparameter configurations swept for each optimizer. Given the stochasticity of deep learning, it is crucial to report the hyperparameters swept for each method to demonstrate the performance of different methods under comparable HP budgets.\n\n- **Handling randomness** Following from my first concern above, I find that none of the loss/error curves report error bars. In question 7 of the checklist, the authors mention that “All random seeds are fixed and shared by all experiments”. Simply matching random seeds is not sufficient to control for error in empirical experiments. For instance, a particular method could perform outstandingly better under a particular random initialization and data ordering. I believe that averaging across multiple random seeds is necessary to accurately report the performance of the method. \n\n- **Missing baseline** While AMSGrad was introduced by [1] to ensure the convergence of Adam, the algorithm is rarely used in practice since Adam has less memory overhead and converges when appropriately tuned. I believe this might also be the case for SADL, so including a baseline with Adam as the outer optimizer in the empirical experiments (e.g. a FedAdam [2] baseline) with sketching makes sense and would be beneficial for practitioners. This would also help address my final concern, especially if it can be shown to perform similarly to SADL.\n\n- **misleading writing** The authors claim to introduce “a family of Sketched Adaptive Distributed Learning (SADL) algorithms which can use suitable unbiased gradient sketching for compression with suitable adaptive optimization algorithms” (9-11). This sentence seems misleading as it makes it seem like SADL is broader than a single outer optimizer. However, my understanding is that results have been shown only in the case of an AMSGrad outer optimizer.\n\n---\n**References**\n\n[1] [ON THE CONVERGENCE OF ADAM AND BEYOND; ICLR 2018]\n\n[2] [Adaptive Federated Optimization; NeurIPS 2021]", "was_revised": true, "final_justification": "I have decided to raise by score as the authors have addressed many of my concerns during the rebuttal phase. I recommend accepting the paper, with the caveat that my understanding of the author's theoretical contributions within the context of the existing literature is limited (e.g., the AC should rely on other reviewers for that).\n\n*Issues resolved*\n- reporting the hyperparameters that were swept and ensuring a fair HP budget across methods\n- the addition of error bars\n- results with AMSGrad\n- Clarification of AMSGrad or Adam used in the emprirical study\n- Clarification about gradient clipping", "rebuttal": "We sincerely thank the reviewer for taking the time to provide a detailed and thoughtful review of our paper. We appreciate all the positive comments regarding our contributions, and are especially grateful that the reviewer acknowledges the novelty and soundness of our thoery works. We would like to provide detailed responses regarding the weakness part and questions. \n\n- **Q1: Hyperparameters.**\n\nWe explore a wide range of hyperparameters in our experiments. Specifically, we sweep the server learning rate over the set $[10^{-4}, 5\\times 10^{-4}, 10^{-3}, 5 \\times 10^{-3}, 10^{-2}, 5 \\times 10^{-2}, 0.1, 0.5, 1.0, 2.0]$. The algorithms share all the other hyper-parameters within each set of experiments. The only exception is CocktailSGD since it's unclear  how to incorporate server-side momentum in their framework. To ensure a fair comparison, we tune both server learning rate and the client side momentum for CocktailSGD.\n\n- **Q2: Handling randomness.**\n\nWe observe that virtually all algorithms, including SADL and the baselines, **exhibit stable behavior** during training. We report the statistical significance level of the ResNet experiment in the following table. Each experiment is repeated three times under different random seeds, with sources of randomness including client data partitioning, data shuffling, and sketching. The table presents the validation error during the last communication round, with the values following $\\pm$ indicating $3\\times$standard deviation across the runs. 1bit-Adam and PAQ has a fixed compression rate, which amounts to $1.6 \\times 10^6$ Bytes/epoch.\n\n| Bytes / epoch         | $1.6 \\times 10^5$ | $1.6 \\times 10^6$ |\n|----------|----------|----------|\n| **SADL**       | $( 17.8 \\pm 2.0)$%   | $(16.9\\pm 0.8)$%   |\n| CDAdam   | $(27.7 \\pm 3.1)$%   | $(25.4\\pm 1.9)$% |\n| CAMS   | $(41.2 \\pm 4.9)$%   | $(26.8\\pm 1.0)$%   |\n| CocktailSGD   | $(47.4\\pm 2.3)$%   | $(31.1\\pm 1.7)$%   |\n| MARINA   | $(24.4\\pm 2.8)$%   | $(24.8\\pm 4.1)$%  |\n| FetchSGD   | $(36.2\\pm 2.4)$%   | $(28.2\\pm 3.1)$%  |\n| 1bit-Adam   | N/A   |  $(32.3\\pm 16.8)$%  |\n| PAQ   | N/A  |$( 24.6\\pm 1.7)$%  |\n\n\n- **Q3: \"misleading writing\" and AMSGrad vs Adam.**\n\nWe agree with the reviewer that Adam has been more popular than AMSGrad in practice, and that actually motivates us to choose **Adam as the server optimizer** in all experiments. We will highlight this aspect in the main paper. \n\nFor completeness, we provide the experimental results of running AMSGrad on all three experimental settings in the paper. The following table presents the validation error during the last communication round, with the values following $\\pm$ indicating $3\\times$standard deviation across three independent runs. We can observe that Adam and AMSGrad has comparable performance on all experimental settings, indicating that SADL is insensitive to the choice of optimizers.\n\n|          | ResNet ($1.6 \\times 10^5$) | ResNet ($1.6 \\times 10^6$) |\n|----------|----------|----------|\n| Adam     | $(17.8 \\pm 2.0)$%   | $(16.9 \\pm 0.8)$%  |\n| AMSGrad   | $(18.6 \\pm 1.5)$%   | $(17.3 \\pm 2.1)$%  |\n\n|          | ViT ($3 \\times 10^5$) | ViT ($3 \\times 10^6$) |\n|----------|----------|----------|\n| Adam     |  $(2.1 \\pm 0.1)$%   | $(1.6 \\pm 0.1)$%  |\n| AMSGrad   | $(1.8 \\pm 0.3)$%   | $(1.6\\pm 0.3)$%  |\n\n\n|          | SST ($8 \\times 10^5$) | SST ($8 \\times 10^6$) |\n|----------|----------|----------|\n| Adam     | $(8.7 \\pm 0.2)$%   | $(8.1\\pm 0.3)$%  |\n| AMSGrad   | $(8.8 \\pm 0.2)$%   |$(8.7\\pm 0.3)$%  |\n\nIn addition, since the mainstream adaptive optimizers share similar designs in utilizing momentum and preconditioners, the theory we developed can also be easily extended to Adam and other adaptive optimizers.\n\n- **Q4: Avoiding bounded gradient assumption and gradient clipping.**\n\nWe agree that the widely used gradient clipping methods guarantee that the gradients are bounded. However, we incorporate the analysis in Section 3 without assuming bounded gradients or using gradient clipping for the following reasons:\n1. Establishing the convergence result without assuming bounded gradients arguably **broadens the applicability** of SADL, particularly for the cases when the objective function is squared loss, which naturally leads to unbounded gradients.\n2. We consider the derivation in Section 3 to be a *technically novel* and **potentially valuable contribution**. We believe it may be of independent interest to the broader optimization and machine learning community.\n3. While gradient clipping effectively addresses the issue of unbounded gradients in practice, it also **introduces bias** to the gradients, which might pose challenges in the theoretical analysis. Understanding and bounding such clipping-induced bias is an important direction of research, especially for the differential privacy community.\n\n- **Q5: Communication bandwidth.**\n\nWe thank the reviewer for the insightful comment. In this optimization focused work, we adopt the classical framework where a central parameter server aggregates local updates, as in MARINA[1]. Therefore, our experimental results (Figure 1-3) align the number of parameters (in bytes) that each worker transmits to the server. This alignment protocol follows common practices in prior literature [2, 3]. Furthermore, the adopted framework is general and can be seamlessly applied to federated learning scenarios where direct peer-to-peer communication is often infeasible. \n\nWe acknowledge that SADL can benefit from ring all-reduce to mitigate network contentions and improve bandwidth utilization. Practical deployment of SADL will depend on the specific system setting and is an important consideration for industrial applications.\n\nWe thank you again for providing the constructive reviews to our work. We hope our responses have provided clarity and addressed your concerns. We look forward to engaging with you in the discussion period. \n\n[1] Gorbunov, Eduard, et al. \"MARINA: Faster non-convex distributed learning with compression.\" International Conference on Machine Learning. PMLR, 2021.\n\n[2] Stich, Sebastian U. \"Local SGD converges fast and communicates little.\" arXiv preprint arXiv:1805.09767 (2018).\n\n[3] Reddi, Sashank, et al. \"Adaptive federated optimization.\" arXiv preprint arXiv:2003.00295 (2020)."}, {"review_id": "hdTx9Hv37B", "reviewer": "Reviewer_iDpm", "summary": "In this paper, authors propose the Sketched Adaptive Distributed Learning (SADL) algorithm to combine gradient sketching with adaptive optimizers for reduced communication overhead and faster convergence. The main contributions lie in that authors provide rigorous theoretical proof on the convergence guarantee of the proposed method. Such proof does not rely on previously commonly used assumption of bounded gradient norm; instead, they prove that gradients are implicitly bounded along the optimization trajectory. In addition, they explicitly consider the impact of three types of noises, namely local mini-batch training, sketch-related compression error, and aggregated noise over training horizon. Finally, authors demonstrate the effectiveness of the proposed SADL algorithm for optimizing ResNet, ViT, and BERT models of 40-100 million parameters.", "questions": "1. Section 4, line 303. It remains unclear which gradient sketching operator (i.i.d. isotrpoic Gaussian, subsampled randomized Hadamard transform, and count-sketch) and adaptive optimization algorithm is used by SADL in these experiments.\n2. Section 4. The ablation study on different combinations of gradient sketching operators and adaptive optimization algorithms is missing. What is the recommended combinations between these two componenets?\n3. Is it possible to further extend the proposed method into asynchronous distributed learning setting? This will be significant different from the current setup, and is well out-of-the-scope of this paper, but some brief discussion on this direction might be valuable in the rebuttal.", "limitations": "The empirical results section may be further improved to include detailed ablation study on different combinations of gradient sketching operators and adaptive optimization algorithms.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 4, "strengths_and_weaknesses": "Strengths:\n* The proposed SADL algorithm combines gradient sketching with adaptive optimization algorithm. This is more closely related to real-world distributed learning scenarios, where adaptive optimizers rather than vanilla SGD are adopted. This also represents a huge improvement compared to other gradient compression methods that build their convergence analyses based on vanilla SGD.\n* Instead of providing in-expectation analysis of different sources of noise (stochastic gradient, compression error, and sequential dependencies), authors provide high probability convergence bounds as a more convincing guarantee of stable training in practice.\n* The convergence rate of the proposed SADL algorithm depends over logarithmically on the ambient dimension $d$ (number of model parameters), which improves over previous results of linear dependency on the ambient dimension. This is established upon the intrinsic dimension of the loss Hessian, which is significantlly smaller than the ambient dimension as proved in several related researches.\n* The proposed method works well even with extremenly low communication budget, and there is no convergence issue under any sketch size as reported in the experiments.\n\nWeaknesses:\n* Limited experimental results for larger models. Since the bottleneck effect of communcation overhead in distributed learning is more severe for models with larger number of parameters, e.g., billion-scale, it would be nice to see whether the proposed algorithm also works well for such models. However, considering the limited time and possibly insufficient computation resources, it is totally fine to not include such results in the rebuttal.\n\nNotes: I did not fully go through all the detailed theoretical proofs in the appendix, and may not be qualified to justify the correctness of these proofs.", "was_revised": true, "final_justification": "The proposed SADL algorithm combines gradient sketching with adaptive optimization algorithm. This is more closely related to real-world distributed learning scenarios, where adaptive optimizers rather than vanilla SGD are adopted. This also represents a huge improvement compared to other gradient compression methods that build their convergence analyses based on vanilla SGD and/or in-expectation analysis. The proposed method works well even with extremenly low communication budget, and there is no convergence issue under any sketch size as reported in the experiments. Authors have also addressed my previous concerns in the rebuttal.", "rebuttal": "We would like to extend our sincere appreciation to the reviewer for all the constructive and positive feedback on our contribution. We are especially grateful that the reviewer acknowledges technical strength of our work in both theory and experiments. We would like to provide detailed responses regarding the weakness part and questions. \n\n- **Q1: Choice of sketching algorithms and optimizers.**\n\nIn all experiments, we choose the combination of Adam as the server optimizer and SRHT as the sketching algorithm. Isotropic Gaussian requires creating a dense matrix of size $b \\times d$, where $b$ is the sketch size and $d$ is the ambient dimension, which is not memory-efficient, so we do not use that. \n\nWe extend our experimental results under the combination of [Adam, AMSGrad] $\\times$ [SRHT, CountSketch] and we present the ablation studies for the ResNet experiment in the following tables. The table presents the validation error during the last communication round.\n\nWhen the communcation cost per-round is $1.6 \\times 10^5$ Bytes/epoch,\n|          | SRHT | CountSketch |\n|----------|----------|----------|\n| Adam     | $17.0$%   | $23.8$% |\n| AMSGrad   | $18.1$%   | $25.1$% |\n\nWhen the communcation cost per-round is $1.6 \\times 10^6$ Bytes/epoch,\n|          |  SRHT | CountSketch  |\n|----------|----------|----------|\n| Adam     | $16.7$%   | $23.6$% |\n| AMSGrad   |  $17.0$% | $24.7$% |\n\nRegarding the choice of server optimizer, Adam achieves slightly better performance than AMSGrad. Regarding the choice of sketching methods, SRHT consistently outperforms CountSketch, while the latter is still competitive with other baseline methods in Fig. 1. This observation aligns with the error bound of CountSketch (Lemma B.3), which differs by a $1/\\sqrt{b}$ factor with SRHT (Lemma B.1). In summary, Adam + SRHT will be an ideal combination in practice. \n\n\n- **Q2: Extension to larger models and asynchronous training paradigms.**\n\nWe appreciate the thoughtful suggestion on the possible extension of our work. Deploying SADL in the wild is definitely an important long-term goal for us. \n\nFirst, we would like to extend our experimental settings to larger models, as suggested by the reviewer. However, several baseline methods that involve error feedbacks face OOM (out of memory) issues. This is because error feedback requires storing error terms which are of the same size as the full model. We finetune a Deepseek-R1-coder model (1.3B parameters) on the Spider dataset (a complex text-to-SQL generation task) and report the validation accuracy in the following table. We select $b=10^7$, which amounts to 1% compression rate. The only two algorithms that survive the experimental settings are SADL and MARINA. We can observe that SADL outperforms and shows high potentials in the compatibility with large models.\n\n|          |  Accuracy |\n|----------|----------|\n| SADL     |  54.7% |\n| MARINA   |  48.4% |\n\n\nSecond, mechanism-wise, SADL naturally extends to asynchronous settings. Our SADL framework can be combined with the generic asynchronous training algorithms such as the one described in Algorithm 1 of [1]. Specifically, lines 3 and 6 in the algorithm block can be modified to incorporate our sketching and desketching operators respectively so that the stochastic gradient is compressed before the transmission. Developing a theoretical analysis on such algorithm will be an interesting extension of our current work.\n\nWe thank you again for providing the constructive reviews to our work. We hope our responses have provided clarity and addressed your concerns. We look forward to engaging with you in the discussion period. \n\n[1] Stich, Sebastian, Amirkeivan Mohtashami, and Martin Jaggi. \"Critical parameters for scalable distributed learning with large batches and asynchronous updates.\" International Conference on Artificial Intelligence and Statistics. PMLR, 2021."}, {"review_id": "6tqeLtoPgE", "reviewer": "Reviewer_UgYH", "summary": "This paper investigates sharper bounds for sketching for distributed optimization, under AMSGrad. A key contribution is to avoid reliance on d, the ambient dimension, and instead reduce to log d. This critically depends on the assumption of low intrinsic dimension, which is reasonable for many large DNNs.", "questions": "- the \"third source of noise\", is that really noise or bounding error? is it in practice or purely in the proof? \n\n - The idea of the limited intrinsic dimension is interesting, and I agree that it is likely true in practice. But can you discuss how baseline / competing sketching bounds compare to yours, if they too adopted that assumption? How much is this lack of reliance on $d$ a direct result of this assumption?", "limitations": "yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "This paper seems to have substantial overlap with Song et al 2023; e.g. a lot of key theorems, and even the overall philosophy, seem borrowed from it. I guess this paper gives the novelty of investigating AMSGrad rather than gradient descent, which is significant, but my comment is that there should be a deeper discussion/review of this prior paper in the introduction, so that this novelty is clear. It also helps justify why the authors are so obsessed with AMSGrad, by showing that the GD method was already previously studied. \n\nOverall, on a surface level, it seems a good contribution, although I will rely on others to comment on novelty.\n\nMore comments below:\n\nIt would be good ot have a more extensive discussion of the baseline algorithms\n\n\nIn terms of mathematical correctness, there were too many pages to check everything carefully. I did look through C1 and immediately found a few (likely minor) mistakes, e.g.\n\n - line 897: sign error in middle term\n - line 898: disappearing $\\kappa V_t^{-1/2} m_{t-1}$ term, extra $R_t$ in front of $g_{t,k}^c$ (sketched variable and gradient should have same dimension)\n - line 900: missing nabla (definitely minor)\n\nHowever, overall, the claims seem reasonable and I did not find anything too damaging in the overall logical flow of the proofs. \n\n\ntypos:\n - alg 1: parameters\n - Line 339: capitalization\n - 846: spacing\n- Line 867: spacing\n\nAlthough it is ok to give the results of Song et al 2023, it would be good to specify clearly the constants hidden in the $\\Theta$ rate.", "was_revised": true, "final_justification": "Looks good, I'll keep my score. Thanks for the discussion, it was very interesting to read.", "rebuttal": "We would like to extend our sincere appreciation to the reviewer for all the constructive and positive feedback on our contribution. We are especially grateful that the reviewer acknowledges the significance of our theoretical advancements. We would like to provide detailed responses regarding the weakness part and questions. \n\n- **Q1: Relationship with (Song et al. 2023).**\n\nSong et al. 2023 is an inspiring work that introduces the sketch-and-desketch framework to distributed learning. Our algorithm is a natural extension of the framework with **adaptive optimization**, with necessary adaptations to incorporate sketching in the precondtioners. Our theoretical analysis is grounded in the properties of standard sketching methods, as also outlined in Song et al 2023.\n\nHowever, the convergence bound derived in Song et al. is limited since _their communication cost scales **linearly** with the number of model parameters $d$_, i.e., $O(d)$, which is also there in prior results using other compression algorithms. \n\nThe major advance in our work is to show an **exponential improvement** in the communication cost from $O(d)$ to $O(\\log d)$, i.e., logarithmic in the number of parameters. The improvement is especially important for modern deep learning models with large number of parameters.\n\nDeriving the theoretical result in the context of _adaptive_ distributed learning is technically challenging, which requires more careful analysis on the effect of sketching in both the precondtioner term and noise terms. For the discussion on the preconditer term, please refer to our response to Q2 (below). For the analysis on noise terms, which is also not covered in Song et al 2023 or prior works, please refer to our response to Q3 (below). \n\nWe will add a discussion in the introduction to properly give credit to Song et al., 2023 and highlight the connection and difference between our work and theirs.\n\n\n- **Q2: Significance of analyzing adaptive methods.**\n\nThe use of adaptive optimizers poses unique challenges in leveraging the **anisotropic** structure (Definition 2.1) of deep learning optimization in the analysis. Adaptive optimizers naturally re-scales the update step to be **isotropic** over the ambient dimension via preconditioning. Our sharp analysis shows that the adaptive steps preserve the structure of the sketching noise so that the variance term can be bounded (Lemma 2.4).\n\n- **Q3: Does the intrinsic dimensionality suffice to derive the convergence bound?**\n\nAlthough intrinsic dimension is one of the key ingredients to derive our convergence bound, we would like to highlight that this property is only used in addressing the second-order term in optimization. \n\nFirst, the convergence rate is sensitive to the noise from every aspect of the optimization procedure. Reducing the communication cost unavoidably introduces **larger noise**, the norm of which is in-expectation **linear** to $d$. Our analysis manages to address the possible conflicts in the existing analysis of adaptive distributed deep learning by demonstrating that SADL can achieve both sufficient descent (Lemma 2.5) and that the aggregate sketching noise remains bounded (Lemma 2.6). Both terms arise from the first-order component in optimization and thus does not depend on the intrinsic dimensionality.\n\nSecond, the current analysis is restricted to un-biased gradient compressors. A vast majority of the existing works (CDAdam, FetchSGD, CAMS etc.) rely on **biased gradient compressors** and error feedback mechanism. The dependence on $d$ in their convergence bound originates from the distortion term (Section C.1 in [1]), where the intrinsic dimensionality is not directly relevant.\n\n- **Q4: Discussion of baseline algorithms.**\n\nWe thank the reviewer for the advice. We have included Table 7 in Appendix that summarizes the communication cost and convergence bounds of each baseline algorithm. We have also discussed and categorized the baselines in Appendix A. We will add a more detailed discussion in the final version of this paper.\n\n- **Q5: The third source of noise.**\n\nThe third source of noise is essentially a temporal aggregation error which exists in practice. Accumulating stochastic noise naturally gives you an $O(\\sqrt{T})$ error by Azuma's inequality. \n\n- **Q6: Typos**\n\nWe appreciate the reviewer’s attention in pointing out the typos. We will carefully review the paper and make multiple passes to correct all identified and remaining typos in the final version.\n\nWe thank you again for providing the constructive reviews to our work. We hope our responses have provided clarity and addressed your concerns. We look forward to engaging with you in the discussion period. \n\n[1] Richtárik, Peter, Igor Sokolov, and Ilyas Fatkhullin. \"EF21: A new, simpler, theoretically better, and practically faster error feedback.\" Advances in Neural Information Processing Systems 34 (2021): 4384-4396."}]

Author rebuttal:

Meta-review: This submission combines sketching with adaptive distributed deep learning methods. 
The authors improve the convergence bound dependce from d to log d compared with other compression methods. After engaging discussion some reviewers raised their scores. While I agree with reviewer Dwby that the paper could be improved and that the dependen on d is unfortunate, I also think that the improvement from d to log d is sufficiently interesting and recommend acceptance.
Decision: Accept (poster)