Paper: A Quadratic Synchronization Rule for Distributed Deep Learning (ICLR-2024)

Reviewer comment (Novelty):
The motivation of finding 'H’ given a learning rate schedule is questionable. Both learning rate and H should be tunable hyperparameters in distributed learning. Solely determining H given the learning rate schedule, yet without asking about what the optimal learning rate is, is a less valuable question. While prior work [1] already shows that the generalization benefits arise only if H = Ω( 1/η ) when η → 0, this work specifically sets H to 1/η^2, which is incremental, reinforcing the same results. In fact, the learning rate is of greater significance in deep learning and distributed learning; for many problems, one can boost generalization and convergence performance solely by carefully tuning the learning rate. A more valid question that I believe should be asked is what is the optimal dynamic learning rate schedule under a given schedule of H, e.g. H=1 or 2 or Post-local SGD with variable H. There is no comparison to this learning rate question under a fixed H schedule, which further saves communication and is easier to implement in reality. Therefore, I think the problem studied here to recommend how to set H is kind of artificial.

Author response:
**W1 (1)**: The motivation for finding 'H’ given a learning rate schedule is questionable. Solely determining H given the learning rate schedule, yet without asking about what the optimal learning rate is, is a less valuable question. What the optimal dynamic learning rate schedule is under a given schedule of H, e.g., H=1 or 2 or Post-local SGD with variable H, is a more valid question.

**A**:
1. As clarified in our response to M2, we have explored the most popular and standard learning rate schedulers and carefully tuned the learning rate for each baseline algorithm, showing that QSR can yield consistent generalization benefits across different learning rate schedules. 

2. Tuning learning rates for existing $H$ schedules (as suggested by the reviewer) cannot fully exploit the generalization benefit of local steps. According to our theory and experiments, to maintain the regularization effect, $H$ should grow to a very large value in the late training phase when the learning rate decays to a small value. For example, when training ViT-B with batch size 4096, $H$ can grow to over 8000 at the end of training. However, existing schedules of $H$ typically use very small values even in the late phase (e.g., the largest H the reviewer suggests seems to be 2).

**W1 (2)**: While prior work [1] already shows that the generalization benefits arise only if $H=\Omega(\eta^{-1})$ when $\eta\to 0$, this work specifically sets $H$ to $\eta^{-2}$, which is incremental.

**A**: QSR is novel compared to previous work [1]. Specifically, [1] focuses on the theoretical side, and only studies the setting of **using constant $\eta$ and $H$ throughout training**. No concrete $H$ schedules are proposed in [1]. By contrast, this paper proposes a specific $H$ schedule, QSR, where $H$ dynamically grows in proportion to $\eta^{-2}$ during training as $\eta$ decays over time. The setting of a decaying learning rate schedule studied in this paper is a common choice in practice (e.g., [2-4]), but it is not included in [1].

Reviewer comment (Experiments):
The baseline for comparison is not clear. For instance, what is the Parallel AdamW? How does the Parallel AdamW work? I can’t find the rigorous definition of it given by the paper. Because of the second momentum terms, Parallel AdamW that assumably aggregates the model parameters every H=1 step is not equivalent to Centralized AdamW when only setting H = 1. The second momentum can not be simply averaged. Even the definition of Parallel SGD is not provided. The paper simply says that refers to Local SGD when H=1. At least a reference or a definition should be provided.

Author response:
**W2**: The baseline for comparison is not clear. What is Parallel AdamW/Parallel SGD?

**A**: Parallel SGD/AdamW refers to the standard data parallel implementation of SGD/AdamW, which is introduced in the very first paragraph: at each step, each worker first computes gradients on their own local batches of data. Then, they average the local gradients via an All-Reduce operation and update the model parameter using the averaged gradient and SGD/AdamW optimizer. In the updated version, we have explicitly defined “Parallel OPT” as the data parallel implementation of an optimizer OPT in the first paragraph and included the pseudo-code for Parallel OPT in Appendix B according to the reviewer’s suggestion.

Reviewer comment (Experiments):
Some experimental settings are questionable. What has motivated the adoption of only SGD for ResNet experiments and AdamW only for VIT experiments? Furthermore, while all experiments employ momentum-based optimizers and underscore the effectiveness of QSR with such optimizers, the paper doesn't deeply explore the role of momentum in this methodology.

