Surge Phenomenon in Optimal Learning Rate and Batch Size Scaling
Authors: Shuaipeng Li
Penghao Zhao
Hailin Zhang
Samm Sun
Hao Wu
Dian Jiao
Weiyan Wang
Chengjun Liu
Zheng Fang
Jinbao Xue
Yangyu Tao
Bin CUI
Di Wang
Venue: NeurIPS 2024 poster (2024)
Keywords: Optimal learning rate
Scaling law
Primary area: learning_theory
TL;DR: Make theoretical research corrections to the relationship between optimal learning and batch size of models trained using Adam-style optimizers.

Abstract: In current deep learning tasks, Adam-style optimizers—such as Adam, Adagrad, RMSprop, Adafactor, and Lion—have been widely used as alternatives to SGD-style optimizers. These optimizers typically update model parameters using the sign of gradients, resulting in more stable convergence curves. 
The learning rate and the batch size are the most critical hyperparameters for optimizers, which require careful tuning to enable effective convergence. Previous research has shown that the optimal learning rate increases linearly (or follows similar rules) with batch size for SGD-style optimizers. However, this conclusion is not applicable to Adam-style optimizers. 
In this paper, we elucidate the connection between optimal learning rates and batch sizes for Adam-style optimizers through both theoretical analysis and extensive experiments. 
First, we raise the scaling law between batch sizes and optimal learning rates in the “sign of gradient” case, in which we prove that the optimal learning rate first rises and then falls as the batch size increases. Moreover, the peak value of the surge will gradually move toward the larger batch size as training progresses.
Second, we conduct experiments on various CV and NLP tasks and verify the correctness of the scaling law.

Surge Phenomenon in Optimal Learning Rate and Batch Size Scaling
Shuaipeng Li Penghao Zhao Hailin Zhang ..22 Xingwu S.n+.13, Hao Wu Dian Jiao', Weiyan Wang', Chengjun Liu Zheng Fang'. Jinbao Xue Yangyu Tao Bin Cui2.4t, Di Wang'
Tencent Hunyuan
School of Computer Science & Key Lab of High Confidence Software Technologies (MOE), Peking University
University of Macau
Institute of Computational Social Science, Peking University (Qingdao)
Abstract
In current deep learning tasks, Adam-style optimizers- Adam, Adagrad, RMSprop. Adafactor, and have been widely used as alternatives to SGD- style optimizers. These optimizers typically update model parameters using the sign of gradients, resulting in more stable convergence curves. The learning rate and the batch size are the most critical hyperparameters for optimizers, which require careful tuning to enable effective convergence. Previous research has shown that the optimal learning rate increases linearly (or follows similar rules) with batch size for SGD-style optimizers. However, this conclusion is not applicable to Adam- optimizers. In this paper, we elucidate the connection between optimal learning rates and batch sizes for Adam-style optimizers through both theoretical analysis and extensive experiments. First, we raise the scaling law between batch sizes and optimal learning rates in the "sign of gradient" case, in which we prove that the optimal learning rate first rises and then falls as the batch size increases. Moreover, the peak value of the surge will gradually move toward the larger batch size training progresses. Second, we conduct experiments on various CV and NL.P tasks and verify the correctness of the scaling law.
Introduction
Deep learning techniques, initiated by Stochastic Gradient Descent (SGD) learning on large datasets, have significantly revolutionized various real-world applications [1]. Over the past decade, numerous optimizers, such as momentum [2], Adagrad [3], ADADELTA [4], RMSprop [5], Adam [6], Adafac- [7], and Lion [8], have been introduced to stabilize the iterative learning process and expedite convergence. Among them, the Adam optimizer is the most widely used across various domains
including Computer Vision (CV) [9-11] 11], Natural Language Processing (NLP) [12-15] and many others [16, 17]. It retains the first and second moment information of parameters to facilitate adaptive learning step size. Unlike SGD-style optimizers that use the raw gradient to determine the learning direction and step size, Adam and its variants (Adagrad, RMSprop. Lion, etc.) employ the sign of gradient for this purpose, thereby ensuring greater robustness [18].
Beyond specific byper-parameter in optimizer configurations, the batch size and the learning rate are the most critical hyperparameters influencing convergence, As the scale of training datasets in various
'Equal contribution. 1Corresponding author.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Focdl 11.0
0.0 0.5
2.0
2.5
105
Figure I:
The relationship between the optimal learning rate and the batch size is different between Adam and SGD. The orange line represents the tendency of the optimal learning rate to converge to non-zero value when the batch size large enough.
workloads (e.g. CV [19, 20), NLP [21, 14), and others) continues to grow, there is an increasing demand for large batch size training across multiple data parallel workers. However, large batch training presents great challenges for robust training and meticulous tuning. The learning rate, which determines the actual step size each learning iteration, is highly dependent on the batch size used, Prior research has explored methods to determine an optimal learning rate according to the batch size in scenarios utilizing SGD optimizers, including square root scaling [22], linear scaling [19, 23], and others [24]. Among these, the empirical model focusing on large batch training [24] yields convincing results in both theoretical and empirical contexts, proposing the following rule depict the relationship between the optimal learning rate and the batch size:
Gmaz
Copt B) =
(1)
B Bagize
For Adam-style optimizers, though existing works also provide some approximation [24, 23], they fail to capture the true scaling law of optimal learning rates with batch sizes. For illustrative purposes, Figure presents curves that simulate the optimal learning rates for the Adam optimizer. We find that, in scenarios involving small batch sizes, the optimal learning rate initially increases and then decreases, resembling a surge in the sea, as depicted by the dashed orange line. For large batch sizes, identify value to which the optimal learning rate converges. The solid orange line represents schematic based on our findings for both small and large batch sizes, showing that the learning rate tends to rise initially, then decrease, and subsequently gradually ascend to asymptotically approach stable value. For clarity in visualization, we have omitted the final asymptotic portion of the curve.
In this paper, aim to elucidate and formalize the connection between optimal learning rates and batch sizes for Adam-style optimizers. By following the notation from the empirical model [24] and conducting more in-depth theoretical analysis, discover that the relationship between the optimal learning rate and the batch size in the above parameter update formular satisfies:
Copt (B)
mmar
Brasee
B
B BRoise
which differs from SGD, especially when the batch size is not too large. Here the meaning of Bnoise is consistent with papers of scaling laws [24, 25], representing the trade-off point between training speed and data efficiency. When the batch size equal Broise: the optimal learning rate reaches local maximum in accordance with Eq 2. Furthermore, we provide additional proof that when the batch size becomes significantly large, the optimal learning rate gradually converges to a non-zero value. We also prove that the previous conclusions about training speed and data efficiency are still valid for Adam-style optimizers, and the variable Broise gradually increases as the training progresses. It is worth noting that when B < Broise. for SGD, the scaling law of optimal learning rates with

---

