Surge Phenomenon in Optimal Learning Rate and Batch Size Scaling

We are truly grateful for the time you have taken to review our paper and your insightful reviews. Below, we address your comments in detail.

W1.

A1. Due to paper page limitation, we were unable to delve deeply into references [1-2]. We acknowledge this limitation and will provide a more comprehensive analysis in a future edition.

Regarding the analysis of training speed and data efficiency in lines 129-131, we provide a detailed exposition in A6. In summary, Eq. 20 plays a pivotal role in both [1] and [2], particularly in its modification of the scaling law in [2]. If we demonstrate that Eq. 20 holds in the context of Adam, the modified scaling law conclusion in [2] naturally follows.

W2.

A2. Following your suggestion, we conducted finer search and presented the results in Fig. 2 of the attached file. The outcomes from these extended experiments provided clearer result which aligns with our theoretical analyses.

W3.

A3. In Appendix A of our paper, we explain why the theoretical results in sign SGD can be generalized to Adam. In both our original experiments in our paper (section 3.1 and 3.3) and our supplementary Experiment-1, we set both hyperparameters beta1 and beta2 to non-zero values, and still observed the proposed surge phenomenon in the experimental results. More details and results be found in the global rebuttal and the appendix.

Q1.

A4. Our motivation for investigating the scaling law of learning rates with respect to batch sizes is to gain deeper insights into the training dynamics of deep learning models. This understanding can aid in fine-tuning hyperparameters, enhancing convergence speed, and avoiding exhaustive grid searches. Leveraging prior knowledge that the optimal learning rate decreases after reaching its peak, researchers and engineers can effectively adjust the learning rate to achieve more efficient training.

Q2.

A5. Given that the peak shifts during training, a straightforward approach is to periodically adjust hyperparameters to redetermine the scaling law, including hyperparameters $\epsilon_{max}$ and $\mathcal B_{noise}$. Our observation may serve as inspiration for developing adaptive learning rate or batch size methods in the future, which can effectively harness this knowledge.

Q3.

A6. Section 2.3 of [1] demonstrates that integrating its Eq. 2.7 during the optimization process results in Eq. 2.11 (Eq. 20 in our paper). This implies that if the delta loss after each optimization step follows $\frac{\Delta L_{max}}{1 + \frac{\mathcal{B}}{B}}$, similar conclusions to Eq. 2.11 can be derived. Consequently, proving Eq. 18 ensures the validity of Eq. 20. This suggests that, Adam and SGD only affect the value of $\Delta L_{max}$, and do not involve other interference within the relation between the delta loss and batch size.

Q4.

A7. Considering the substantial training costs associated with grid search in our experiments, we utilize random 10k samples to mitigate these expenses. This approach accelerates model convergence while yielding consistent conclusions. Notably, when attempting to use the original dataset, we encountered the need for larger models to fit the data, resulting in significantly longer training times—far beyond what we can afford.

Beyond the considerations of training costs, selecting a distinct value for the hyperparameter $\beta_{i}$ does not compromise the validity of our conclusions. We provide a comprehensive discussion on the relationship between signSGD and Adam in Appendix-A. Furthermore, our paper and supplementary file include experiments that explore outcomes related to the Adam scenario. More details and results can be found in our paper (section 3.1 and 3.3) and our supplementary Experiment-1.

Q5.

A8. Thank you for your observation. In Fig. 2, we deliberately chose scenarios with higher loss to emphasize the decrease and subsequent increase learning rates parts. In these cases, the batch size corresponding to the peak of the curve (the red dashed lines) is relatively small in this case, leading to the grid search not fully capturing the rising range of the curve. However, our other experiments, documented in both the main paper and the supplementary file, reveal curves that more are closely align with the trend depicted in Fig. 1. We include these scenarios to gain a clearer understanding of the underlying dynamics.

L1.

A9. Thank you for your insightful suggestions. Our experiments, including the supplementary ones, demonstrate that the conclusions drawn from the second-order expansion approximation of the loss function effectively predict the surge phenomenon observed in most mainstream scenarios. In our final paper revision, we will thoroughly review and discuss the referenced articles. Additionally, we recognize the potential of exploring higher-order approximations with respect to scaling laws as promising future work.

L2.

A10. As observed in Fig. 2(a) of our paper and Fig. 4 in our supplementary file, the distribution of gradients along the dataset direction approximately adheres to Gaussian distributions. Furthermore, as demonstrated in Equation 11, this equation indeed possesses a functional extremum as long as the distribution satisfies ($f(B) \propto B$). This implies that, under broader conditions, both the batch size and learning rate still exhibit a surge during the scaling process. We appreciate the suggestions provided in L1-2 and plan to explore this issue more thoroughly in future work, incorporating the ideas discussed therein.

