Paper: Towards a Better Theoretical Understanding of Independent Subnetwork Training (ICLR-2024)

Reviewer comment (Experiments):
**W.1.** The provided experimental results are not conclusive enough and placed in the end of the Appendix.

Author response:
We comment on this in our general comment ([link](https://openreview.net/forum?id=HhVns87e74&noteId=Y3PkEuz0tL)) on experiments to all reviewers.

We would like to highlight that the title of our work is _“Towards a Better **Theoretical** Understanding of IST”_. That is why, in the main part of the paper, we focus on contributions from theoretical analysis. Another reason is the limited space allowed by the conference submission format. If the reviewer believes that moving the experimental results to the main part will strengthen the points we are making, we can do this in the camera-ready version.

Reviewer comment (Experiments):
**Q.1.** Can this work be extended to neural networks such as [1]?  

[1] Dun, Chen, Cameron R. Wolfe, Christopher M. Jermaine, and Anastasios Kyrillidis. "Resist: Layer-wise decomposition of resnets for distributed training." In Uncertainty in Artificial Intelligence, pp. 610-620. PMLR, 2022.

Author response:
As far as we know, convergence of ResNets training is not adequately understood even for much simpler optimization algorithms such as gradient descent. This lack of understanding is compounded in the case of IST due to the biased nature of the gradient estimator and the complexities inherent in distributed training. Given these challenges, a  straightforward extension of our current work to ResNets or similar architectures is not feasible at this stage.

Reviewer comment (Theory):
Section 2.2 lists 3 assumptions for the theoretical analysis. The authors also discuss the necessity of each assumption. Could the authors try to remove one of them? Specifically, is it possible to discuss other problems other than the specific quadratic one?

Author response:
We address this weakness in our general response ([link](https://openreview.net/forum?id=HhVns87e74&noteId=fjCbL9Q2IZ)) to all reviewers regarding assumptions.

We note that the assumption of exact submodel gradient computation can be easily generalized to stochastic estimators with bounded variance. Regarding the choice of a quadratic model. We chose this approach due to the clear theoretical insights it provides, as detailed in our study. Analyses of similar methods with non-quadratic loss functions have often led to restrictive and impractical assumptions, resulting in unsatisfactory convergence bounds, as discussed in Section 4.3 and Appendix C. The complexity and lack of a suitable theoretical framework for a general class of $L$-smooth functions, particularly with challenges like biased gradient estimators, further motivated our decision.

While we acknowledge the value of extending our analysis to non-quadratic settings, the current theoretical landscape and the intricacies of such an extension necessitated our initial focus on the quadratic model. This foundational work sets the stage for future exploration in more complex loss functions as the field advances.

Reviewer comment (Experiments):
The paper focuses on the theoretical aspects of IST. It would be insightful to discuss the practical implications of the findings for real-world applications and provide guidance on effectively utilizing IST in various distributed training scenarios.

Author response:
Thank you for your note on the practical implications of our findings. We discussed it in our general comment ([link](https://openreview.net/forum?id=HhVns87e74&noteId=Y3PkEuz0tL)) on experiments to all reviewers.

Briefly, our study indicates that IST can be highly efficient in both homogeneous and heterogeneous settings, given interpolation conditions typical for large neural networks. However, in the more general case, we found that the convergence of naïve IST is influenced by the level of heterogeneity among the distributed nodes. Specifically, decreasing the learning rate can be especially helpful in heterogeneous scenarios for reducing error.

Reviewer comment (Writing):
At the bottom of Page 2, $\mathbf{R}^d \rightarrow \mathbf{R}$

Author response:
Thank you for pointing this out. We already fixed it in our revision.

Reviewer comment (Presentation):
Figures 1a and 1b have different vertical axis, relative error and absolute error. Could the authors provide both relative and absolute error for these two cases.

Author response:
We would like to note that it was done on purpose. The first (1a) plot illustrates that the method (functional gap) converges to the neighborhood of the solution in the first case. While the second (1b) one shows that the method’s iterates converge to a fixed point different from the optimum. We can provide additional results in the camera-ready version of the paper.

Reviewer comment (Theory):
**Restrictive (maybe impractical) assumption the work uses that assumes performing only one gradient descent step during local training in IST**. Performing only one gradient computation at each node requires very frequent communication with the server, which will incur high communication costs. Hence, performing multiple gradient descent steps in the local mode is more desired as in the IST work [1] and [2]. I admit taking multiple steps may hurt the accuracy, while the main motivation of IST is saving communication costs under the accuracy-efficiency trade-off.

Author response:
We address this comment in our general response ([link](https://openreview.net/forum?id=HhVns87e74&noteId=fjCbL9Q2IZ)) to all reviewers regarding assumptions.

We understand that this assumption might seem restrictive and potentially impractical due to the increased communication costs it could entail. However, it is required to isolate and analyze specific properties of IST more effectively. We would like to highlight that our theory could potentially be extended to scenarios involving multiple local steps, as suggested in the approach of Khaled et al. (2019). However, the current state of research indicates that no conclusive analysis demonstrates the advantages of multiple local steps in the worst-case scenarios. We believe that the insights gained from our approach are crucial for guiding future research that might explore more complex scenarios, including those involving multiple local steps.

Reviewer comment (Theory):
**Why is the gradient sparsification introduced in Eq.6 for IST?** The original IST work formulated the IST training method as Eq(69) without the gradient sparsification operator. The authors state that “it can create a significant disparity between theory and practice” in Appendix D. I didn’t notice that IST uses any form of gradient sparsification, as IST is orthogonal to gradient sparsification techniques. The motivation for introducing this operator in IST formulation is unclear.

Author response:
Thank you for bringing this up! Let us clarify the point.

Gradient sketching is introduced to better represent submodel computations. Note that when the gradient is taken with respect to the submodel $\nabla f(\mathbf{C} x)$ it can result in a non-sparse update, which contradicts the idea of IST. To illustrate this, imagine a logistic regression problem, for which computing a gradient for compressed weights vector (even for zero) may result in a non-sparse gradient (potentially without any zeros at all). That is why we believe that sparsification of the gradient is an essential component of correclty representing IST. Thus, our work improves upon the analysis of the original paper.

Reviewer comment (Theory):
Can the author further discuss why the assumption that performing only one gradient computation is reasonable, as it is not a common choice for efficient IST? It would be better if the author could show the theoretic analysis still holds with a more relaxed assumption, like performing twice or more.

Author response:
We address this comment in our general response ([link](https://openreview.net/forum?id=HhVns87e74&noteId=fjCbL9Q2IZ)) to all reviewers regarding assumptions.

We understand that this assumption might seem restrictive and potentially impractical due to the increased communication costs it could entail. However, it is required to isolate and analyze specific properties of IST more effectively. We would like to highlight that our theory could potentially be extended to scenarios involving multiple local steps, as suggested in the approach of Khaled et al. (2019). However, the current state of research indicates that no conclusive analysis demonstrates the advantages of multiple local steps in the worst-case scenarios. We believe that the insights gained from our approach are crucial for guiding future research that might explore more complex scenarios, including those involving multiple local steps.

Reviewer comment (Theory):
Can the authors provide a clear discussion on the reason for introducing the gradient sparsification operator compared to the original formula? Why is this operator necessary when analyzing the convergence of IST if those two techniques are orthogonal?

Author response:
Thank you for bringing this up! Let us clarify the point.

Gradient sketching is introduced to better represent submodel computations. Note that when the gradient is taken with respect to the submodel $\nabla f(\mathbf{C} x)$ it can result in a non-sparse update, which contradicts the idea of IST. To illustrate this, imagine a logistic regression problem, for which computing a gradient for compressed weights vector (even for zero) may result in a non-sparse gradient (potentially without any zeros at all). That is why we believe that sparsification of the gradient is an essential component of correclty representing IST. Thus, our work improves upon the analysis of the original paper.

Reviewer comment (Theory):
Theoretical understanding seems to be constrained by the assumptions, which might separate the current analysis from the real use cases.

Author response:
We address this comment in our general response ([link](https://openreview.net/forum?id=HhVns87e74&noteId=fjCbL9Q2IZ)) to all reviewers regarding assumptions.

We want to emphasize that in our study, the algorithm and problem setting are intentionally simplified to precisely isolate and understand the unique effects of combining data with model parallelism, especially in the context of distributed submodel gradient computations. We are aware that assumptions we make may deviate from practical settings. However, our goal was not to replicate real-world scenarios closely but to develop a fundamental understanding of certain theoretical aspects of IST.

Reviewer comment (Experiments):
Although the suggested analysis of the convergence and bias sounds interesting and useful, the limited experimental validation would limit the application of the proposed observation in real distributed training scenarios. In particular, the authors have emphasized the need for the theoretical understanding of a wide-spread parallelization and co-design of communication and training algorithms for large-scale training, but the limited validation would hinder the application of the findings from this work.

Author response:
We comment on this Weakness in our general comment ([link](https://openreview.net/forum?id=HhVns87e74&noteId=Y3PkEuz0tL)) on experiments to all reviewers.

In addition, let us stress that the main contribution of our work is from the theory side, and that is why, in our experiments, we focus on well-controlled settings that satisfy the assumptions in our paper to provide evidence that our theory translates into observable predictions. These are well-designed experiments that do support our theory and core claims. Since our results guarantee that the methods work, we do not need to test them extensively on large or complicated datasets and models to show that they do (which is necessary for heuristics unsupported by any theory).

Reviewer comment (Experiments):
Can we see a more realistic distributed training scenario (e.g., training ImageNet deep neural networks) to validate the key observations of this work?

Author response:
We consider extending our experimental results with neural network training for the camera-ready version of the paper. Due to time constraints, it is not possible to provide them during the rebuttal period.