Author response:
**W3 (1)**:  Some experimental settings are questionable. What has motivated the adoption of only SGD for ResNet experiments and AdamW only for VIT experiments?

**A**: Adopting SGD for ResNets and Adam/AdamW for ViTs is the standard practice (e.g., [5-10]). This enables us to run strong baselines with existing training recipes [6, 9] and showcase that QSR can be seamlessly integrated into standard training recipes for both convolutional and transformer-based models, enhancing both communication efficiency and generalization performance.
It is also worth noting that no standard training recipe is known for training ViTs with SGD. Directly running SGD on transformer-based models typically results in a much slower convergence speed [15].

**W3 (2)**: While all experiments employ momentum-based optimizers and underscore the effectiveness of QSR with such optimizers, the paper doesn't deeply explore the role of momentum in this methodology.

**A**: 

1. Adding momentum is an orthogonal direction to adjusting H for improving optimization and generalization. Our theoretical analysis builds upon the Slow SDE framework originally developed only for vanilla SGD [16]. This is novel and highly nontrivial. Understanding the role of momentum in this method is an interesting and challenging problem we plan to investigate in our future research.

2. Our experimental results on SGD/Local SGD with momentum align with the prediction of our theory, demonstrating that the implication of our theory goes beyond vanilla SGD/Local SGD.

Reviewer comment (Theory):
The theoretical analysis and proof overlap too much with previous work [1]. Both rely on the same SDE technique (which is the novelty of the prior work) for analyzing analogous issues, with the current work lacking additional theoretical insights and development. Additionally, some proof steps are missing.

Author response:
**W4**: The theoretical analysis and proof overlap too much with previous work [1].

**A**: We respectfully disagree with the above comments. The analysis in [1] cannot handle the case where $H \sim \eta^{-2}]$, but our theoretical result can. Specifically, their analysis depends on the condition $H\eta=O(1)$, and many bounds therein explode as $H\eta \to \infty$, e.g., Theorem 3.3, Lemmas I.14 and I.16. In the context of QSR, where $H\eta=\frac{\alpha^2}{\eta}$ goes to infinity as $\eta$ approaches $0$, the condition $H\eta=O(1)$ is violated, rendering the analysis in Gu et al. (2023) ineffective for QSR. To help the reviewer better understand our proof, we have added a proof outline in the updated version.

Reviewer comment (Experiments):
There is no comparison with other practical distributed learning systems and methods, given the vast literature on distributed learning. The comparison is more of a simple self-comparison between different hyperparameter scalings of H, e.g., constant H vs. the QSR of H, assuming the same learning rate schedule is used. This comparison is not convincing because 1) for different H setups, the optimal learning rate schedule will vary; 2) it didn’t compare to other possibilities of H = Ω( 1/η ) making it hard to understand the limitations of the problem domain. This makes the work more of hyperparameter tuning heuristics and lacks insights. The experimental results are more of validating the theorems and lack practical value.

Author response:
**W5**: There is no comparison with other practical distributed learning systems and methods. The comparison is more of a simple self-comparison between different scalings of $H$.  This comparison is unconvincing because 1) for different H setups, the optimal learning rate schedule will vary; 2) it didn’t compare to other possibilities of $H = \Omega(\frac{1}{\eta})$.

**A**: We respectfully disagree with the reviewer’s interpretation of our empirical evaluation as a “self-comparison”.

1. One of our baselines, Parallel SGD/AdamW, is the common and standard algorithm for distributed training (e.g., [6, 7, 10, 11]). Another baseline, setting $H$ as a constant throughout training, is the conventional approach for local gradient methods (e.g., [12-14]). Therefore, we have compared QSR with standard and widely-used algorithms instead of making a self-comparison.

2. As clarified in our response to M2, we have explored the most popular and standard learning rate schedulers and carefully tuned the learning rate for each baseline algorithm, ensuring that our comparison is fair. 

3. As for (2), we have already compared QSR with the scaling $H\sim \eta^{-1}$ in Figure 1 and Section 3. For completeness, we included $H\sim \eta^{-3}$ in Appendix G of the updated version, which shows that $H\sim \eta^{-3}$ does not provide consistent improvements over QSR.