Thanks again for appreciating our work and providing constructive suggestions. We hope you will consider raising your score. Please let us know if you have further questions.

---

[1] An Empirical Model of Large-Batch Training, 2018, McCandlish et al. (reference [25] in our paper)

[2] Scaling Laws for Neural Language Models, 2020, Kaplan et al. (reference [26] in our paper)

We are grateful for the time and effort you have taken in reviewing our paper and for your thoughtful feedback. Here, we respond to your comments.

W1. It seems that Lemma 1 can only apply to quadratic problems. In the appendix, the relation is approximately equal. But in Lemma 1, it becomes "equal" without any further assumptions. Q1. Why does the theorem apply to general network optimization?

A1. We would like to clarify that Lemma 1 can be applied to ALL gradient descent-based machine learning processes. In this context, the quadratic term arises from the Taylor expansion of the loss with updated parameters, NOT from quadratic problems. General network optimizations all have the form of gradient descent-based learning. Here we use the "equal" sign for simplicity. The similar analysis and simplification were also proposed in OpenAI's paper[1] Eq. 2.4-2.7. We will improve the clarification and the symbol, and add the reference in Lemma 1 in the final version of our paper.

W2. It is unclear how to select the optimal batch size or learning rate based on the theorem because either S, E or μ, σ is hard to estimate for a large network. Q2. How to make the result of the theorem in use?

A2. Our work extends the conclusions on SGD optimizers in OpenAI's papers[1-2], focusing on Adam-style optimizers. Though in experiments we used grid-search to validate our conclusions, in practice the scaling law only contains two unknown hyper-parameters $\epsilon_{max}$ and $\mathcal B_{noise}$. Following papers [1-2], $\mathcal B_{noise}=B_{crit}$ can be efficiently approximated using scaling law $B_{crit} = \frac{B^*}{L^{\frac{1}{\alpha}}}$. We only need one simple search for a pair of (batch size, optimal learning rate) to determine the last hyper-parameter $\epsilon_{max}$. Therefore, the costly grid-search can be avoided. We will improve the presentation of our paper and add this clarification.

Thanks again for appreciating our work and for your constructive suggestions. We hope you will consider raising your score. Please let us know if you have further questions.

---

[1] An Empirical Model of Large-Batch Training, 2018, McCandlish et al. (reference [25] in our paper)

[2] Scaling Laws for Neural Language Models, 2020, Kaplan et al. (reference [26] in our paper)

Dear Reviewer JT4e,

We would like to express our sincere gratitude for your valuable feedback and for increasing the score of our paper. Your insightful comments and suggestions have greatly improved the quality of our work. We will ensure to incorporate these suggestions into the final manuscript, and we truly appreciate the time and effort you have put into reviewing our submission.

Thank you for taking the time to review our paper and for your valuable insights. We address your comments as follows.

W1. The presentation is not great. A lot is assumed from [1], but it would make reading easier to make things more self-contained.

A1. Due to page constraints, we currently include only essential assumptions and conclusions from related work. In the final version of our paper, we aim to make the contents more self-contained and enhance our presentation.

W1. It is unclear what is going on in Figure 1 between the solid and dashed Adam curves. Q1. In Fig. 1 how is the solid Adam curve generated? Shouldn't the curve eventually asymptote? It doesn't seem like that from the plot.

A2. The solid line of Fig. 1 in our paper is a schematic diagram that we created based on Eq. 13 and 17, while the dashed line of Fig. 1 is an illustration for Eq. 13. According to Eq. 11 and 13, there exists an extreme value of learning rate, resulting in a tendency to increase first and then decrease. As $B$ continues to increase, $\mathcal E_{i}(B)$ grows slowly according to Eq. 10 and 17, and finally asymptotically approaching a value, causing the solid line to gradually rise after a downtrend. For better visualization, we only showed part of the figure without plotting the final asymptote. Since the specific values of the downward and upward inflection points are related to the current loss and are hard be determined, the schematic diagram is only presented here for illustration purposes. Furthermore, in supplementary Experiment-3, we can provide additional evidence to support our conclusions by visualizing the experiment results in accordance with the trend shown in Fig. 1. More details can be found in the global rebuttal. We will provide more detailed explanations in the final version of our paper.

W1. Equations like Eq. 22 should be better explained and plots like in Fig 3 are hard to parse. The process for making the Fig. 3 plot is unclear. Q2. How is Fig. 3 generated? The linear fits look very strange.

