Surge Phenomenon in Optimal Learning Rate and Batch Size Scaling

Summary

The paper analyzes the intricate relationship between optimal learning rate and batch size scaling for adaptive optimizers, such as Sign SGD and Adam. Building on prior analysis for SGD by McCandlish (2018), this work reveals a non-monotonic relationship between optimal learning rates and batch size. The optimal learning rate initially increases, reaches a peak, and then decreases, eventually saturating to a non-zero value, referred to as the surge phenomenon. These theoretical predictions are validated through experiments in both image classification and NLP tasks.

Strengths

* The paper provides new insights into the relationship between optimal learning rate and batch size for Adam optimizers. The drop in the optimal learning rate after the peak is new, to the best of my knowledge, and is relevant given that Adam is the default optimizer choice. * Prior results suggesting a square root scaling for Adam are reproduced, further supporting the findings.

Weaknesses

* At times, the paper assumes that the reader is well-versed with the prior work of McCandlish (2018). For instance, lines 129-131. It would be helpful to reiterate prior results to motivate the analysis. * The empirical results are not very convincing. If we only consider the experiments (for instance Figure 4), without any reference to theoretical results, the surge phenomenon does not seem appreciable. A finer learning rate search has to be performed to demonstrate the surge phenomenon clearly. In Figure 4(b), the optimal learning rate is oscillating around two points. It is unclear if this is due to the surge phenomenon or just random fluctuations. I would request the authors to help me understand their empirical results better. * The theoretical results are derived for sign SGD, while it's known that Adam parameters beta1 and beta2 are crucial hyperparameters. It's unclear why the theoretical results can be generalized to Adam.

Questions

* What is the practical implication of the decrease in the optimal learning rate after the peak? * Given that the peak shifts through training, can the authors propose guidelines for scaling the learning rate and batch size? * What is the intuition behind the training speed and data efficiency relationship being the same as the SGD case? * Why the ResNet model (Figure 3) is trained with random 10k samples at every epoch? This should affect the overall results. Also, why this experiment is performed with sign SGD only? * If the results of Figure 2(b) and 2(c) are combined, then it should predict (1) a drop in the optimal learning rate, (2) saturation, (3) increase, and finally saturation. How does this result align with the main result (Figure 1)?

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

* The theory is built on the quadratic approximation of the loss function. In the last few years, it has been established that modern neural networks are typically trained at large learning rates (Edge of Stability, see Refs. [1-2]), which cannot be captured using quadratic approximations of the loss [3]. * Gaussian distribution for gradients is assumed for the theoretical analysis, whereas it is known that the gradient distribution is not Gaussian, and this is precisely why Adam performs better than SGD in language modeling. It is unclear whether the results hold for such settings. [1] Gradient descent on neural networks typically occurs at the edge of stability Cohen et al. (2021) arXiv:2103.00065 [2] Adaptive gradient methods at the edge of stability Cohen et al. (2022) arXiv:2207.14484 [3] Self-stabilization: The implicit bias of gradient descent at the edge of stability Damian et al. (2022) arXiv:2209.15594 [4] Linear attention is (maybe) all you need (to understand transformer optimization) Ahn et al. (2022) arXiv:2310.01082

Summary

The paper presents a heuristic analysis of the scaling of the optimal learning rate with batch size for Adam-style optimizers in the framework of [1]. The analysis is accompanied by experiments to support the predictions. Notably the authors demonstrate that the optimal learning rate can decrease with batch size in a certain range, a phenomenon that had not been identified previously and which is not present for SGD. The analysis also recovers the square-root scaling rule in the small batch size regime identified in other work. [1] McCandlish, S., Kaplan, J., Amodei, D., & Team, O. D. (2018). An empirical model of large-batch training. arXiv preprint arXiv:1812.06162.

Strengths

The paper tackles a practically important problem using a mix of heuristic theory and experiments. The prior literature on linear scaling rules with large batch training applies to SGD but not to Adam style optimizers which are the dominant optimizers for transformers. The surge phenomenon is interesting and novel.

Weaknesses

The presentation is not great. A lot is assumed from [1], but it would make reading easier to make things more self-contained. 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. It is unclear what is going on in Figure 1 between the solid and dashed Adam curves. The takeaways and consequences for a practitioner are unclear. Minor: the Latex parentheses look sloppy. [1] McCandlish, S., Kaplan, J., Amodei, D., & Team, O. D. (2018). An empirical model of large-batch training. arXiv preprint arXiv:1812.06162.

Questions

My questions are related to the perceived weaknesses 1. 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. 2. How is Fig. 3 generated? The linear fits look very strange. 3. 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. 4. Is there any characterization of the **intermediate** (i.e. large but not infinite batch) behavior of the scaling? It would be helpful to have an analog of Figure 1 for the empirical data.

Rating

6

Confidence

3

Soundness

3

Presentation

1

Contribution

2

Limitations

Yes.

Summary

The paper gives an optimal choice of learning rate and batch size for neural networks. Different from the previous results on SGD-style optimizers. The authors give such solutions for Adam-style-optimizers.

Strengths

1. Batch size and learning rate will affect the performance largely and cost a lot to select a good one. It is important to understand the optimal batch size and learning rate for Adam-style optimizers. 2. The experimental results match the theorem proposed by the authors.

Weaknesses

1. 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. 2. It is unclear how to select the optimal batch size or learning rate based on the theorem because either S, E or $\mu,\sigma$ is hard to estimate for a large network.

Questions

1. Why does the theorem apply to general network optimization? 2. How to make the result of the theorem in use?

Rating

6

Confidence

3

Soundness

3

Presentation

2

Contribution

2

Limitations

N/A

Summary

This work provides a scaling law between learning rate and batch size for Adam. Namely, this work finds that the optimal learning rate increases and then decreases as the batch size becomes larger; and, the peak of this curve corresponds to the trade-off between training speed and data efficiency.

Strengths

- The paper is well motivated. While prior works studying the relationship between batch size and learning rate have focused on SGD, this work focuses on Adam (which is more popular / widely used). - The paper is written clearly and is well organized. I appreciate the “summary” notes included by the authors - The paper includes empirical evidence for CV and NLP tasks to support theoretical claims.

Weaknesses

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

Questions

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?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

It would be interesting to see results on a larger variety of modern and widely used datasets and architectures.