Sketched Adaptive Distributed Deep Learning: A Sharp Convergence Analysis

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.

• Q1: Relationship with (Song et al. 2023).

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

However, 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.

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

Deriving 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).

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

• Q2: Significance of analyzing adaptive methods.

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

• Q3: Does the intrinsic dimensionality suffice to derive the convergence bound?

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

First, 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.

Second, 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.

• Q4: Discussion of baseline algorithms.

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

• Q5: The third source of noise.

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

• Q6: Typos

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

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

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

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.

• Q1: Choice of sketching algorithms and optimizers.

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

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

When the communcation cost per-round is $1.6 \times 10^5$ Bytes/epoch,

When the communcation cost per-round is $1.6 \times 10^6$ Bytes/epoch,

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 a $1/\sqrt{b}$ factor with SRHT (Lemma B.1). In summary, Adam + SRHT will be an ideal combination in practice.

• Q2: Extension to larger models and asynchronous training paradigms.

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

First, 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.

Second, 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.

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

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

Dear Reviewer iDpm,

We thank you again for your valuable feedback and your endorsement of our work. We hope our response has addressed your questions and further strengthened your confidence in the contributions of our work. Please feel free to let us know if you have any further questions—we would be happy to clarify them.

Submission 23560 Authors

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.

• Q1: Hyperparameters.

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

• Q2: Handling randomness.

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

• Q3: "misleading writing" and AMSGrad vs Adam.

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

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

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

• Q4: Avoiding bounded gradient assumption and gradient clipping.

We 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:

1. 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.
2. 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.
3. 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.

• Q5: Communication bandwidth.

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

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

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

[1] Gorbunov, Eduard, et al. "MARINA: Faster non-convex distributed learning with compression." International Conference on Machine Learning. PMLR, 2021.

[2] Stich, Sebastian U. "Local SGD converges fast and communicates little." arXiv preprint arXiv:1805.09767 (2018).

[3] Reddi, Sashank, et al. "Adaptive federated optimization." arXiv preprint arXiv:2003.00295 (2020).

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.

• Q1: "limited motivation for all the complexity"

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

1. 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).
2. 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).

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

• Q2: "numerical evaluations are rather limited"

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

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

Papers 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].

• Q3: SADL's $O(1/T)$ convergence for small $T$.

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

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

[1] Douillard, Arthur, et al. "Streaming diloco with overlapping communication: Towards a distributed free lunch." arXiv preprint arXiv:2501.18512 (2025).