batch sizes transitions into linear scaling, consistent with previous conclusions [19, 23]:
Eopt
(B) Pd
Emaz
B;
(3)
B Bnoise
while for Adam, the relationship transitions into square root scaling, aligning with previous approxi- mations [23, 24]:
2mmoz
Copc Copc
VB.
VBnosse
In addition to theoretical analysis, our extensive empirical study on various CV and NLP workloads further validate the conclusions. The true optimal learning rate, across different Adam configurations, exhibits a clear downward trend after reaching its peak value as the batch size increases. This behavior contradicts previous research, but demonstrates the correctness and generalizability of our theory. The experiments also reveal a gradual increase in the variable Bnoise. corresponding to the peak optimal learning rate, as training progresses.
Theorems
2.1 Batch Size and Optimal Learning Rate
In this section, we theoretically derive the optimal learning rate for given batch size. Initially, we introduce an approximation of Adam-style optimizers. In alignment with the insights elucidated in [26], a thorough examination of the Adam optimizer and its variants reveals that their primary distinction from SGD resides in the utilization of the gradient's sign for updates during each iteration, as opposed to the gradient itself:
0++1 0, sigin sign(Gest))) (Gest),
(5)
where 0, is the parameter at time I, Gest is the gradient estimated via mini-batch, and is the learning rate. As the batch size increases, the expected value of the update amount tends to saturate. For example, assuming that the mean value of the gradient is positive, when the accumulated gradient of the mini-batch is positive, increasing the batch size will have no contribution to the signed update amount. This is significantly different from the behavior of SGD where the larger the batch size, the more accurate the gradient estimate. In Appendix A, we provide a detailed discussion on this approximation for the Adam optimizer. Next, we derive the optimal learning rate that maximizes the loss improvement. And then we establish a lemma that addresses the optimal learning rate given an estimated mini-batch gradient:
Lemma 1. Suppose that we are updating the parameter O using the mini-bateh gradient V, with the true gradient being G and the true Hessian being H. Then the optimal learning rate that maximizes the decrease in loss is:
G E Copt. orgmar, E[AL] H cVo( V HE
(6)
and the corresponding loss improvement AL is:
G
Aopp 2
Copt.
(7)
The proof is in Appendix B. Although conclusion is based approximation, adopt the equal sign here to simplify the analysis, following the notation used in previous work [24].
Now let's consider the case where V = sign(Gest). and assuming that the estimated gradient Gest follows Gaussian distribution. The Gaussian distribution assumption is motivated by the following: if the mini batch size is sufficiently large, we can invoke the Central Limit Theorem (CLT) and approximate the distribution as Gaussian- a common assumption in previous research [27-30]. Furthermore, our experimental results confirm that the gradient distributions closely approximate Gaussian distributions, as illustrated in Figure 8 of Appendix H. We have the following theorem:
Theorem 2. Suppose the gradient of parameter for each sumple follows Gaussian distribution with mean μa and variance of the expected loss improvement is:
E.E, E,E,pug A A _1 H &,E,H,
(8)
3

---

and the corresponding optimal learning rate is
Eipe
Copt E,(1-17)EH E,E,&
(9)
where & is a function (derived from the Gauss error function) with respect the batch size B:
E(B)= VT
d
(10)
T 2B
+
We prove the above theorem in Appendix C.
An important observation from the proof is that, not only is the covariance matrix of sign(Gest) related to B, but its expected value also depends on B. This implies that in Eq 6 the numerator is the first-order form of the function about B, and the denominator is the second-order form of the function about B:
Bf(B) J(B)² +y
B
(B)
(11)
f(B)
7(3
Therefore, the conclusion in the case Adam cannot be derived by simply following the form mentioned in [24]:
c(B) +
(12)
(1
Buuiss
Then aim to derive the specific expression for the optimal learning rate with respect to the batch size through the following theorems.
Theorem 3. When B
<
the optimal learning rate is a function with respect to batch size B: 2n²
xa²
Buuisa 2
B Buuise
Copt B))
Σ,
Hiss
Σ.
2x
Hig
ac
(13)
Buster
+
2
B
Broiee
where Broise is a variable unrelated to batch size B:
Σ, His
Bnoise
Pilli
i i
(14)
0:05
Hij
i=
Defining Bpeak as the batch size at which the optimal learning rate reaches a peak value, it is obvious that:
B Bpeak
Bnoise-
(15)
The peak value is:
Breue 2n
a2 Σi
(16)
Σ, Hi,
We prove the thcorem Appendix D. From the theorem we can finally get Eq 2, which implics that there is interval where the batch size becomes larger and the optimal learning rate needs to be
reduced. Considering that 2nd ककद is much larger than normal batch sizes in research and industry (as
shown in Figure 2), this theorem can cover most of the scenarios. To make the conclusion more comprehensive, we also derive the following theorem:
xa² 2p²
Theorem 4. When B >
the optimal learning rate becomes:
Eapt
Σ sign(1+)sign((u)) Hij
(17)

---

We prove the theorem in Appendix E.
Therefore, when B increases infinitely, the optimal learning rate will eventually converge non-zero value. If we make an (unrealistic) assumption that Mi R sign(1s)) we will find that the lower bound of Cmax in Theorem 3 will become the one in Theorem 4, which means that the local peak value of the optimal learning rate is larger than the final convergence value. However, considering that the variance of the gradient in the later stages of training is very small, which makes the above conclusion d, R sign(u:) difficult establish, the stable value in the later stages of training is more likely to exceed the local maximum. We provide reference curve in Figure 1.
2.2 Data/Time Efficiency Trade-off
Following the empirical model for large-batch training [24], we also review the trade-off off between data and time efficiency during batch size selection. We have the following theorem:
Theorem 5. When B < xo² the derived loss improvement with respect to the batch size 2p²²
Ammas
op (B) =
(18)
Buoise B
where Almaz is defined as
HiKi
E, Dj
ALmar
0,0;
Milly
i j
(19)
2 E
oiaj
Hi.
0
i=
We prove the theorem in Appendix F. This result aligns with the conclusion drawn in the SGD situation [24], indicating that many related conclusions also remain valid.
It has been concluded in previous work [24] that, when using the SGD optimizer with the same form as Eq 18, the relationship between training speed (number of steps S) and data efficiency (number of samples E) is given by: ($
S
E
Smin -
=1.
(20)
Emin
Here S(min) represents training speed, the actual (minimum) possible number of steps taken reach a specified model performance; and E(min) represents data efficiency, the actual (minimum) possible number of training examples processed to reach that same level of performance. For more details, please refer to the Eq 2. 11 and the Appendix D in [24]. Additionally, as referenced in the Eq 2.12 in [24] and Eq 1.4 in [25], Buoise is the balance point between training speed and data efficiency:
Emin
B.
Bnoise Rs Berit
(21)
Smin
I
Since in Adam optimizer we arrive at the same Eq 18 as in SGD optimizer, the above equations 20 and 21 still hold. In Adam scenarios, Bpeak Broise is not only the local maximum of the optimal learning rate, but also the balance point between training speed and data efficiency. Moreover, as training progresses and the loss decreases, according to Eq 21, Bpeak will gradually becomes larger.
2.3 Summary
In this section, we have drawn several conclusions from our theoretical analysis, which are summa- rized as follows:
1. As the batch size increases, the optimal learning rate demonstrates decreasing trend within a specified range (Eq 2).
2. The batch size that corresponds to the local maximum optimal learning rate is consistent with the balance point of training speed and data efficiency (Eq 21). As the training progresses and the loss decreases, Byeak will gradually becomes larger.
5

---

3
Experiments
In this section, we carry out series of experiments to corroborate the theoretical scaling law proposed in Section and detail the experimental workloads and configurations in Section 3.1. The process for deriving the estimated variables from our theory is elucidated in Section 3.2. We also showcase and dissect the applicability of our scaling law across a variety of workloads in Section 3.3.
3.1 Experimental Setup
Workloads. In our empirical study, we incorporate 4 open-source workloads that are extensively utilized: (1) training 5-layer CNN model on the Fashion-MNIST [31], which typical CV test case to start with. It consists of 60000 28x28 grayscale images in 10 classes; (2) training a ResNet- 18 model [32] on the Tiny-ImageNet dataset [33], which contains 100000 images of 200 classes (500 for each class) downsized to 64x64 colored images. In each epoch
train the model with
random 10k samples to reduce the overall complexity; (3) training a dense Transformer model [12] (simplified DistilGPT2 [34]) the ELI5 ELI5-Category dataset [35], which is a smaller but newer and categorized version of the original ELIS dataset [36]. It contains 10.5k complex, diverse questions that require explanatory multi-sentence answers. (4) training fine-gratd fine-grained Mixture-of -freeeertst (MoE)
model, similar in structure to Mistral-Mof [37] but with shared experts [38], on the RedPajama- dataset [39]. which contains 30 trillion filtered and deduplicated tokens (100+ trillions raw) from 84 CommonCrav dumps covering 5 languages, along with 40+ pre-computed data quality annotations. These workloads are popular in both academia and industry, covering typical deep learning tasks in the domains of CV and NL..
Batch sizes and learning rates. To showcase the optimal learning rate for each batch size configu- ration, we leverage a grid-search-style experiments set. Each point in the grid search corresponds to certain round with the same configuration but different random number seed. The start point, stop point, and the interval of different workloads are listed in Table 1. In NLP tasks, the term "batch size" refers to the number of tokens in a batch, practiced in related works [25].
Table Grid search configurations.
Adam
Learning Rate Stop
Batch Size
Workload
Round
B1
82
Start
Step
Start
Stop
Step
CNN
0.9
0.999
le-4
le-3
le-4
1
12
1
100
CNN
0.9
0.999
le-4
1e-3
le-4
64
1164
100
100
DistilGPT2
0.9
0.999
1e-5
1.09e-3
1.2e-4
4
114
10
30
DistilGPT2
0.0
0.0
le-5
5.5e-4
66-5
4
114
10
30
ResNet18
0.0
0.0
le-4
7.876e-4
7.65e-5
16
376
33
100
MoB
0.9
0.999
2e-6
6e-5
26-6
192k
12M
1.2M
10
Hyper-parameters. Since we derive the theorems on Adam-style optimizers, we conduct experiments using the Adam optimizer. We experiment on both the "sign of gradient" configuration (81 0.
B2 0) and the default hyper-parameters (31 = 0.9, 32 = 0.999), as shown in Table 1.
Hardware environment. We execute each round of experiments utilizing an NVIDIA A100 card. The training time of each round for the datasets are approximately hour for -NNIST, 1.5 hours for TinylmageNet 2 hours for ELI5-Category and 11 hours for RedPajama-v2. Given our primary focus on the convergence process, the specific hardware environment does not matter in our experiments. Our theoretical and empirical findings can be generalized to other hardware settings. Additionally, some system optimizations [40-43] are also beneficial to enhancing training efficiency.
3.2 Variable Estimation
We try to estimate the value of Bnoise and the expectation of emaz through curve fitting. After using Eq 21 to simplify Eq 20 (see Appendix G for details), we can record the actual possible number of steps taken S and the actual possible number of training examples processed E to reach specified level of performance corresponding to the optimal learning rate of each batch size in the grid search results, and then perform linear fitting to obtain the estimated value of B Broise:
S = Bnoise E
+
(22)
Smin

---

Subsequently, we use the optimal learning rate and batch size of the grid search results to estimate the max optimal learning rate of Adam-style Elemax] Adam
Copt
B Buoise
B Bnoise
E (maz Adam Adom
E
(23)
and SGD-style [emar]sGnn
B Broise
E Blemaz]sCr SGD
E
Copt
1
(24)
B
Previous research [24] represents the SGD optimizer and the Adam optimizer using Eq 24 with a = 1 and a 0.5, respectively. We include these fitted curves comparisons in the following section.
While we use grid search to estimate the value of Buoises practice we can efficiently approximate using the scaling law from previous studies [24, 25]. where Bnoise Berit L. With this approximation, we only need simple search for a pair of (batch size. optimal learning rate) to determine the final hyperparameter (mar- Therefore, the costly grid-search can be avoided.
3.3 Results
Following Section 3.2, we first estimate the variables and fit the curves using observations, then conduct grid-search-style experiments for learning rates and batch sizes.
Figure 2. 3, 3,4, illustrate the experimental results of CNN-FashionMNIIST ResNet18- ResNet18-TinylmageNe DistiGPTT-ELLSCattggn and MoE-Reddaaaam--2, respectively. Each figure is divided into two parts: the left subfigure illustrates the grid-search results for batch sizes and learning rates, and these data points are utilized fit the curve of Eq 22 in the right subfigure. In order to estimate the variables, we train models from scratch using different learning rates and batch sizes, then record the number of steps S and examples in each experiment to achieve an equivalent training loss. Using the recorded S and E, we fit the curve in the right subfigure and obtain the estimated Bnoise. In the left subfigure, upon achieving the desired training loss, all experiments continue to train the same number of steps. Any subsequent decrease in training loss is represented through different colors, as indicated in the color bar. For each batch size, we highlight the optimal learning rate that results in the most significant reduction in training loss. We also plot the batch size Bnoiee that corresponds to the peak optimal learning rate, the fitted SGD curves with a = 0.5 and a = 1 as derived from previous research [24], and the fitted Adam curve derived from our theorems.
For the CNN- ANN-FaahionMNNIST workload, we train exactly 10 more step after achieving the desired training loss. As shown in Figure 2(a), the batch size bound no? task. Given the simplicity of the CNN-FashionMNNIST workload, commonly-used batch sizes are usually smaller than the batch size bound. We plot the situations corresponding Theorem 3 and Theorem 4 in Figure 2(b) and 2(c), respectively. In both cases, the trend predicted by our theory consistent with the actual optimal learning rate performance, showing declining range at small batch sizes and saturation range at large batch sizes.
2nd for Theorem 3 is around 800 in this
For the ResNet 8- Imaelt workload, we train 50 more steps after achieving the desired training loss. We plot the figures for Theorem 3 at different achieved training losses, which represent the progress of training, as shown in Figure 3. The observed optimal learning rates primarily exhibit downward trend after the batch size exceeds the estimated BRoise- Although the SGD curve with a= 0.5, which is claimed by [24] to represent the Adam optimizer, serves as a good approximation in certain cases, it fails to capture the peak optimal learning rate as our Adam curve docs. Comparing that the estimated Broise gradually increases as the training progresses (i.c. training loss decreases), which corroborates the second conclusion in Section 2.3.
the red dashed lines in different figures,
For the DistilGPT--BliSCategge EliSCategory workload, we train 50 more steps after achieving the desired training loss. As shown in Figure 4, we test on two distinct Adam configurations for Theorem 3: the first with 31 = 0.0, 32 0.0, and the second with 3, 0.9, 32 = 0.999. In both configurations, promising learning rates that lead to substantial decrease in loss are consistent with our Adam curve. worth noting that another curve, SGD with a 0.5 [24], also provides a suitable approximation in this scenario. To more clearly demonstrate the accuracy of our theoretical predictions, present
7

---

167
CNN
1.2
1.0-
0.8
0.2
0.0
700
725
750
775
800
825
850
875
900
mo?
(a) Statistical histogram
2
Mayr
D4y 0000
175
Fam 1-
1.00
03
100335
322
410
1200
(b) The situation of small batch size in Theorem 3.
(c) The situation of large batch size in Theorem
Figure 2: Batch size versus optimal learning rate within the context of CNN trained FaShionMNIST.
detailed results from finer-grained grid search in Figure of Appendix H. These experiments demonstrate that our theorems can be generalized to different optimizer configurations, validating the analysis in Appendix A.
For the sparse MoE model using the RedPajama-v2 dataset, we train 300 more steps after achieving the desired training loss. Figure 5 demonstrates that our predictions optimal learning rate are both accurate and appropriate.
In addition to the above workloads, we also conduct an analysis of experimental results from third parties, confirming that our conclusions remain valid. Detailed results are presented in Figure 7 of Appendix H.
4
Discussion
We have carried out empirical studies representative workloads using the Adam optimizer. Our investigation into the scaling laws of learning rates relative to batch sizes has provided deeper insights into the training dynamics of deep learning models. This understanding tune hyperparameters, enhance convergence speeds, and circumvent exhaustive grid searches. By leveraging prior knowledge that the optimal learning rate decreases after reaching peak, researchers and engineers can more effectively adjust the learning rate to achieve efficient training outcomes.
help fine-
In real-world applications, there are numerous different learning workloads [44-47]. Other factors, beyond the scope of this paper, may influence the learning process - the specific optimizer used, weight decay, gradient clipping, etc. While assert that our theorem can be applied to numerous practical scenarios, it may not fully encompass all situations involving intricate training configurations.
As one of our conclusions points out, the variable Bnoise will gradually increases as the training progresses. natural to implement adaptive learning rates (and batch sizes) possible, to speed up the training process. As mentioned in [24], using adaptive batch sizes and warmed-up learning rates brings considerable benefits. Fully exploring the potential of batch size and learning rate scheduling requires meticulous design, which we leave future work.

---

1.22
Bew 24.0
25
1010
25
24-
Laaniig
0.
0.-
04.3
135
1:00
100
230
350
03
03
300
(a) Leraining
5.213 + 0.001
(b) Leraining = 5.146 + 0.001
Maai
Rooe 1832
C4
2-13
1. Male
2-100
64
0.2
203
216-0
200
Batch Als
Andbrocalatt
(c) Tatrining = 5.113 0.001
(d) Leraining = 4.863 0.001
2
6.8
1
35
0.2
0060
Sssth dbs
Liraining 4.741 0.001
Liraining 4.633 + 0.001
Figure 3: The relationship between batch sizes and optimal learning rates within the context of ResNet- 18 trained on TinylmageNet. The red dashed line accurately predicts the peak value, and as the training loss decreases, the peak value gradually shifts to the right.
1.2
Mae
2.00
2-05
1.75
Eap 2.0
12
1.00
675
0.50
04153
08
10.0922
20
40
65
82
100
322
10tTh
KoY produ
Butch des
Rasgeocal
(a)
31
0.0,
32 0.0
(b) 31 0.9, 32 0.999
Figure 4: The relationship between batch sizes and optimal learning rates within the context of DistilGPT2 trained on Eli5Category.

---

0.0006
Normafadd Trainin LOSS-4
0.0004
Loaming rate
0.0003
0.0001
Ours
0.2
Linear
Square
0.0000
0.0
0.2
0.4
0.6
0.8
1.5
12
Batch size
107
Figure 5: Grid search results for the MoE [37, 38) structure model.
Our theory is based on the quadratic approximation of the loss function. Experimental results demonstrate that conclusions drawn from this second-order expansion effectively predict the surge phenomenon observed in most mainstream scenarios. However, recent studies [48-50] have proposed that quadratic approximations do not accurately capture the loss scenarios involving large learning rates. We acknowledge the potential benefits of exploring higher-order approximations and consider this promising direction for future research.
Related Work
Aiming to accelerate convergence, our work analyzes the scaling law of optimal learning rates with respect to batch sizes for Adam-style optimizers. Numerous related studies have been proposed to enhance the convergence of deep learning tasks by investigating optimal learning rates, developing new optimizers, and analyzing gradient noise.
Earlier works have proposed various scaling laws tune learning rates for SGD-style optimizers, including square root scaling [22], linear scaling [19, 231, and others [24]. They also obtained a scaling law for Adam-style optimizers [24, 23] through approximation, revealing square root-like relationship where the optimal learning rate monotonically increases with the batch size, However, as illustrated in Section and their analysis holds only for small batch sizes, whereas the true scaling law exhibits greater complexity, with the optimal learning rate reaching a peak value at balanced batch size.
There are many meticulously designed optimizers for various tasks and scenarios. First-order optimizers dominate nowadays deep learning models, including adaptive methods [3-7], sign-based methods [18, 8), wise methods (for large-batch training) [51 52]. Second-order optimizers [53- 55], though with stronger theoretical guarantees, quadratic complexity with respect to the number of parameters. Despite the emergence of new optimizers, empirical evidence confirms that Adam has remained the most widely used and effective optimizer over the past decade.
not efficient for large-scale models due to
Our analysis is inspired by the empirical model of large-batch training [24], which predicts the useful batch size using the gradient noise scale. Gradient noise can help with learning rate determination [56], batch size selection [57, 58], and gaining deeper insights into the convergence process [59-62].
6 Conclusion
In this paper, we established scaling law between optimal learning rates and batch sizes for Adam- style optimizers. We theoretically proved the optimal learning rate initially increases and then decreases as the batch size grows, and that the peak value of the surge represents trade-off point between training speed and data efficiency. Through extensive experiments, we validated our theory on diverse deep learning models and datasets.
10

---

Acknowledgments and Disclosure of Funding
This work is supported by National Science and Technology Major Project (20222D0116155)) National Natural Science Foundation of China (U22B2037, U23B2048), Beijing Municipal Science and Technology Project (2231 10001033300))) research grant No. SH-2024JK29, PKU-Tencent joint research Lab, and High-pertin -erfhpeeformacce Computing Platform of Peking University. Bin Cui, Shuaipeng Li and Di Wang are the corresponding authors.
References
[1] Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. nature, 521(7553):436 444, 2015.
[2] Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initialization and momentum deep learning. ICML, pages 1139-1147, 2013.
[3] John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. JMLR, 12(7), 2011.
[4] Matthew Zeiler, Adadelta: an adaptive learning rate method. arXiv preprint 12Xiv:1222.57000 2012.
[5] Tijmen Tieleman and Geoffrey Hinton. Rmsprop: Divide the gradient by running average of its recent magnitude. coursera: Neural networks for machine learning. COURSERA Neural Networks Mach. Learn, 17, 2012.
[6] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR, 2015.
[7] Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In ICML, pages 4596 4604, 2018.
[8] Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsich, Yifeng Lu, et al. Symbolic discovery of optimization algorithms. NeurIPS, 2023.
[9] Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Ne u r IPS,
30, 2017.
[10] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x 16 words: Transformers for image recognition scale. In ICLR, 2021.
[11] René Ranftl, Alexey Bochkovskiy, and Vladlen Koltun. Vision transformers for dense prediction.
In ICCV, pages 12179-122 12188, 2021.
[12] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. NeuriPS, 30, 2017.
[13] Jacob Devlin, Ming- Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. NAACL-HLT, pages 4171-4186, 2019.
[14] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint 13X, 2023.
[15] An Wang, Xingwu Sun, Ruobing Xie. Shuaipeng Li, Jiaqi Zhu, Zhen Yang, Pinxue Zhao, IN Han. Zhanhui Kang, Di Wang, et al. Hmoe: Heterogencous mixture of experts for language modeling. arXiv preprint arXiv:2408. 10681, 2024.
[16] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.