Reviewer comment (Experiments):
Have you considered comparing the performance of $H\sim \eta^{-3}$, $H\sim \eta^{-4}$? While prior work [Gu et al.](https://arxiv.org/pdf/2303.01215.pdf) already shows that the generalization benefits arise only if H = Ω( 1/η ) when η → 0, this work specifically sets H to $η^{-2}$. This result seems to be quite heuristic and not a thorough investigation of the problem domain. What understanding about the limitations or insights of the method has been derived?

Author response:
**Q1(1)**: Have you considered comparing the performance of QSR with $H\sim \eta^{-3}$ or $H\sim \eta^{-4}$? While prior work [1] already shows that the generalization benefits arise only if $H = \Omega( 1/\eta )$ when $\eta \to 0$, this work specifically sets $H$ to $\eta^{-2}$. This result seems to be quite heuristic and not a thorough investigation of the problem domain. 

**A**:

1. As clarified in our response to W1(2), QSR is novel compared to [1]. Specifically, [1] studies the setting of using **constant $\eta$ and $H$** throughout training. By contrast, in QSR, $H$ dynamically grows in proportion to $\eta^{-2}$ during training as $\eta$ decays over time. The setting of a decaying learning rate schedule studied in this paper is a common choice in practice (e.g., [2-4]) and not studied in [1]. 
2. QSR goes far beyond a heuristic, and its strong implicit bias for sharpness reduction is well supported by our theory. To explore other choices of power, our paper has already included a thorough comparison of QSR with $H\sim \eta^{-1}$ both theoretically and empirically. Following the reviewer’s suggestion, we included $H\sim \eta^{-3}$ in Appendix G of the updated version, which shows that $H\sim \eta^{-3}$ does not provide consistent improvements over QSR. 

**Q1(2)**: What understanding of the limitations or insights of the method has been derived?

**A**: QSR offers an important insight that improved generalization performance and communication efficiency can be simultaneously achieved by adopting an appropriate $H$ schedule. The limitation of QSR, as discussed in Section 5, is that it may not consistently deliver noticeable generalization improvements for smaller models with shorter horizons. However, training in this regime is not costly, making it less of a critical concern.

Reviewer comment (Theory):
The proof provided for Lemma E.2 appears to be too succinct. Furthermore, in Lemma E.2, it is stated that 'we can apply the Itô-Taylor expansion from Lemma B.7 of Malladi et al. (2022) to derive (16), (17), and (18).' Is there a typo in this reference to equations? I suspect you meant to refer to equations (11), (12), and (13). Also, could you please elucidate how you derived equations 11, 12, and 13?

Author response:
**Q2**: The proof provided for Lemma E.2 appears to be too succinct. Furthermore, in Lemma E.2, it is stated that 'we can apply the Itô-Taylor expansion from Lemma B.7 of Malladi et al. (2022) to derive (16), (17), and (18).' Is there a typo in this reference to equations? I suspect you meant to refer to equations (11), (12), and (13). Also, could you please elucidate how you derived equations (11), (12), and (13)?

**A**: We thank the reviewer for pointing out the typo; we have corrected it in the updated version. As stated in the proof of Lemma E.2, to compute the moments of ${\phi}^{(s)}_{k, H}-{\phi}^{(s)}$ (eq (11)-(13) in the original submission and (13)-(16) in the updated version), it suffices to compute the moments of ${\zeta}(\alpha^2)-{\phi}^{(s)}$ since their moments differ in only $o(\mathrm{poly}(\alpha))$. The omitted step is a simple application of the standard Itô-Taylor expansion to the moments of ${\zeta}(\alpha^2)-{\phi}^{(s)}$. To provide the reviewer with a clearer understanding of our proof, we have enhanced the proof of Lemma E.2 with additional details and included the itô-Taylor expansion Lemma in the updated version.

Reviewer comment (Evaluation):
How do you calculate the communication time in Table 4, why don't you opt for using communication volume instead?

Author response:
**Q4**: How do you calculate the communication time in Table 4? Why not opt for using communication volume instead?

**A**: We derive the communication time from the difference in total training time across runs with various communication frequencies since the asynchronous nature of CUDA makes it hard to directly measure the communication time. Please refer to Appendix F in the updated version for details. Most of our figures and tables have already reported the communication volume, e.g., Figures 1 to 3, and Tables 1 to 3. The wall-clock time analysis aims to showcase that, compared with Parallel SGD/AdamW, QSR can indeed accelerate training by significantly reducing communication overhead.

Reviewer comment (Evaluation):
**No confidence intervals in experiments**: given the extensive hyper-parameter search Appendix C claims, I would have expected that the experimental results in the main paper would come with confidence intervals, and were presented as average over several runs to confirm the reproducibility of the results.

Author response:
**A**: We have repeated our main experiments and updated Table 1. We are unable to repeat all experiments due to high training costs (~20 hours for a single run on 16 NVIDIA GeForce RTX 3090 GPUs), but the improvement is consistent across different settings (Tables 2 and 3, Figures 2 and 3).

Reviewer comment (Experiments):
**No visualization of the effective $H$ scheduler**: Tab. 4 hints that using constant $H$ and QSR actually leads to very similar training times, suggesting that using the QSR scheduler for the values of $\alpha$ retained does not lead to major deviations compared to using constant $H$. In effect, what does the evolution of $H$ look like throughout training, and how different is it from using a constant $H$?

Author response:
**A**: Following the reviewer's suggestion, we have added a visualization (Figure 5) in the appendix, showing that it differs greatly from using a constant $H$. For example, in Figure 5, when we use $\alpha=0.0175$ for a 300-epoch cosine learning rate schedule with peak learning rate 0.008, $H$ starts to grow to 30 at epoch 183, and reaches 8142 at epoch 273 when$\eta$ decays to $1.9\times 10^{-4}$. Indeed, the difference in the $H$ schedule is significant enough to improve the test accuracy from 79.32% to 80.98%.

Reviewer comment (Theory):
**Unclear why $\gamma =2$ is optimal:** Section 2 claims that *"$H^{(s)}$ could have been set to $H^{(s)}:= \max \left \\{ H_\text{base},  \left \lfloor \left ( \frac{\alpha}{\eta_t}  \right )^{\gamma} \right \rfloor \right \\}$ for any $\gamma$"*. However, it still remains unclear to me why setting $\gamma = 2$ is optimal.

Author response:
**A**: As shown in eq (3) of our paper, the drift term gets stronger as $H$ increases, driving the iterate to drift faster towards flatter minima. However, $H$ cannot grow arbitrarily large; otherwise, the slow SDE loses track of Local SGD, as reflected in the error bound in Theorem 3.1 (3) and Theorem 3.3 in [1]. Therefore, to optimize the generalization performance, we should set $H$ to the largest value without making the slow SDE approximation ineffective. $H\sim \eta^{-2}$ is the largest possible $H$ that one can find a valid slow SDE approximation near the minimizer manifold. Furthermore, we empirically demonstrate the optimality of $\gamma=2$ via comparisons with $\gamma=1$ (Section 3) and $\gamma=3$ (Appendix G in the updated version).

Reviewer comment (Experiments):
**Comparison with related work lacking:** Appendix A cites the *"Stochastic Weight Averaging in Parallel (SWAP) algorithm"* [[2]](https://openreview.net/pdf?id=rygFWAEFwS) as another method that varies the value of $H$ during training, leading to a decrease in training time and better generalization. Yet, no comparison with the method is performed in the paper.

Author response:
**A**: We have added the comparison with SWAP in Appendix H of the updated version, where we showcase QSR's superior generalization performance over SWAP. Specifically, the original SWAP proposed in [7] is less communication efficient than QSR since it uses SGD for the majority of the training process and only switches to local updates at some $t_0$ near the end. To compare SWAP with QSR at a similar level of communication volume, we experiment with a slightly modified SWAP, which starts with Local SGD/AdamW using a constant communication period $H_{\mathrm{base}}$ and, after some time $t_0$, lets workers perform local updates with a final model averaging. As shown in Figure 9, QSR outperforms SWAP, though we have tuned $t_0$ carefully for the latter.

Reviewer comment (Experiments):
**Additional hyper-parameter to tune**: *"choosing the right value of $H$"* has been traded for *"choosing the right value of $\alpha$"*, thus QSR does not "automatically" set $H$ and no hyper-parameter tuning has been saved.

Author response:
**A**: We did not claim that QSR can ease the burden of hyperparameter tuning. Rather, we would like to highlight that, compared with $H$-schedules in previous works, e.g., post-local SGD, QSR improves test accuracy without increasing the number of hyperparameters to tune.

Reviewer comment (Experiments):
A network speed of 25Gbps is used for the experiments. Would your algorithm still help in HPC settings where fast communications can happen over 100 Gbps Infiniband?

Author response:
**A**: The communication time is expected to be 1/4 of that reported in Table 4 as the bandwidth is increased to 4 times.  However, we do not know the exact speedup for our experiments over 100 Gbps Infiniband since we do not have access to such a network.

Reviewer comment (Experiments):
Are both the model size **and** number of epochs important factors for QSR to work? For example, using a smaller model (such as ResNet18) but with a decent amount of epochs (e.g., 300 epochs on CIFAR10), would QSR have an impact?

Author response:
**A**: 

1. Large models are typically coupled with long horizons (e.g., several hundred epochs on ImageNet [3, 4, 8, 9]) to fully exploit the model’s capability. Training large models with long horizons is the setting where distributed training is necessary, and the computation cost is notoriously high, with the communication cost further exacerbating the situation. In our experiments, we follow standard training recipes of relatively large models, ResNet-152 [3] and ViT-B [4], and showcase that QSR can reduce communication costs and improve test accuracy.

2. By contrast, the small-scale experiments, e.g., ResNet-18 on CIFAR-10 mentioned by the reviewer, are cheap and fast, and do not require distributed training, let alone the communication cost.

3. The theoretical reasons are twofold. First, QSR is inspired by the Slow SDE, which, as interpreted in Section 3, captures the long-term dynamics of SGD/Local SGD. Ablation studies in Section 2 of [1] also show that the training budget should include enough iterations for the generalization benefits of Local SGD to arise, and that the generalization improvement over SGD becomes more significant as we train longer. Second, the model should be overparamterized, where the local minimizers of the training loss can form a manifold [2, 5, 6].

Reviewer comment (Experiments):
The experiments focuses on vision tasks and models. Although ViT model uses the transformer-like architecture, I'm still very interested in how QSR performs in NLP tasks with transformer architectures.

Author response:
**A**: Thank you for your suggestion on exploring QSR on NLP tasks. As discussed in Section 5, local gradient methods may have the potential to improve the transferability from pretraining to downstream tasks in NLP. However, given that the large-scale experiments on vision tasks are already costly and time-consuming (~20 hours for a single run), we leave the exploration of QSR on NLP tasks for future work.

Reviewer comment (Theory):
Many theoretical analysis suggests a learning rate scheduler $\eta \propto \frac{1}{\sqrt{T}}$, where $T$ is the number of steps. What if we detach the choice of $H$ from the learning rate, and simply use $H \propto T$ (with some rounding of course)?

Author response:
**A**: You are correct in pointing out that $H$ is proportional to $T$ if QSR is applied to a square root learning rate decay scheduler. However, this relationship does not universally apply to other widely used schedulers, e.g., cosine, linear and step decay, where $\eta\propto \sqrt{\frac{1}{T}}$ does not hold. Hence, we choose to express $H$ as the function of $\eta$, which offers a general form remaining consistent across different learning rate schedulers.

Reviewer comment (Theory):
Is it possible to extend the theoretical analysis to local Adam?

Author response:
**A**: It is possible but very challenging. Though QSR is motivated by the slow SDE for SGD/Local SGD, it can also improve the test accuracy when applied to Local AdamW, indicating that there may be some similar implicit bias near the minimizer manifold for Local AdamW/Adam. However, notice that the analyses for Local SGD in our paper and [1] has already been quite complicated. Extending this to Local Adam/AdamW requires the characterization of the dynamics of the precondition matrix and momentum, further complicating the proof. Given these difficulties, we leave the analysis of Local Adam/AdamW to future work.