Understanding outer learning rates in Local SGD

Thank you so much for your review and positive evaluation of our work.

1. (Non-convex analysis). Our theorem mainly relies on the two lemmas Lemma B.6 and Lemma B.7; Lemma B.6 holds regardless of convexity, but Lemma B.7. does not. We would need to relax Lemma B.7. to handle non-convex objectives, but this lemma in turn relies on the fact that gradient descent is contractive (Lemma B.1.), which is not true for arbitrary non-convex functions. We believe it is possible to use similar tools as in [10, Theorem 4.1], but this requires adding the bounded gradients assumption and using a different approach for the theorem derivation that starts with the smoothness inequality. Alternatively, we can use the approach of [12] of bounding this sequence by a sum of gradients at the cost of an extra $H$ factor in the Local SGD variance term.

2. (LR grids) The local learning rate is always fixed in our experiments is fixed to $\frac{\eta_{\mathrm{baseline}}}{\sqrt{M}}$ where $\eta_{\mathrm{baseline}}$ is the local learning rate that was tuned on the minibatch SGD baseline for each model scale and $M$ is the number of nodes (see Table 1 in the main paper). We only varied the outer learning rate otherwise, since what is important for our analysis is the ratio of outer to local learning rate. We have included a table with all the experiments we’ve conducted varying the outer learning rate in our response to Reviewer advP.

3. (Implications of Thm. 3.3) Theorem 3.3 shows there are two sources of noise for Local SGD: the second term in equation (3) scales with $\eta$ and $\max(0, \gamma-1)$ but not $H$, and the third term scales with $\eta^2$ and also $H$. In addition, we also have to satisfy the constraint $\eta L (1 + (\gamma-1)_+ H) \leq 1$. For tuning the outer learning rate, if either $\eta$ or $H$ are larger, we are forced to use smaller $\gamma$ in the same way. However, the convergence of Local SGD will suffer more under larger $\eta$ than larger $H$, all else being equal, because larger $\eta$ will increase both noise terms while larger $H$ only increases the latter.

4. (Minimizing FLOPs) To minimize iteration complexity given a fixed number of clients $M$ and local steps $H$, we ought to choose the outer learning rate to make up for a small inner learning rate and to vary with the number of local steps. Depending on the magnitude of the gradient variance, choosing an outer learning rate larger than 1 can also be useful. We found that outer learning rates can be tuned at a smaller scale and transfer as-is to larger models. For example, with the 150M model at H=50, M=4, it achieves a perplexity of 17.75, while Nesterov with lr=0.7 reaches 17.25 and SF-SGD with lr=2.0 achieves 16.88. These represent improvements of 2.8% and 4.9% respectively compared to the baseline with no outer learning rate. The performance gains transfer and become even more pronounced at larger model scales: for the 1B model, vanilla SGD achieves 13.67 perplexity while Nesterov (lr=0.7) and SF-SGD (lr=2.0) reach 12.51 and 12.40 respectively, representing improvements of 8.5% and 9.3%. Of course, in practice, we control the number of clients $M$ and also have to consider FLOP efficiency from a systems perspective, e.g. normally increasing $M$ comes at the cost of decreasing the local batch size or alternatively increasing $\sigma^2$. This complicates the trade-off; Figure 4 shows the result of varying the number of clients/replicas in practice on FLOP efficiency. A study of FLOP efficiency would have to take into account the specific hardware we have available, and what the minimum batch size needed to saturate the GPU computation is.

5. “For Figure 9 in the Appendix, why are we seeing high gradient similarity near the end of the training for the left plots and why are we seeing low gradient similarity for the right plots?” For all of them except SF-SGD in Figure 9 (c), we do observe low similarity towards the end of training (lower than 0.1). For SF-SGD in Figure 9 (c), we do observe a spike in similarity towards the end but are not sure why this is the case.

6. “Furthermore, line 283 states "... noise-dominated regime, which we may expect to observe towards the end of the training process." Why is that so?” This is an empirical observation made in prior work [11], and the intuition is that towards the end of training the gradient norms are smaller so the relative magnitude of the noise is higher compared to the gradient norms.

7. “f. For Figure 1b, why does the loss start high for low variance setting of $\sigma = 0.001$?” All of the losses are reported after one communication round– the figure should start at r=1 not at r=0. They would all start at the same value if we started at r=0. Since we cannot upload PDFs or figures, here are the trajectories in tabular form:

Figure 1 (b) tabular

5. “g. For Table 1, SF-SGD case, why are outer learning rates with magnitudes >2 detrimental?” They weren’t really detrimental, they just did not help very much. Here are the results of our grid search on a 150M parameter model with $H=50$ and $M=4$. Very small outer learning rates were a lot more detrimental, as they did not allow for enough progress.

Additional Learning Rate Sweeps (150M, H=50, M=4)

[10] Glasgow, M. R., Yuan, H., & Ma, T. (2022, May). Sharp bounds for federated averaging (local sgd) and continuous perspective. In International Conference on Artificial Intelligence and Statistics (pp. 9050-9090). PMLR. [11] Faghri, F., Duvenaud, D., Fleet, D. J., & Ba, J. (2020). A study of gradient variance in deep learning. arXiv preprint arXiv:2007.04532.