---

[17] Hailin Zhang, Yujing Wang, Qi Chen, Ruiheng Chang, Ting Zhang, Ziming Miao, Yingyan Hou, Yang Ding, Xupeng Miao, Haonan Wang, al. Model-enhanced vector index. NeurIPS, 36, 2024.
[18] Jeremy Bernstein, Yu-Xiang Wang, Kamyar Azizzadeneshelii and Animashree Anandkumar. signsgd: Compressed optimisation for non-convex problems. In ICML, pages 560-569. PMLR, 2018.
[19] Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in 1 hour. arXiv preprint arXiv: 011.17006722777 2017.
[20] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In /CML, pages 8748-8763. PMLR, 2021.
[21] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidr, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023.
[22] Alex Krizhevsky. One weird trick for parallelizing convolutional neural networks. arXiv preprint arXiv: 1404597,, 2014.
[23] Diego Granziol, Stefan Zohren, and Stephen Roberts. Learning rates as a function of hatch size:
A random matrix theory approach to neural network training. JMLR, 23(173):1-65, 2022.
[24] Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAl Dota Team, An empirical model of large-batch training. arXiv preprint 1rXiv:1812.06. 2018.
[25] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXhv:2001.08361, 2020.
[26] Lukas Balles and Philipp Hennig. Dissecting adam: The sign, magnitude and variance of stochastic gradients. In ICML, volume 80, pages 413-422, 2018.
[27] Stephan Mandt, Matthew Hoffman, and David Blei. A variational analysis of stochastic gradient algorithms. ICML, pages 353. PMLR, 2016.
[28] Stanislaw Jastrzıbski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors inftuencing minima in sgd. arXiv preprint arXiv:1711.04623. 2017.
[29] Qianxiao Li, Cheng Tai, and E Weinan. Stochastic modified equations and adaptive stochastic gradient algorithms. In ICML, pages 2101-2110, 2017.
[30] Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects. In ICML, volume 97, pages 7654-7663, 2019.
[31] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. CoRR, abs/1708.07747, 2017.
[32] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In CVPR, pages 770-778, 2016.
[33] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In CVPR, pages 248-255, 2009.
[34] Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. In NeurIPS EMC Workshop, 2019.
[35] Jingsong Gao, Qingren Zhou, and Rui Qiu. ELI5-Category: a categorized open-domain qa dataset. 2021.
12