A3. Eq. 22 is derived from Eq. 20-21, and the detailed derivation is in Appendix G. We will add more explanations in the final version of our paper. In Fig. 3, the left part is the grid search results of batch sizes and learning rates, while the right part is the fitted curve of Eq. 22 using the S and E data from corresponding (batch size, best learning rate) pairs in the left part. We will improve the presentation of the experiments.

W1. The takeaways and consequences for a practitioner are unclear. Q3. Do the results suggest using any modification of the square root scaling rule in practice? If so it would be helpful to have a comparison to understand potential benefit.

A4. In our experiments (see Fig. 2-4 in our paper and also the figures in the attached file), our fitted curve outperforms the square root scaling (corresponding to the alpha=0.5 curve). We have discussed the benefits of this approach in Section 3.3 and global rebuttal. Notably, while square root scaling may perform well in small batch size scenarios, it is essential to recognize that in large batch size scenarios, the optimal learning rate may no longer scale linearly and could decrease. For practitioners, we recommend following the papers [1-2] to approximate $\mathcal B_{noise}=B_{crit}$ using scaling law $B_{crit} = \frac{B^*}{L^{\frac{1}{\alpha}}}$, then use one simple search for a pair of (batch size, optimal learning rate) to determine the last hyper-parameter $\epsilon_{max}$. Subsequently, the conclusion of our paper can be applied to avoid costly grid-search. We will clarify the process in the final version of our paper.

Q4. Is there any characterization of the intermediate (i.e. large but not infinite batch) behavior of the scaling?

A5. In our supplementary Experiment-3, we validated our theory by referring to publicly available experiment results from another paper, which included data on scaling in the intermediate transition phase. We have also visualized these results to further support our findings. More details and results be found in the global rebuttal and the appendix file.

W1. Minor: the Latex parentheses look sloppy.

A6. We will improve the readability of our paper.

Thanks again for appreciating our work and for your constructive suggestions. We hope you will consider raising your score. Please let us know if you have further questions.

---

[1] An Empirical Model of Large-Batch Training, 2018, McCandlish et al. (reference [25] in our paper)

[2] Scaling Laws for Neural Language Models, 2020, Kaplan et al. (reference [26] in our paper)

Thank you for your thorough review and constructive suggestions. We address your comments in the following paragraphs.

W1. I think this paper could benefit from experiments with more popular architectures for the NLP tasks (e.g. maybe it would be useful to include some experiments on tasks with llama or mistral models).

A1. We provided additional experimental analyses on both sparse MoE and dense Transformer models for NLP tasks. The details can be found in the global rebuttal. For the sparse MoE structure, we experimented on the fine-grained MoE model with shared experts[1], which is a popular model and has a similar structure to the well-known Mistral-MoE[2]. As for the dense model, we validated our proposed conclusions by referring to another paper's experiments[3]. Please refer to the global rebuttal for the configurations, results, and analyses of these experiments.

W2. I also think it would be useful to have experiments with more datasets. Recent work shows that the data itself matters. E.g., for fine tuning LLMs, many factors wrt the data (e.g., data quality, variable sequence lengths, deduplicated data) can affect training. It would be interesting to see if the surge phenomena is agnostic to these factors or not.

A2. In the global rebuttal and our attached files, we incorporated additional datasets to enhance the robustness of our findings. Specifically, the Experiment-1 utilized the RedPajama-v2 dataset, and Experiment-3 used a combination of Deepseek's in-house data and OpenWebText2. Please refer to global rebuttal for datasets' details. We acknowledge that data-related factors may affect training. However, due to their complexity and abundance, we leave the impact of these factors on surge phenomena as future work, which could lead to a separate, valuable research endeavor.

Q1. For NLP related tasks, the number of tokens in a batch might vary. Would the surge phenomena still apply for learning rate vs. num tokens?

A3. In our paper and previous analysis[4], "batch size" is actually "token batch size", which considers the total number of tokens in a batch. In our experiments, we use the packing strategy to splice the data in each batch, so the number of tokens in all batches can be considered equal. We will address the misuse of this terminology in the final version.

Thanks again for appreciating our work and for your constructive suggestions. Please let us know if you have further questions.

---

[1] DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models, 2024, Dai et al.

[2] Mixtral of Experts, 2024, Jiang et al.

[3] DeepSeek LLM Scaling Open-Source Language Models with Longtermism, 2024, Bi et al.

[4] An Empirical Model of Large-Batch Training, 2018, McCandlish et al. (reference [25] in our paper)