Thank you so much for your review and positive evaluation of our work.

1. (Reproducibility) We did not report all of the hyperparameter sweep results in line with reporting in prior work which reported the best hyperparameters only. We are happy to report all the experimental results we have done. Here are the results of the hyperparameter sweeps we will include in the appendix:

Additional Learning Rate Sweeps (150M, H=50, M=4)

SF-SGD Beta (b) Parameter Sweep (150M, H=50, M=4, lr=2.0)

2. (Impact) Yes, it is true that our theory as it is provides limited practical prescriptions. We are primarily concerned with understanding the role of the outer learning rate and how it should be tuned, but have not directly provided a theory-based schedule for it; this task we leave to future work. One crucial takeaway from our theory, which is not common in practice, is that it sometimes pays off to try outer learning rates greater than 1.0. To the best of our knowledge, the current literature (e.g. [7, 8]) did not sweep for such learning rates, going up to 1.0 and no higher. Our work shows that for Schedule-Free SGD as the outer optimizer, a learning rate of 2.0 performed best. We will highlight this recommendation further.

3. (Captions) We will expand the captions to be more self-contained in the revised manuscript. Thank you for pointing this out.

4. "Is the use of a schedule-free outer optimizer instead of SGD with a schedule is due to the theory not allowing for scheduled outer learning rates? If so, mentioning this could be helpful for the reader." The main reason we used it was to remove the variable of (outer) learning rate scheduling. The theory can be extended to scheduled outer learning rates, but in its current form it does not allow for them. We will mention this.

5. "How is perplexity calculated in Figure 2? (e.g., what data, what batch size, etc)" Perplexity is calculated on the C4 validation set, with sequence length 1024, and batch size 512 for all model scales. The tokenizer is SentencePiece.

6. "Why does Theorem 3.5 not use an EMA for the momentum?" Our way of writing the momentum update is equivalent to using EMA for momentum. [9, Lemma 7.2] shows it is equivalent to the update $m_{t} = \mu m_{t-1} + g_t$ and $x_{t+1} = x_t - \eta m_t$. By unrolling, this corresponds to $x_{t+1} = x_t - \eta \sum_{k=0}^t \mu^k g_{t-k}$. For classical EMA written as $m_t = (1-\alpha) m_{t-1} + \alpha g_t$ and $x_{t+1} = x_t - \gamma m_t$, unrolling we have $x_{t+1} = x_t - \gamma \alpha \sum_{k=0}^t (1-\alpha)^k g_{t-k}$. Thus putting $\mu = 1-\alpha$ and $\gamma (1-\mu) = \eta$ we can see that they are equivalent up to a rescaling of the learning rate.

7. "For LLM experiments, were the hyperparameters swept at each scale or were they swept only at a smaller scale?" We swept them at the 150M scale, with a small number of additional trials at higher scale.

Thank you so much for your detailed comments.

[7] Douillard, Arthur, et al. "Diloco: Distributed low-communication training of language models." arXiv preprint arXiv:2311.08105 (2023). [8] Liu, B., Chhaparia, R., Douillard, A., Kale, S., Rusu, A. A., Shen, J., ... & Ranzato, M. A. (2024). Asynchronous local-sgd training for language modeling. arXiv preprint arXiv:2401.09135. [9] Garrigos, G., & Gower, R. M. (2023). Handbook of convergence theorems for (stochastic) gradient methods. arXiv preprint arXiv:2301.11235.

Thank you so much for your review.

1. "Could the authors discuss the challenges of extending their analysis to the heterogeneous setting?" Our primary setting is large language model training, as in DiLoCo and similar works, and there the training data is i.i.d. Additionally, we would like to point out that in the heterogeneous setting, the guarantees for Local SGD are very pessimistic; that is, small degrees of heterogeneity result in much worse theoretical guarantees for Local SGD compared to Minibatch SGD, as discussed in [3, p. 5-6], and it is not clear under which assumptions the algorithm should be analyzed (see [12]). Because we wanted to focus on when Local SGD is superior to the alternatives, this is one more reason we focused on the i.i.d. setting.

That said, we believe extending our analysis to the heterogeneous setting is possible. Lemma B.6. still holds, but Lemma B.7. should be adjusted to handle function heterogeneity. If we make the same assumption as [1, Assumption 1.2 (3)], then this is possible at the cost of an extra $H$ factor. Modifying the rest of the proof of Theorem 3.3 with the revised Lemma B.7. is straightforward. The biggest drawback is that the conclusion will be too pessimistic compared to Minibatch SGD. We believe it might be possible to use assumptions such as second-order similarity and second-order smoothness, as in [4], to get around this, but this requires significant modifications not just to Lemma B.7 but to the main proof strategy of Theorem 3.3, and we leave it to future work.