---

[36] Angela Fan, Yacine Jernite, Ethan Perez. David Grangier, Jason Weston, and Michael Auli. 2019.
EL15: long form question answering. In ACL, pages
[37] Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv preprini arXiv:2401 2401.00088,. 2024.
[38] Damai Dai, Chengqi Deng, Chenggang Zhao, et al. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models. In ACL, pages 1280-1297, 2024.
[39] Together Computer. Redpajama: an open dataset for training large language models, 2023.
[40] Guozheng Wang, Yongmei Lei, Zeyu Zhang, and Cunlu Peng. A communication efficient admm-based distributed algorithm using two-dimensional torus grouping allreduce. Data Science and Engineering, 8(1):61-72, 2023.
[41] Xian-He Sun and Xiaoyang Lu. The semory-bounded. speedup model and its impacts in computing. Journal of Computer Science and Technology, 38(1):64-79, 2023.
[42] Keshi Ge, Yiming Zhang, Yongquan Fu, Zhiquan Lai, Xiaoge Deng, and Dongsheng Li. Accelerate distributed deep learning with cluster-aware sketch quantization. Science China Information Sciences, 66(6):162102, 2023.
[43] Xupeng Miao, Xiaonan Nie, Hailin Zhang, Tong Zhao, and Bin Cui. Hetu: A highly efficient automatic parallel distributed deep learning system. Science China. Information Sciences, 66(1):117101, 2023.
[44] Penghao Zhao, Hailin Zhang, Qinhan Yu, Zhengren Wang, Yunteng Geng, Fangcheng Fu, Ling Yang, Wentao Zhang, and Bin Cui. Retrieval-augmented generation for ai-generated content: A survey, arXiv preprint arX11:2402.19473, 2024.
[45] Fanglue Zhang, Junhong Zhao, Yun Zhang, and Stefanie Zollmann. A survey on 360 images and videos in mixed reality: algorithms and applications. Journal of Computer Science and Technology, 38(3):473 2023.
[46] Hailin Zhang, Penghao Zhao, Xupeng Miao, Yingxia Shao, Zirui Liu, Tong Yang, and Bin Cui Experimental analysis of large-scale learnable vector storage compression. VLDB.
[47] Hailin Zhang, Zirui Liu, Boxuan Chen, Yikai Zhao, Tong Zhao, Tong Yang, and Bin Cui. Cafe: Towards compact, adaptive, and fast embedding for large-scale recommendation models. Proceedings of the ACM on Management of Data, 2(1): 2024.
[48] Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In ICLR, 2021.
[49] Jeremy M Cohen, Behrooz Ghorbani, Shankar Krishnan, Naman Agarwal, Sourabh Medapati, Michal Badura, Daniel Suo, David Cardoze, Zachary Nado, George E Dahl, et al. Adaptive gradient methods at the edge of stability. arXiv preprint arXiv:2207 144844, 2022.
[50] Alex Damian, Eshaan Nichani, and Jason D Lee. Self-stabilization:: The implicit bias of gradient descent at the edge of stability. arXiv preprint arXiv:2209. 15594, 2022.
[51] Yang You, Zhao Zhang, Cho-Jui Hsieh, James Demmel, and Kurt Keutzer. Imagenet training in minutes. In Proceedings of the 47th international conference on parallel processing, pages
1-10, 2018
[52] Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song. James Demmel, Kurt Keutzer, and Cho-Jui Hsich. Large batch optimization for deep learning: Training BERT in 76 minutes. In ICLR, 2020.
[53] Zhewei Yao, Amir Gholami, Sheng Shen, Mustafa Mustafa, Kurt Keutzer, and Michael Mahoney. Adahessian: An adaptive second order optimizer for machine learning. In AAAI, volume 35, pages 10665-10673, 2021.
13