2. "Are there any new frameworks or techniques in the theoretical analysis compared to prior works like [1][2]? If so, the authors should highlight these differences more explicitly." Yes, there are. First, our algorithmic framework is significantly different; Neither [1] nor [2] consider the use of outer learning rates or outer momentum, with both using simple averaging (equivalent to fixing the outer learning rate at $1$) and [1] considering local momentum. Our framework enables us to study the conditions under which using outer learning rates and momentum is useful. This difference in framework is the primary reason why we compare against [5, 6] in the discussion after Theorem 3.3 and not [1, 2]. Second, the analysis of Local SGD always involves bounding the so-called "client drift", or how far the local iterates stray from each other between synchronizations. In [1] this is controlled in Lemma 4 and in [2] this is controlled in Lemma 3.3. Both of these control the drift by bounding the maximum squared norm. We control the drift using a combination of the local regret against the starting point (Lemma B.6) and the maximum squared norm (Lemma B.7). This more fine-grained control allows us to show the benefit of using a large outer learning rate. To the best of our knowledge, a regret guarantee like Lemma B.6. has not been used in the literature before. In order to use this lemma, we also have to conduct the analysis per communication round rather than per training step as in [1, 2]. We will revise the next version of the manuscript to highlight these novel contributions.

3. "For optimization methods, it’s not recommended to introduce new hyperparameters if not necessary, due to the cost of hyperparameter tuning. Gen-loc-SGD has more learning rates that need to be tuned compared with vanilla local SGD. Are there some dramatic improvements due to the introduction of the outer learning rates?" Yes, the performance degradation when not using an outer learning rate is significant. In our experiments, vanilla Local SGD consistently underperforms compared to using outer optimizers. For example, with the 150M model at H=50, M=4, it achieves a perplexity of 17.75, while Nesterov with lr=0.7 reaches 17.25 and SF-SGD with lr=2.0 achieves 16.88. These represent improvements of 2.8% and 4.9% respectively. The performance gains become even more pronounced at larger model scales: for the 1B model, vanilla SGD achieves 13.67 perplexity while Nesterov (lr=0.7) and SF-SGD (lr=2.0) reach 12.51 and 12.40 respectively, representing improvements of 8.5% and 9.3%. Because we can tune the outer learning rate at a smaller scale (this is what we did here) and therefore at a much smaller cost, we believe the additional hyperparameter is justified. There are very few reliable ways of squeezing out an additional 8%-10% performance at the 1B scale at the cost of one additional hyperparameter that can be tuned at a smaller scale.

4. Thank you for your comments on improving clarity. We will define the learning rates $\gamma$ and $\eta$ in Algorithm 1 and use bold faces for the vectors.

We thank you for your review and hope our discussion above sufficiently addresses your concerns.

[1] Hao Yu, et al. On the Linear Speedup Analysis of Communication Efficient Momentum SGD for Distributed Non-Convex Optimization. ICML 2019. [2] Stich, Sebastian U. Local SGD Converges Fast and Communicates Little. ICLR 2019. [3] Woodworth, B. E., Patel, K. K., & Srebro, N. (2020). Minibatch vs local sgd for heterogeneous distributed learning. Advances in Neural Information Processing Systems, 33, 6281-6292. [4] Zindari, A., Luo, R., & Stich, S. U. (2023). On the convergence of local sgd under third-order smoothness and hessian similarity. In OPT 2023: Optimization for Machine Learning. [5] Karimireddy, S. P., Kale, S., Mohri, M., Reddi, S., Stich, S., & Suresh, A. T. (2020, November). Scaffold: Stochastic controlled averaging for federated learning. In International conference on machine learning (pp. 5132-5143). PMLR. [6] Jhunjhunwala, D., Wang, S., & Joshi, G. (2023). Fedexp: Speeding up federated averaging via extrapolation. arXiv preprint arXiv:2301.09604. [12] Patel, Kumar Kshitij, et al. "The limits and potentials of local sgd for distributed heterogeneous learning with intermittent communication." The Thirty Seventh Annual Conference on Learning Theory. PMLR, 2024.

Thank you so much for your review and positive evaluation of our work.

1. We agree that our use of "adaptive" in Section 3.3, as we do not mean adaptivity in the same sense as Adam or AdaGrad (i.e. algorithmic adaptivity), but rather in the sense of obtaining data-dependent guarantees. We will modify the term "adaptive" to instead be "data-dependent" and thank the reviewer for pointing this out.
2. We will reorganize the presentation of Theorems 3.5 and 3.6 to highlight the key takeaways and defer the more technical discussions. The core observation from Theorem 3.5, is that momentum allows us to use an effectively larger stepsize but does not alter the outer stepsize tradeoff shown by Theorem 3.3. Similarly, Theorem 3.6 shows that when noise magnitude dominates the optimization error (e.g. in later training stages), using $\gamma < 1$ helps maintain convergence. Conversely, when the optimization term is larger than noise, $\gamma > 1$ acts like momentum to accelerate convergence. When both terms are comparable, $\gamma = 1$ (simple averaging) is optimal.
3. We thank the reviewer for pointing this out. Figure 2 shows the total number of training steps (i.e. $R \times H$ where R is the number of communication rounds and $H$ is the number of local steps). We will clarify this and add a short table or glossary with all the abbreviations we use.