---

[54] Jimmy Ba, Roger Grosse, and James Martens. Distributed second-order order optimization using kronecker-factorne approximations. In Iclr, 2022.
[55] Hong Liu, Zhiyuan Li, David Hall, Percy Liang, and Tengyu Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training. In /CLR, 2024.
[56] Tom Schaul, Sixin Zhang, and Yann LeCun. No more pesky learning rates. In ICML, pages 343-351, PMLR, 2013.
[57] Lukas Balles, Javier Romero, and Philipp Hennig. Coupling adaptive batch sizes with learning rates. In UAI, 2017.
[58] Samuel Smith and Quoc V. Le. bayesian perspective on generalization and stochastic gradient descent. In ICLR, 2018.
[59] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In ICLR, 2017.
[60] Ravid Shwartz-Ziv and Naftali Tishby. Opening the black box of deep neural networks via information. arXiv preprint arXiv:/703. /00.0000000000, 2017.
[61] Umut Simsekli, Levent Sagun, and Mert Gurbuzbalaban. A tail-index analysis of stochastic gradient noise deep neural networks. ICML, pages 5827-5837, 2019.
[62] Christopher Shallue, Jachoon Lee, Joseph Antogmini. Jascha Sohl-Dickstein Roy Frostig, and George E Dahl. Measuring the effects of data parallelism on neural network training. JMLR, 20(1 (((112)11444 2019.
[63] Xiao Bi, Deli Chen, Guanting Chen, Shanhuang Chen, Damai Dai, Chengqi Deng, Honghui Ding, Kai Dong, Qiushi Du Zhe Fu, al. Deepseek Ilm: Scaling open-source language models with longtermism. arXiv preprint arXiv:2401 00244494,. 2024.
14

---

A Parameter Update Amount in the Adam Optimizer
The update amount in the Adam optimizer consists of two parts, the first moment:
mh = 31m4-1 + (1 - B1 Gest,t
m
(25)
m 1- 81
and the second moment:
U/ 12
+ (1 B2) 32)
Ut
(26)
02 1
where my = and th = 0. The final update amount is
1-8 I-3
Σ 87 GGst,
my
(27)
+ Adam
1- 21 0/ V
G²
Geatt CAdam
1-
After ignoring the role of Adams when B1 - and B2 1, Eq 27 transforms to
Σ
Ee
(28)
E.
Gestl
V1
(aaee Geal
The equation is obtained by using var,(G = E, E2 Gesel Note that the expected value
E, [Gest] is over the iteration distribution, over the data distribution. Obviously, when the variance of Gest is small, the update amount is approximately sign(Gest).
On the other hand, when B1 - 0 and $2 + 0. Eq 27 can be simplified to
Gest
sign(Grst)
(29)
V est
Therefore, we can approximate the parameter update amount of the Adam optimizer as sigm (Gest). without affecting the theoretical conclusion.
B Proof of Lemma 1
Proof. We can perturh the parameters by some vector V with learning rate e, and approximate the true loss using quadratic expansion in terms of via Taylor expansion:
L(e - e V) = L( - e + 2e 2 HV.
(30)
Consider the expected value of loss improvement over a data distribution p(x) over data points æ:
1
EJAL] = (0) L(O v)] ~ e [[
- HV].
(31)
2
By maximizing the expected loss improvement, we obtain the optimal learning rate in Eq and the optimal loss improvement in Eq 7.
C Proof of Theorem 2
Proof. Consider random variable NNw, o2), let Fー and -
then
1
g dx
E E[sign(x))
sign(x)
o 0 nT
-0
sign(p + "oyee 3/4 dy
v 2m
(32)
1/4 dy
c c dy)
V2T
= (1 - (())) (22) = er
H
V20
15

---

and the variance is
var(sign(x)) = E(sign(a))] - E[sign(x)] = erf
(33)
/20
where + represents the cumulative distribution function of the standard normal distribution, and er the Gauss error function that is defined as erf(z) = / S0 e dt.
Given that the gradient G2 of any data point x in the data distribution p(x) relative to certain parameter 8; follows Gaussian distribution with mean b4 and variance o. then
B Gz
of p4:
(e))
N
(34)
B
B
Therefore, when V sign(Gest), using Eq and Eq 33 we can get the expectation
B[V]
erf (V
(35)
2 2
and the covariance matrix
cou(V)
1 er (v Bui
2
(36)
Given that the real gradient satisfies:
G
Pi
(37)
and define E, as function with respect to the token batch size B, based on the Gauss error function:
(v
B (
2
p4
=(B)=erff erf
dt,
(38)
2 o.
substituting the expectation and the variance of V in Eq 6 and Eq 7 from Lemma using the above equations, can get Eq 8 and Eq respectively.
To simplify the subsequent computation, we approximate the function & by other sigmoid-lin analytical forms. Specifically, we find that (B) ~
di
which results in Eq 10.
2B
D Proof of Theorem 3
Proof. When B <
Eq 10 reduces to: 2/
T
2B Pi
E,(
R
(39)
T
Substituting the function E, in Eq 9 using the above equation, and defining Broise and eman in Eq and 16, we can deduce the relationship between the optimal learning rate and the token batch size:
16

---

Ei
2B Mr 14 01
Copt (B)
Milli
i i j
Σ, His + 2B Σ, E,
0:dj
Hisj
i=
Σ, V 22 m7
"
Σ, Hi,
Milli
i i Had
2E, Σ,
0:0g
i= j
1+ B
7 -2 H4,2
(40)
VB
V
2 Σ
1 1+
B
Σ
H4g
Busise
Boiie
2 Σ
Σ, His
Bagian
B
B
Broise
Emax
S Cmaz.
B Bagise
B Brosse
should be noted that since Bnoise
and
Cmax then both E, Hes and
Milly D:og
itj
Σ,Σ,
H4s greater than O. Therefore, Emaz can be expressed as form that
==
does not contain Bnoise:
V Brosse E, u? 25
Emax
E; Hipi
Ez, His
Σ P
Σ, H
Lilli O109
i j
2,
Σ, E,
His
i i= j
Σ. !
(41)
Mili auiyy
itj
2
2 Σx E,
His VE, Hi.
i=j
Bally a.o
i i j
Σ,E,
Hi.j
i=
The last incquality is derived using the AM-GM inequality (a² + 62 2 2ab).
E
Proof of Theorem 4
Proof. When
Eq 10 converges to:
E sign 14
(w:).
(42)
0,
Substituting the function E, in Eq 9, obtain Eq 17.
17

---

Proof of Theorem 5
Proof. When B <
2n we have the approximate results in Eq 39 from Theorem 3. Substituting &;
and Booise using Eq 39 and 14, and defining ALmaz in 19, the optimal loss improvement in Eq can be expressed as:
I
E,
Aloopt (B)
i
i
H. Σ,E,
0100
His
i=
Σ E₁ 0:01
Σ. Mi. + +2 +2E,E,
i
His
B
10
i= j
(43)
1213
E Σ3
1
0.03
Buaus B
+ 1
+ 5 H.s) i=
Ei
0:03
A Jmma
ALmma
Baase
+
B
Following the Appendix D in [24], this allows both the total number of steps and data examples processed to still be written as
s / (1 Bnoise
ds
B
(44)
F / (Buoise
B)ds
Therefore the conclusion in [24]
Smin = f
ds
(45)
Emin Buoiseds
and Eq 20, 21 still hold.
G Variable Estimation for Data/Time Efficiency Relationship Equation
When considering S, E, Smin and Emin > Bq 20 can be simplified to
S -11
---1)==
E
=1
Smin
Emin
SE Smin F. SFmin + S SuinFimn
Smin + SEmin SE
(46)
Smin
Emin
=
E
1
Emin 1
5
Smin E
Smin
so linear fit to the relationship between & and E can estimate Smin and E Emin-
H Additional Experiments Results
18

---

0,00040
0.00035
5.25
0.00100
5.24
0,00025
5.23
0.00020
5.22
100
115
Figure 6: Finer-grained grid search results for the experiments shown in Figure 4(b).
0.0050
2.675
Optimal Lcarning rats cyd18
Sydl8
3.0
2.600
0.0040
4.8
2,625
2.575
0,0025
0 0015
4.2
2.500
0.0015
4.0
0.000
2.500
0.0010
0.0055
2.475
0.2
0.4
0.5
0.8
1.0
1-2
Batth viee 2
1cG
(a) 1e17 FLOPs (177M FLOPs/token)
(b) 1e20 FLOPs (2.94B FLOPs/token)
Figure The optimal learning rates, based on the results presented Deekseek paper [63], align with our theorems.
Figure 8: Examples of gradient distributions observed during the training of an MoE structure model, which approximate Gaussian distributions.
19

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The claims provide coherent and consistent overview of the research objec-
tives, methodologies, and findings.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: See Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
20

---

Justification:
See Section 2 and Appendix A, B, C.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: See Section 3.1.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
21

---

The authors should provide scripts to reproduce all experimental results for the new reproducible, they Answer: [NA]
Justification: The model structure and data required for the experiments are publicly avail- able. And the necessary experimental configuration has been provided in Section 3.1 for reproduction.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (kttps://nips.ce public/ Seellii/ggiee///oooo SooodeSubiisiionooe
for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / //nips. wwtpcccccpggpuucccceecccc wwtpe//iiigii///eeeeeeeeeeeeeeee htte://oooSuumisssSeeeeee ssionPolicy) for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification:
The necessary experimental configurations have been provided in Section 3.1
for reproduction.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer:
[Yes]
Justification: See Section
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
22

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: See Section 3.1.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https pettps:/nnerre http///eerrrpp.eeeeeeeeee httppubccc//hhicccccuueeeeeee
Answer: [Yes]
Justification:
The research conforms the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: There no societal impact of the work performed. This work a foundational research and is not tied to particular applications.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
23

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We cite the creators and provide necessary information in Section 3.1.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
25


Reviews:
[{"review_id": "gVhOVMYQvJ", "reviewer": "Reviewer_ytaJ", "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.", "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?", "limitations": "It would be interesting to see results on a larger variety of modern and widely used datasets and architectures.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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).\n- The paper is written clearly and is well organized. I appreciate the “summary” notes included by the authors\n- 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). \n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough review and constructive suggestions. We address your comments in the following paragraphs.\n\n> 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).\n\nA1. 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.\n\n\n> 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.\n\nA2. 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.\n\n\n> 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?\n\nA3. 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.\n\n\nThanks again for appreciating our work and for your constructive suggestions. Please let us know if you have further questions.\n\n---\n\n[1] DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models, 2024, Dai et al.\n\n[2] Mixtral of Experts, 2024, Jiang et al.\n\n[3] DeepSeek LLM Scaling Open-Source Language Models with Longtermism, 2024, Bi et al.\n\n[4] An Empirical Model of Large-Batch Training, 2018, McCandlish et al. (reference [25] in our paper)", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VGYrCBRvX0", "reviewer": "Reviewer_JT4e", "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.", "questions": "1. Why does the theorem apply to general network optimization?\n\n2. How to make the result of the theorem in use?", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "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.\n\n2. 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.\n\n2.  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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n\n> 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.\n> Q1. Why does the theorem apply to general network optimization?\n\nA1. 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.\n\n\n> W2. 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.\n> Q2. How to make the result of the theorem in use?\n\nA2. 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.\n\n\nThanks 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.\n\n---\n\n[1] An Empirical Model of Large-Batch Training, 2018, McCandlish et al. (reference [25] in our paper)\n\n[2] Scaling Laws for Neural Language Models, 2020, Kaplan et al. (reference [26] in our paper)", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wui3XNkofR", "reviewer": "Reviewer_ASLs", "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. \n\n[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\n\n1. 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.\n2. How is Fig. 3 generated? The linear fits look very strange. \n3. 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.\n4. Is there any characterization of the **intermediate** (i.e. large but not infinite batch) behavior of the scaling?\n\nIt would be helpful to have an analog of Figure 1 for the empirical data.", "limitations": "Yes.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 2, "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. \n\n[1] McCandlish, S., Kaplan, J., Amodei, D., & Team, O. D. (2018). An empirical model of large-batch training. arXiv preprint arXiv:1812.06162.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our paper and for your valuable insights. We address your comments as follows.\n\n> W1. The presentation is not great. A lot is assumed from [1], but it would make reading easier to make things more self-contained.\n\nA1. 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.\n\n\n> W1. It is unclear what is going on in Figure 1 between the solid and dashed Adam curves.\n> 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.\n\nA2. 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.\n\n\n> 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.\n> Q2. How is Fig. 3 generated? The linear fits look very strange.\n\nA3. 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.\n\n\n> W1. The takeaways and consequences for a practitioner are unclear.\n> 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.\n\nA4. 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.\n\n\n> Q4. Is there any characterization of the intermediate (i.e. large but not infinite batch) behavior of the scaling?\n\nA5. 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.\n\n\n> W1. Minor: the Latex parentheses look sloppy.\n\nA6. We will improve the readability of our paper. \n\n\nThanks 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.\n\n---\n\n[1] An Empirical Model of Large-Batch Training, 2018, McCandlish et al. (reference [25] in our paper)\n\n[2] Scaling Laws for Neural Language Models, 2020, Kaplan et al. (reference [26] in our paper)", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "O4S7VgLQBT", "reviewer": "Reviewer_PmLA", "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.", "questions": "* What is the practical implication of the decrease in the optimal learning rate after the peak? \n* Given that the peak shifts through training, can the authors propose guidelines for scaling the learning rate and batch size? \n* What is the intuition behind the training speed and data efficiency relationship being the same as the SGD case? \n* 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?\n* 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)?", "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]. \n* 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.\n\n\n[1] Gradient descent on neural networks typically occurs at the edge of stability\nCohen et al. (2021)\narXiv:2103.00065\n\n[2] Adaptive gradient methods at the edge of stability\nCohen et al. (2022)\narXiv:2207.14484\n\n[3] Self-stabilization: The implicit bias of gradient descent at the edge of stability\nDamian et al. (2022)\narXiv:2209.15594\n\n[4] Linear attention is (maybe) all you need (to understand transformer optimization)\nAhn et al. (2022)\narXiv:2310.01082", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "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.\n* 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.\n* 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.\n* 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n\n> W1. \n\nA1. 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.\n\nRegarding 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.\n\n\n> W2. \n\nA2. 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.\n\n\n> W3. \n\nA3. 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.\n\n\n> Q1. \n\nA4. 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.\n\n\n> Q2. \n\nA5. 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.\n\n\n> Q3. \n\nA6. 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. \n\n\n> Q4. \n\nA7. 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.\n\nBeyond 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.\n\n\n> Q5. \n\nA8. 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.\n\n\n> L1. \n\nA9. 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.\n\n\n> L2. \n\nA10. 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.\n\n\nThanks 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.\n\n---\n\n[1] An Empirical Model of Large-Batch Training, 2018, McCandlish et al. (reference [25] in our paper)\n\n[2] Scaling Laws for Neural Language Models, 2020, Kaplan et al. (reference [26] in our paper)", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Dear reviewers,

Thank you for your reviews and constructive suggestions. We have incorporated additional experimental analyses to strengthen our conclusions.


**Experimental Analyses**
We made experimental analyses on both sparse MoE and dense structures. For the sparse MoE structure, in Experiment-1, we experimented on a fine-grained MoE model with shared experts[1], which has a similar structure to Mistral-MoE[2], a representative of the sparse models. As for the dense model, in Experiment-2 and Experiment-3, we validated our proposed conclusions with a more fine-grained experiment on DistilGPT-2, as well as an experiment from another paper[3], which showed aligned results with our theorems. The following is our specific experimental configurations and results.

---

**Experiment - 1**

We performed a grid search on a 500M parameter MoE model with 16 experts using the RedPajama-v2 dataset.

| Key                 | Value          | Key                 | Value          |
|---------------------|----------------|---------------------|----------------|
| `VOCAB_SIZE`        | 32000          | `Adam Beta1`        | 0.9            |
| `N_LAYERS`          | 12             | `Adam Beta2`        | 0.999          |
| `MAX_SEQ_LENGTH`    | 4096           | `Weight decay`      | 0.00001        |
| `D_MODEL`           | 768            | `MAX_GRAD_NORM`     | 1              |
| `N_HEADS`           | 12             | `AUX_LOSS_ALPHA`    | 0.1            |
| `HEAD_DIM`          | 64             | `WSD warmup`        | 2000           |
| `EXPERTS_TOTAL_DIM` | 12288          | `WSD decay ratio`   | 0.1            |
| `TOP_K_ROUTER`      | 4              | `Dataset`           | RedPajama-v2   |
| `N_EXPERTS`         | 16             |


In the experiments, we aimed to find the largest best learning rate across token batch sizes (token bs, i.e., considering the number of tokens in a batch). For simplicity, we fitted the final training loss at different learning rates using a quadratic function of the form $f(x) = Ax^2+Bx+C$, which served as a good approximation. Subsequently, we utilized the extremum $-\frac{B}{2A}$ of the quadratic function to estimate the optimal learning rate.

We have presented the fitted curve in the attached file Fig. 1. Our method aligns well with the observed data points. Other methods, including linear scaling and square root scaling, cannot align with the best learning rates.

The results reaffirm our original conclusions: the optimal learning rate first increases and then decreases as batch size grows. This further validates our scaling law between learning rate and batch size for Adam.

---

**Experiment - 2**

We conducted a detailed grid search on DistilGPT-2 using the ELI5 category dataset. The results are plotted as Fig. 2 in the attached file.

---

**Experiment - 3**

Additionally, we note that a publicly available paper[3] provides grid search results that align with our theoretical predictions. We have plotted the fitted curves in Fig. 3 of the attached file. For better visualization, we removed the data points with large loss values in the experimental results. Notably, the $B_{crit}$ value is approximately 0.1M for (a) and 1.5M for (b), estimated using the Scaling Law $B_{crit} = \frac{B^*}{L^{\frac{1}{\alpha}}}$ , matching the peak position of the optimal learning rate observed in their experiments. This observed trend also conforms to the pattern depicted by the orange solid line in Fig. 1 of our paper, further validating our theoretical results. 

---

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

Meta-review: This paper elucidate the connection between optimal learning rates and batch sizes for Adam-style optimizers through both theoretical analysis and extensive experiments (I quoted). All reviewers agree to accept this paper because (1) the method is well motivated, (2) experimental results are impressive and support the theory, (3) providing new insights to the scaling rule in Adam, etc. I agree with them, and thus recommend an accept.
Decision: Accept (poster)