Large Batch Optimization for Deep Learning: Training BERT in 76 minutes
Authors: Yang You
Jing Li
Sashank Reddi
Jonathan Hseu
Sanjiv Kumar
Srinadh Bhojanapalli
Xiaodan Song
James Demmel
Kurt Keutzer
Cho-Jui Hsieh
Venue:  (2020)
Keywords: large-batch optimization
distributed training
fast optimizer

Abstract: Training large deep neural networks on massive datasets is  computationally very challenging. There has been recent surge in interest in using large batch stochastic optimization methods to tackle this issue. The most prominent algorithm in this line of research is LARS, which by  employing layerwise adaptive learning rates trains ResNet on ImageNet in a few minutes. However, LARS performs poorly for attention models like BERT, indicating that its performance gains are not consistent across tasks. In this paper, we first study a principled layerwise adaptation strategy to accelerate training of deep neural networks using large mini-batches. Using this strategy, we develop a new layerwise adaptive large batch optimization technique called LAMB; we then provide convergence analysis of LAMB as well as LARS, showing convergence to a stationary point in general nonconvex settings. Our empirical results demonstrate the superior performance of LAMB across various tasks such as BERT and ResNet-50 training with very little hyperparameter tuning. In particular, for BERT training, our optimizer enables use of very large batch sizes of 32868 without any degradation of performance.  By increasing the batch size to the memory limit of a TPUv3 Pod, BERT training time can be reduced from 3 days to just 76 minutes.

Published as a conference paper at ICLR 2020
LARGE BATCH OPTIMIZATION FOR DEEP LEARNING: TRAINING BERT IN 76 MINUTES
Yang You", Jing Li', Sashank Reddi', Jonathan Hseu', Sanjjv Kumar", Srinadh Bhojanapalli Xiaodan Song'. James Demmel', Kurt Keutzer?, Cho-Jui Hsieh 1.3
Yang You student researcher Google Brain. This project done when he Google Brain.
Google', UC Berkeley", UCLA
(youyang. demmel, keutaer) (ecs. berkeky.dddd (jingli, sasbank, sanjivk, berinadh, xiaodansong, chojui) google.com
ABSTRACT
Training large deep neural networks on massive datasets is computationally very challenging. There has been recent surge in interest in using large batch stochastic optimization methods to tackle this issue. The most prominent algorithm in this line of research is LARS, which by employing layerwise adaptive learning rates trains RESNET on ImageNet in a few minutes. However, LARS performs poorly for attention models like BERT, indicating that its performance gains are not consistent across tasks. In this paper, we first study a principled layerwise adaptation strategy to accelerate training of deep neural networks using large mini-batches. Using this strategy, develop new layerwise adaptive large batch optimization technique called LAMB; we then provide convergence analysis of LAMB as well as LARS, showing convergence to a stationary point in general nonconvex settings. Our empirical results demonstrate the superior performance of L.MBB across various tasks such as BERT and RESNET-50 training with very little hyperparameter tuning. In particular, for BERT training, our optimizer enables use of very large batch sizes of 32868 without any degradation of performance. By increasing the batch size to the memory limit of TPUv3 Pod, BERT training time can be reduced from days to just 76 minutes (Table 1). The LAMB implementation is available online'.
1 INTRODUCTION
With the advent of large scale datasets, training large deep neural networks, even using computation- ally efficient optimization methods like Stochastic gradient descent (SGD), has become particularly challenging. For instance, training state-of-the-an deep learning models like BERT and ResNet-50 takes 3 days on 16 TPUv3 chips and 29 hours on 8 Tesla P100 gpus respectively (Devlin et al., 2018; He al., 2016). Thus, there growing interest to develop optimization solutions tackle this critical issue. The goal of this paper is to investigate and develop optimization techniques to accelerate training large deep neural networks, mostly focusing approaches based on variants of SGD.
Methods based on SGD iteratively update the parameters of the model by moving them in a scaled (negative) direction of the gradient calculated on minibatch. However, SGD's scalability is limited by its inherent sequential nature. Owing to this limitation, traditional approaches to improve SGD training time in the context of deep learning largely resort to distributed asynchronous setup (Dean al., 2012; Recht et al., 201 1). However, the implicit staleness introduced due the asynchrony limits the parallelization of the approach, often leading to degraded performance. The feasibility of computing gradient large minibatches in parallel due to recent hardware advances has the resurgence of simply using synchronous SGD with large minibatches as an alternative to asynchronous SGD. However, naively increasing the batch size typically results in degradation of generalization performance and reduces computational benefits (Goyal al., 2017).
Synchronous SGD on large minibatches benefits from reduced variance of the stochastic gradients used in SGD. This allows one to use much larger learning rates in SGD, typically of the order square root of the minibatch size. Surprisingly, recent works have demonstrated that up to certain minibatch sizes, linear scaling of the learning rate with minibatch size be used to further speed up the
'nttps: :/ /// //ggttttubbb.. http:///eoooeeooooo/oo//o//////e optim mi zers/ amb.py

---

Published as a conference paper at ICLR 2020
training Goyal et al. (2017). These works also elucidate two interesting aspects enable the use of linear scaling in large batch synchronous SCD: (i) linear scaling of learning rate is harmful during the initial phase; thus, hand-tuned warmup strategy of slowly increasing the learning rate needs to be used initially, and (ii) linear scaling of learning rate can be detrimental beyond a certain batch size. Using these tricks. Goyal et al. (2017) was able to drastically reduce the training time of ResNet-50 model from 29 hours to hour using batch size of 8192. While these works demonstrate the feasibility of this strategy for reducing the wall time for training large deep neural networks, they also highlight the need for adaptive learning rate mechanism for large batch learning.
Variants of SGD using layerwise adaptive learning rates have been recently proposed to address this problem. The most successful in this line of research is the LARS algorithm (You et al., 2017), which was initially proposed for training RESNET. Using LARS, ResNet-50 can be trained on ImageNet in just few minutes! However, it has been observed that its performance gains are not consistent across tasks. For instance, LARS performs poorly for attention models like BERT. Furthermore, theoretical understanding of the adaptation employed in LARS is largely missing. To this end, we study and develop new approaches specially catered to the large batch setting of our interest.
Contributions. More specifically, we make the following main contributions this paper.
Inspired by LARS, we investigate a general adaptation strategy specially catered to large batch learning and provide intuition for the strategy.
Based on the adaptation strategy, we develop new optimization algorithm (LAMB) for achieving adaptivity of learning rate in SGD. Furthermore, we provide convergence analysis for both LARS and LAMB to achieve stationary point in nonconvex settings. We highlight the benefits of using these methods for large batch settings.
We demonstrate the strong empirical performance of LAMB across several challenging tasks. Using LAMB we scale the batch size in training BERT to more than 32k without degrading the performance; thereby, cutting the time down from 3 days to 76 minutes. Ours is the first work to reduce BERT training wall time to less than couple of hours.
We also demonstrate the efficiency of LAMB for training sfate-of-the-art image classification models like RESNET. To the best of our knowledge, ours is first adaptive solver that can achieve state-of-the ahet accuracy for RESNET-50 as adaptive solvers like Adam fail to obtain the accuracy of SGD with momentum for these tasks.
1.1 RELATED WORK
The literature on optimization for machine learning is vast and hence, we restrict our attention to the most relevant works here. Earlier works on large batch optimization for machine learning mostly focused on convex models, benefiting by a factor of square root of batch size using appropriately large learning rate. Similar results can be shown for nonconvex settings wherein using larger minibatches improves the convergence to stationary points; albeit at the cost of extra computation. However, several important concerns were raised with respect to generalization and computational performance in large batch nonconvex settings. It was observed that training with extremely large batch was difficult (Keskar et al., 2016; Hoffer et al., 2017). Thus, several prior works carefully hand-tune training pyper-parameters, like learning rate and momentum, to avoid degradation of generalization performance (Goyal et al., 2017; Li, 2017; You et al., 2018; Shallue et al., 2018).
(Krizhevsky, 2014) empirically found that simply scaling the learning rate linearly with respect to batch size works better up to certain batch sizes. To avoid optimization instability due to linear scaling of learning rate, Goyal et al. (2017) proposed a highly hand-tuned learning rate which involves a warm-up strategy that gradually increases the LR to a larger value and then switching to the regular LR policy (e.g. exponential or polynomial decay). Using LR warm-up and linear scaling, Goyal et al. (2017) managed to train RESNET-50 with batch size 8192 without loss in generalization performance. However, empirical study (Shallue et al., 2018) shows that learning rate scaling heuristics with the batch size do not hold across all problems across all batch sizes.
More recently, to reduce hand-tuning of hyperparameters, adaptive learning rates for large batch training garnered significant interests (Reddi et al., 2018; Zaheer et al., 2018; Zhang et al., 2019). Several recent works successfully scaled the batch size to large values using adaptive learning rates without degrading the performance, thereby, finishing RESNET-SE training on ImageNet in a few minutes (You et al., 2018: 8; landola et al., 2016; Codreanu et al., 2017; Akiba et al., 2017; Jia et al.,
2

---

Published as a conference paper at ICLR 2020
2018; Smith al., 2017; Martens & Grosse, 2015; Devarakonda et al., 2017; Mikami al., 2018; Osawa et al., 2018; You et al., 2019; Yamazaki et al., 2019). To the best of our knowledge, the fastest training result for RESNET-50 on ImageNet is due to Ying et al. (201 8), who achieve 76+% % top-1 accuracy. By using the LARS optimizer and scaling the batch size 32K on a TPUv3 Pod. Ying et al. (2018) was able to train RESNET-50 on ImageNet in 2.2 minutes. However, it was empirically observed that none of these performance gains hold in other tasks such BERT training (see Section 4).
PRELIMINARIES
Notation. For any vector It E R. either 21, or [zely are used to denote its j"h coordinate where j E [d]. Let be the d x d identity matrix, and let I = l,., In] be its decomposition into column submatrices I1; = d x da. For x E R let ( be the block of variables corresponding to the columns of I, ie.., (0))) = I x E R4 for i = {1.2, .h.. For any function f : F + R. we use V.f(r) to denote the gradient with respect to و! For any vectors u, U E Rd, we use u² and u/v to denote elementwise square and division operators respectively. We use ||-|| and ||-I11 to denote nn-m and 4 l--rorm of a vector respectively.
We start our discussion by formally stating the problem setup. In this paper, we study nonconvex stochastic optimization problems of the form
min f(x) == Eg-e [(( s) + |>]|
(I)
where é is smooth (possibly nonconvex) function and P a probability distribution on the domain C R Here, x corresponds to model parameters, is the loss function and is an unknown data distribution.
We assume function ((x) is Ly-smooth with respect to 7th block, i.e., there exists a constant L, such that
[V v,l(x 1,, s)ll s 1, Vr IR, R. and E S,
for all i E [h]. We use L (L...... _A) T to denote the h-dimensional vector of Lipschitz constants.
We use Los and Lavg denote max; L, and Σ, 4 respectively. We assume the following bound on the variance stochastic gradients: EE| V Furthermore, we also assume [][[V[(x,,))
v.f(x) of for x E IR4 and i E [M]] [+()))+11? S o? for all x E Rdd and [d]. We use (o1.-.., .x) and ā (ō), Fa) denote the vectors of standard deviations of stochastic gradient per layer and per dimension respectively. Finally, we assume that the gradients are bounded i.e., |[V[(x, s)ly G for all [d], Rd and S. Note that such assumptions typical in the analysis of stochastic first-order methods (cf. (Ghadimi & Lan, 2013a: 3a; Ghadimi et al.,
2014; Reddi et al., 2016; 2018)).
Stochastic gradient descent (SGD) one of the simplest first-order algorithms for solving problem in Equation 1. The update at the th iteration of SGD is of the following form:
1
F++1 F2 Ne
Vi(2s, st) + Aree
(SGD)
S
where S, is of b random samples drawn from the distribution P. For very large batch settings, the following is a well-known result for SGD.
Theorem ((Ghadimi & Lan, 2013b)). With large batch T and using appropriate learning rate, we have the following for the iterates of SGD:
E (IV < (((()) (((x))-((z')(aa T(x')))
||0|/3 T
T
where xt is an optimal solution the problem equation and xa an iterate uniformly randomly chosen from {01, ... IT)
However, tuning the learning rate π) SGD, especially in large batch settings. is difficult in practice. Furthermore, the dependence on Loo (the maximum of smoothness across dimension) can lead to significantly slow convergence. In the next section, we discuss algorithms circumvent this issue.
3

---

Published as a conference paper at ICLR 2020
3 ALGORITHMS
In this section, we first discuss general strategy to adapt the learning rate large batch settings. Using this strategy, we discuss two specific algorithms in the later part of the section. Since our primary focus is on deep learning, our discussion is centered around training a h-layer neural network.
General Strategy. Suppose we use an iterative base algorithm A (e.g. SGD or ADAM) the small batch setting with the following layerwise update rule:
F++1 = Tt + 7)
where u, the update made by A time step We propose the following two changes the update for large batch settings:
1. The update is normalized to unit (y-norm. This is ensured by modifying the update to the form ue/ /e/l|u,]l Throughout this paper, such a normalization done layerwise i.e., the update for each layer is ensured to be unit 1o-mmm..
2. The learning rate is scaled by o( [[=e]] for some function R normalization, such a scaling is done layerwise.
R
Similar to the
Suppose the base algorithm A SGD, then the modification results in the following update rule:
2+ zt
9t
(3)
11
for all layers E [h] and where o, and Di are the parameters and the gradients of the 7th layer at time step f. The normalization modification is similar to one typically used in normalized gradient descent except that is done layerwise. Note that the modification leads to biased gradient update; however, in large-batch settings, it can be shown that this bias is small. It is intuitive that such normalization provides robustness exploding gradients (where the gradient can be arbitrarily large) and plateaus (where the gradient can be arbitrarily small). Normalization of this form essentially ignores the size of the gradient and is particularly useful in large batch settings where the direction of the gradient is largely preserved.
The scaling term involving & ensures that the norm of the update is of the same order as that of the parameter. We found that this typically ensures faster convergence in deep neural networks. In practice, we observed that a simple function of ó(z) min qain(man mmax(2. (z, n 7 works well. It is instructive to consider the where (=z)
In this scenario, the overall change in the learning
rate is [|= gradient (see equation 2). We now discuss different instantiations of the strategy discussed above. In particular, we focus on two algorithms: LARS (3.1) and the proposed method, LAMB (3.2).
"1 which can also be interpreted as an estimate on the inverse of Lipschitz constant the
3.1 LARS ALGORITHM
The first instantiation of the general strategy is LARS algorithm (You al., 2017), which is obtained by using momentum optimizer as the base algorithm A in the framework. LARS was earlier proposed for large batch learning for RESNET ImageNet. In general, it is observed that the using (heavy-ball) momentum, one can reduce the variance the stochastic gradients at the cost of little bias. The pseudocode for LARS is provide in Algorithm 1.
We now provide convergence analysis for LARS in general nonconvex setting stated in this paper. For the sake of simplicity, we analyze the case where 3 = 0 and À = 0 in Algorithm 1. However, our analysis should extend to the general case well. We will defer all discussions about the convergence rate to the end of the section.
Theorem 2. Let 7)- = =
20( ag L for all t E [T], b= T, a S Φ(v) S On for all u > O where an, ag 0. Then for x, generated using LARS (Algorithm 1). we have the following bound
2 0o(
(f(x)) - -(xx)) Laan
||0
E
VV
T
Th
where x" is an optimal solution to the problem in equation and n. is an iterate uniformly randomly chosen from (xr,'.. ,IT].
4

---

Published as a conference paper at ICLR 2020
Algorithm 2 LAMB
Input: 21 E R,, learning rate (1: Y2 parameters
0 < 81, 32 < 1, scaling function di e
Set mo = 0, vo =
Algorithm LARS
for f = 1 to 7 do
Draw samples S. from P.
Input: E RR" learning rate (n 1 parameter
Compute 9:
[3.] E Vi(xi,
0 < 31 < 1, scaling function
my 81 3,m (1 D1 gt
Set mo
for to 7 do
V = 320 + (1- 32)92
Draw samples S, from
my m,/(1 (1)
Compute 9:
[8:] 2.,es: ES, Vi( 8t)
8t = 2//(1 - -35)
my
= 81m
81 (g+ Ame)
Compute ratio ri
vo
m
#41 = 24 q
edIz) 1) mi for all (b]
1 = x. 7 lr. +xx (c)
o(I (1)
(r( + +xxe
Im?
end for
end for
3.2 LAMB ALGORITHM
The second instantiation of the general strategy is obtained by using ADAM as the base algorithm A. ADAM optimizer is popular in deep learning community and has shown to have good performance for training state-of-the-an language models like BERT. Unlike LARS, the adaptivity of LAMB is two-fold: (i) per dimension normalization with respect to the square root of the second moment used in ADAM and (ii) layerwise normalization obtained due to layerwise adaptivity. The pseudocode for LAMB is provided in Algorithm When 31 and B2 O, the algorithm reduces to be Sign SGD where the learning rate is scaled by square root of the layer dimension (Bernstein et al., 2018).
The following result provides convergence rate for LAMB in general nonconvex settings. Similar to the previous case, we focus on the setting where 31 = and À 0. As before, our analysis extends to the general case; however, the calculations become messy.
Theorem 3. Let Q) = 1
ag LIITT for all t E [T]. b T, d. d/ all E [h], and on S. (vv) S aaa for all u > O where on, On > 0. Then for xa generated using LAMB (Algorithm 2), we have the following bounds:
2(f(n)) ) f(="))
When 32 = 0, we have
IV so
(f((z)) -f(f(())))) f(a"))Lavu
|1ā1 Th
E
T
2.
When 32 O, we have
G2d x 12(f(ro)-
-fxxx
|لll|
[[[V](=E)11 so
(v(1--B) h (
T
where x² is an optimal solution to the problem in equation I and To is an iterate uniformly randomly chosen from ..,...",T]
Discussion on convergence rates, We first start our discussion with the comparison convergence rate of LARS with that of SGD (Theorem 1). The convergence rates of LARS and SGD differ in two ways: (1) the convergence criterion (E)>> h [|]+ opposed E|]| fll²] in SGD and (2) the dependence on L and a in the convergence rate. Briefly, the convergence rate of LARS is better than SGD when the gradient is denser than curvature and stochasticity. This convergence rate comparison is similar in spirit to the one obtained in (Bernstein et al., 2018). Assuming that the convergence criterion in Theorem and Theorem 2 is of similar order (which happens when gradients are fairly dense), convergence rate of LARS and LAMB MB depend on Lavg instead of Lisc and are thus, significantly better than that of SGD. A more quantitative comparison is provided in Section C of the Appendix. The comparison of LAMB (with 3g with SGD is along similar lines. We obtain slightly worse rates for the case where 3g > 0: although, we believe that its behavior should be better than the case B2 0. We leave this investigation to future work.
5

---

Published as a conference paper at ICLR 2020
4 EXPERIMENTS
We now present empirical results comparing LAMB with existing optimizers on two important large batch training tasks: BERT and RESNET-50 training. We also compare LAMB with existing optimizers for small batch size (< 1K) and small dataset (c.g. CIFAR, MNIST) (see Appendix).
Experimental Setup. To demonstrate its robustness, we use very minimal hyperparameter tuning for the LAMB optimizer. Thus, it is possible to achieve better results by further tuning the hyperparameters. The parameters 31 and 32 in Algorithm 2 are set to 0.9 and 0.999 respectively in all our experiments; we only tune the learning rate. We use a polynomially decaying learning rate of 1): = 1/o (1 (1-) in Algorithm 2), which is the same as in BERT baseline. This setting also works for all other applications in this paper. Furthermore, for BERT and RESNET-50 training, we did not tune the hyperparameters of LAMB while increasing the batch size. We use the square root of LR scaling rule to automatically adjust learning rate and linear-epoch warmup scheduling. We use TPUv3 in all the experiments. A TPUv3 Pod has 1024 chips and can provide more than 100 petaflops performance for mixed precision computing. To make sure we are comparing with solid baselines, we use grid search to tune the hyper-parameters for ADAM, ADAGRAD, ADAMW (ADAM with weight decay), and LARS. We also tune weight decay for ADAMW. All the hyperparameter tuning settings are reported in the Appendix. Due to space constraints, several experimental details are relegated to the Appendix.
4.1 BERT TRAINING
We first discuss empirical results for speeding up BERT training. For this experiment, we use the same dataset Devlin al. (2018), which is concatenation of Wikipedia and BooksCorpus with 2.5B and 800M words respectively. We specifically focus on the SQuAD task? in this paper. The Fl score on SQuAD-vI 1 is used as the accuracy metric in our experiments. All our comparisons are with respect to the baseline BERT model by Devlin et al. (2018). To train BERT, Devlin et al. (2018) first train the model for 900k iterations using sequence length of 128 and then switch to a sequence length of 512 for the last 100k iterations. This results in training time of around days on 16 TPUv3 chips. The baseline BERT model achieves a Fl score of 90,395, To ensure fair comparison, we follow the same SQuAD fine-tune procedure of Devlin et al. (2018) without modifying any configuration (including number of epochs and hyperparameterr). As noted earlier, we could get even better results by changing the fine-ture tune configuration. For instance, by just slightly changing the learning rate in the fine-tune stage, we can obtain a higher Fl score of 91.688 for the batch size of 16K using LAMB. We report Fl score of 91.345 in Table 1, which is the score obtained for the untuned version. Below we describe two different training choices for training BERT and discuss the corresponding speedups.
For the first choice, we maintain the same training procedure as the baseline except for changing the training optimizer to LA MB. We run with the same number of epochs as the baseline but with batch size scaled from 512 to 32K. The choice of 32K batch size (with sequence length 512) is mainly due to memory limits of TPU Pod. Our results are shown in Table 1. By using the LAMB optimizer, we are able to achieve a Fl score of 91. 11.460 in 15625 iterations for a batch size of 32768 (14063 iterations for sequence length 128 and 1562 iterations for sequence length 512). With 32K batch size, we reduce BERT training time from 3 days to around 100 minutes. We achieved 49.1 times speedup by 64 times computational resources (76.7% efficiency). We consider the speedup is great because we use the synchronous data-parallelism. There is communication overhead coming from transferring of the gradients over the interconnect. For RESNET-50, researchers are able to achieve 90% scaling efficiency because RESNET-SK has much fewer parameters (# parameters is equal to #gradients) than BERT (25 million versus 300 million).
To obtain further improvements, we use the Mixed-Batch Training procedure with LAMB. Recall that BERT training involves two stages: the first 9/10 of the total epochs use sequence length of 128, while the last 1/10 of the total epochs use a sequence length of 512. For the second stage training, which involves a longer sequence length, due to memory limits, maximum batch size of only 32768 can be used on a TPUv3 Pod. However, we can potentially use a larger batch size for the first stage because of shorter sequence length. In particular, the batch size can be increased to 131072 for the first stage. However, we did not observe any speedup by increasing the batch size from 65536 to 131072 for the first stage, thus, we restrict the batch size to 65536 for this stage. By using this strategy, we are able to make full utilization of the hardware resources throughout the training
https://raipuurarrggin httpuuukrarggtub.io.. wwatbbio/SOAA--xxxhphe
Pre-trriined BERT model can be downloaded from http::/ggiiaaaccom///og//reee//e

---

Published as a conference paper at ICLR 2020
Table 1:
We use the F1 score on SQuAD-vl the aceuracy metric. 'he baseline F1 score is the score obtained by the pre-trained model (BERT-Large) provided on BERT's 's public repository (as of February 1st, 2019). We use TPUv3s in experiments. We use the same setting the baseline: the first 9/10 of the total epochs used sequence length of 128 and the last 1/10 of the total epochs used sequence length of 512. All the experiments run the same number of epochs. Dev set means the test data. worth noting that we can achieve better results by manually tuning the hyperparameters.
The data in this table is collected from the untuned version.
Solver
batch size
steps
score on dev set
TPUs
Time
Baseline
512
1000k
90.395
16
81,4h
LAMB
512
1000k
91.752
16
82.8h 8h
LAMB
1k
500k
91.761
32
43.2h
LAMB
2k
250k
91.946
64
21,4h
LAMB
4k
125k
91.137
128
693.6m
LAMB
8k
62500
91.263
256
390.5m
LAMB
16k
31250
91.345
512
200.0m
LAMB
32k
15625
91.475
1024
101.2m
LAMB
64k/32k
8599
90.584
1024
76.19m
procedure. Increasing the batch size is able to warm-up and stabilize the optimization process (Smith et al., 2017), but decreasing the batch size brings chaos to the optimization process and can cause divergence. In our experiments, we found technique that is useful stabilize the second stage optimization. Because we switched to a different optimization problem, it is necessary to re- re-warm-up the optimization. Instead of decaying the learning rate at the second stage, we ramp up the learning rate from zero again in the second stage (re-warm-up). As with the first stage, we decay the learning rate after the re- re-warm-up phase. With this method, we only need 8599 iterations and finish BERT training in 76 minutes (100.2% efficiency).
Comparison with ADAMW and LARS. To ensure that our approach is compared to a solid baseline for the BERT training, we tried three different strategies for tuning ADAMW (Loshchiloy Hutter, 2017): (1) ADAMW with default hyperparameters (Devlin al., 2018) (2) ADAMW with the same hyperparameters as LAMB, and (3) ADAMW with tuned hyperparameters. ADAMW stops scaling at the batch size of 16K because it is not able to achieve the target F1 score (88.1 vs 90.4). The tuning information of ADAMW is shown in the Appendix. For 64K/32K mixed-batch training. fail to get any reasonable result with ADAMW
even after extensive tuning of the hyperparameters, optimizer. We conclude that ADAMW does not work well large-batch BERT training or least hard to tune. We also observe that LAMB performs better than LARS for all batch sizes (Table 2).
Table 2: LAMB achieves higher performance (F1 score) than LARS for all the batch sizes. The baseline achieves a Fl score of 90.390. Thus, LARS stops scaling at the batch size of 16K.
Batch Size
512
1K
2K
4K
8K
16K
32K
LaRS
90.717
90.369
90.748
90.537
90.548
89.589
diverge
LAMB
91.752
91.761
91.946
91.137
91.263
91.345
91.475
4.2 IMAGENET TRAINING WITH RESNET-50.
ImageNet training with ResNer-50 is an industry standard metric that is being used in ! The baseline get 76.3% top-1 accuracy in 90 epochs (Goyal et al., 2017). All the successful implementations are based on momentum SGD (He et al., 2016; Goyal et al., 2017) or LARS optimizer (Ying et al., 2018 8; Jia et al., 2018; Mikami et al., 2018 8; You et al., 2018; Yamazaki et al., 2019). Before our study, we did not find any paper reporting a of accuracy achieved
hhups///lpprff.rg/
7

---

Published as a conference paper at ICLR 2020
by ADAM (Kingma & Ba, 2014), ADAGRAD, ADAMW optimizer. In experiments, even with comprehensive hyper-parameter tuning, ADAGRAD/ADAM/ADAAN (with batch size 16K) only achieves 55.38%//6.004%//6 .7% top-1 accuracy. After adding learning rate scheme of Goyal et al. (2017), the top-1 accuracy of ADAGRAD/ADAM/ADAA However, they are still much lower than 76.3%. The details of the tuning information are in the Appendix. Table 3 shows that LAMB achieve the target accuracy. Beyond batch size of 8K, LAMB'S accuracy is higher than the momentum. LAM s accuracy is also slightly better than LARS. At batch size of 32K, LAMB achieves 76.4% top-1 accuracy while LARS achieves 76.3%. At batch size of 2K, LAMB is able to achieve 77.11 % top-1 accuracy while LARS achieves 76.6%.
was improved to 72.0% 317733448//..... 3.7%.
Table 3: Top-1 validation accuracy of ImageNer/RESNET--b training at the batch size 16K (90 epochs). The performance of momentum was reported by (Goyal et al., 2017). + mcans adding the learning rate scheme Goyal et al. (2017) to the optimizer: (1) 5-epoch warmup to stablize the initial stage; and (2) multiply the learning rate by 0.1 at 30th, 60th, and 80th epoch. The target accuracy is around 0. 0.763 (Goyal et al., 2017). All the adaptive solvers were comprehensively tuned. The tuning information was in the Appendix.
optimizer adagrad/adagrad+
adam/adam+
adamw/adamw+
momentum
lamb
Accuracy
0.5538/0.7201
0.6604/0.7348
0.6727/0.7307
0.7520
0.7666
4.3
HYPERPARAMETERS FOR SCALING THE BATCH SIZE
For BERT and ImageNet training, we did not tune the hyperparameters of LAMB optimizer when increasing the batch size. We use the square root LR scaling rule and linear-epoch warmup scheduling to automatically adjust learning rate. The details can be found in Tables 4 and 5
Table 4: Untuned LAMB for BERT training across different batch sizes (fixed #epochs). We use square root LR scaling and linear-epoch warmup. For example, batch size 32K needs finish 15625 iterations. It uses 0.2x 15625 = 3125 iterations for learning rate warmup. BERT'S baseline achieved a higher F1 score if manually tune the hyperparameters.
F1 score of 90.395. We achieve
Batch Size
512
1K
2K
4K
8K
16K
32K
Learning Rate
5
200x10
27 10
2=0x10
01.0x100
20
20 10
Warmup Ratio
320
160
40
20
10
Fl score
91.752
91.761
91.946
91.137
91.263
91.345
91.475
Exact Match
85,090
85,260
85.355
84.172
84.901
84.816
84.939
Table 5: Untuned LAMB for ImageNet training with RESNET-50 for different batch sizes (90 epochs). We use square root LR scaling and linear-epoch warmup. The baseline Goyal et al. (2017) gets 76.3% top-1 accuracy 90 epochs. Stanford DAWN Bench (Coleman al., 2017) baseline achieves 93% top-5 accuracy. LAMB achieves both of them. LAMB can achieve an even higher accuracy if we manually tune the hyperparameters.
Batch Size
512
IK
2K
4K
8K
16K
32K
Learning Rate
28.0 100
255x 106
22
21
100
21dx
2000 100
2000 x 100
Warmup Epochs
0.3125
0.625
1.25
2.5
10
20
Top-5 Accuracy
0.9335
0.9349
0.9353
0.9332
0.9331
0.9322
0.9308
Top-1 Accuracy
0.7696
0.7706
0.7711
0.7692
0.7689
0.7666
0.7642
5
CONCLUSION
I Large batch techniques are critical to speeding up deep neural network training. In this paper, we propose the LAMB optimizer, which supports adaptive elementwise updating and layerwise learning
8

---

Published as a conference paper at ICLR 2020
rates. Furthermore, LAMB is general purpose optimizer that works for both small and large batches. We also provided theoretical analysis for the LAMB optimizer, highlighting the cases where performs better than standard SGD. LAMB achieves better performance than existing optimizers for a wide range of applications. By using LAMB, we are able to scale the batch size of BERT pre-training to 64K without losing accuracy, thereby. reducing the BERT training time from 3 days to around 76 minutes. LAMB also the first large batch adaptive solver that can achieve state-of-the-ar accuracy on ImageNet training with RESNET-50.
6 ACKNOWLEDGEMENT
We want to thank the comments from George Dahl and Jeff Dean. We want to thank Michael Banfield, Dehao Chen, Youlong Cheng, Sameer Kumar, and Zak Stone for TPU Pod support.
REFERENCES
Takuya Akiba, Shuji Suzuki, and Keisuke Fukuda. Extremely large minibatch sgd: Training resnet-50 on imagenet 15 minutes. arXiv preprint arXiv:/711 1.0425,, 2017.
Yoshua Bengio. Practical recommendations for gradient- based training of deep architectures. In Neural networks: Tricks of the trade, pp. 437-478. Springer, 2012.
Jeremy Bernstein, Yu-Xiang Wang, Kamyar Azizzadeneshelli, and Anima Anandkumar. signsgd: compressed optimisation for non-convex problems. CoRR, abs/1802.04434 2018.
Valeriu Codreanu, Damian Podareanu, and Vikram Saletore. Scale out for large minibatch sgd: Residual network training on imagenet- 1k with improved accuracy and reduced time to train. arXiv preprint arXiv: /7/ 1.04291, 2017.
Cody Coleman, Deepak Narayanan, Daniel Kang, Tian Zhao, Jian Zhang, Luigi Nardi, Peter Bailis, Kunle Olukotun, Chris Ré, and Matei Zaharia. Dawnbench: An end-to-end deep learning bench- mark and competition. Training, 100(101):102, 2017.
Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthieu Devin, Mark Mao, Andrew Senior, Paul Tucker, Ke Yang, Quoc V Le, et al. Large scale distributed deep networks. In Advances in neural information processing systems, pp. 1223-1231, 2012.
Aditya Devarakonda, Maxim Naumov, and Michael Garland. Adabatch: Adaptive batch sizes for training deep neural networks. arXiv preprint arXiv: 2017.
Jacob Devlin, Ming-Was Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 11X0:18100004805, 2018.
Timothy Dozat. Incorporating nesterov momentum into adam. 2016.
Saeed Ghadimi and Guanghui Lan. Stochastic first- and zeroth-order methods for nonconvex stochastic programming. SIAM Journal on Optimization, 23(4):2341-2) -22628,,,,,,,,, 2013a. doi: 10.1137/ 12088081 1..
Saeed Ghadimi and Guanghui Lan. Stochastic first-and zeroth-order methods for nonconvex stochastic programming. SIAM Journal on Optimization, 23(4):23412236 2013b.
Saced Ghadimi, Guanghui Lan, and Hongchao Zhang. Mini-batch stochastic approximation methods for nonconvex stochastic composite optimization. Mathematical Programming, 155(1-2):267-355 2014.
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 I hour. arXiv preprint arXiv: (X06:/7706.0677 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: closing the generalization gap in large batch training of neural networks. arXiv preprint 1rXiv:1705.08741, 2017.

---

Published as a conference paper at ICLR 2020
Forrest N landola, Matthew W Moskewicz, Khalid Ashraf, and Kurt Keutzer. Firecaffe: -inerr acceleration of deep neural network training on compute clusters. In Proceedings of the IEEE Conference on Computer Vision and Patlern Recognition, pp. 2592-2600, 2016.
Xianyan Jia, Shutao Song, Wei He, Yangzihao Wang, Haidong Rong, Feihu Zhou, Liqiang Xie, Zhenyu Guo, Yuanzhou Yang, Liwei Yu, al. Highly scalable deep learning training system with mixed-precision: Training imagenet in four minutes. arXiv preprint 1rXivv1807.11200 2018.
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. arXiv preprint arXiv: 160904460, 2016.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6980 2014.
Alex Krizhevsky, One weird trick for parallelizing convolutional neural networks. arXiv preprint ar41v:1404.5997, 2014.
Mu Li. Scaling Distributed Machine Learning with System and Algorithmm Co-design. PhD thesis, Intel, 2017.
Ilya Loshchilov and Frank Hutter. Fixing weight decay regularization in adam. arXiv preprint arXiv:1711.0510,, 2017.
James Martens and Roger Grosse. Optimizing neural networks with frotecker-factored approximate curvature. In International conference on machine learning, pp. 2408-2417, 2015.
Hiroaki Mikami, Hisahiro Suganuma, Yoshiki Tanaka, Yuichi Kageyama, et al. Imagenet/resnet-50 training in 224 seconds. arXiv preprint arXiv:/811.05235 2018.
Yurii E Nesterov. A method for solving the convex programming problem with convergence rate (1/k* 2). In Dokl. akad. nauk Sssr, volume 269, pp. 543-547, 1983.
Kazuki Osawa, Yohei Tsuji, Yuichiro Ueno, Akira Naruse, Rio Yokota, and Satoshi Matsuoka. Second-order optimization method for large mini-batch: Training resnet-50 on imagenet in 35 epochs. arXiv preprint arXiv: 120019, 2018.
Benjamin Recht, Christopher Re, Stephen Wright, and Feng Niu. Hogwild: A lock-free approach to parallelizing stochastic gradient descent. In Advances in neural information processing systems, pp. 693-701, 2011.
Sashank J. Reddi, Ahmed Hefny, Suvrit Sra, Barnahás Póczos, and Alexander J. Smola. Stochastic variance reduction for nonconvex optimization. In Proceedings of the 33nd International Confer- ence on Machine Learning, ICML 2016, New York City, NY. USA, June 19-24 2016, pp. 314-323, 2016.
Sashank J. Reddi, Satyen Kale, and Sanjiv Kumar. On the Convergence of Adam & Beyond. In Proceedings of the 6th International Conference Learning Representationss, 2018.
Christopher J Shallue, Jaehoon Lee, Joe Antognini, Jascha Sohl-Dickstein, Roy Frostig, and George E Dahl. Measuring the effects of data parallelism on neural network training. arXiv preprint arXiv: wtt.18111.000000 2018.
Samuel L Smith, Pieter-Jan Kindermans, and Quoc V Le. Don't decay the learning rate, increase the batch size. arXiv preprint
2017.
Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initialization and momentum in deep learning. In International conference on machine learning, pp. 1139-1147, 2013.
Masafumi Yamazaki, Akihiko Kasagi, Akihiro Tabuchi, Takumi Honda, Masahiro Miwa, Naoto Fukumoto, Tsuguchika Tabaru, Atsushi Ike, and Kohta Nakashima. Yet another accelerated sgd: Resnet-50 training on imagenet in 74.7 seconds. arXiv preprint arXiv: 12650,, 2019.
Chris Ying, Sameer Kumar, Dehao Chen, Tao Wang, and Youlong Cheng. Image classification at supercomputer scale. arXiv preprint 181 1,06992. 2018.
10

---

Published as a conference paper at ICLR 2020
Yang You, Igor Gitman, and Boris Ginsburg, Scaling sgd batch size to 32k for imagenet training. arXiv preprint arXiv: 100:::108000000 0.0888888, 2017.
Yang You, Zhao Zhang, Cho-Jui Hsich, James Demmel, and Kurt Keutzer. Imagenet training in minutes. In Proceedings of the 47th International Conference on Parallel Processing, pp. 1. ACM, 2018.
Yang You, Jonathan Hseu, Chris Ying, James Demmel, Kurt Keutzer, and Cho-Jui Hsich. Large-batch training for Istm and beyond. arXiv preprint 1rXiv:1901.08256, 2019.
Manzil Zaheer, Sashank J. Reddi, Devendra Singh Sachan, Satyen Kale, and Sanjiv Kumar. Adaptive methods for nonconvex optimization. In Advances Neural Information Processing Systems, pp. 9815-9825, 2018.
Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank J. Reddi, Sanjiv Kumar, and Suvrit Sra. Why ADAM beats SGD for attention models. CoRR, abs/1912.03394, 2019.

---

Published as a conference paper at ICLR 2020
APPENDIX
A PROOF OF THEOREM
Proof. We analyze the convergence of LARS for general minibatch size here. Recall that the update of LARS is the following
Iet1 xt -nollla
(6)
gr
l1ge
(8)
for all i E [h]. Since the function f is L-smooth, we have the following:
h
f(0(+1) S f(x,) + (Z))( 2 - -x( ((0)
(6)
11
+
Fi+1
-
x
2
1=1
(3)
Lon (11x;)))
ge
f((x) -
+11 x [if((z)]
x
+
11ge
2
im
Gml
h
d
(1| ( x [V. x
9r.
V [
[Vf(
[+((a)]] )
n
S s(x))) - 1/x
L1|
119.5
V.
2
((x)) - 1/t
(1|
I1) x [[v,[(
b
d:
1 x
[. x
Pi.
(
V.f(x))
+
na
1111
11g
Vis((a))
2
(4)
The first inequality follows from the lipschitz continuous nature of the gradient. Let A (x)( Then the above inequality can be rewritten in the following manner:
(4)
gt
h
f(x41) S f1x)) - mt
M1 (f)))
x (4 + (z)()))) || + (x.))))
v.f(=:)] ))
s (1 (2)
+
|T
2
i=1 j=1
f(x)) -
(A + V.f(m.) V.f(za) V.f(())
nia
m
0(1=,:"011
1 ( v.f(x))
+ 2
||
f(x) 17
2110 ID) v.f(all
[V.f(a))] (0) + v.f(x))))- (A, v.f(x,) vef(x.))
nha
+
<11 x
+
||
A + V.f(x.))
2
niaa
f(x)) 17
<11 (9) [)) [x)()))|
||||
2
[..f(=))|| 0 v.f(x))]
94 + V.f(m.))
<11= (2) x
+ 17
| (
(zf(x)))
(5)
12

---

Published as a conference paper at ICLR 2020
Using Cauchy-Schward inequality in the above inequality, we have:
(14411 S (11)) 1/6 Σ
(11 +1 [II+/
+ 1/t Σ
((1x(111 [I) x (|V.)(=)) v.f(2))) +|44(11 +
1
2
h
na
S fz1)
7
5(1ا) (())) + 2T Σ
M1 1) x
1 60
||
||
|1
2
Taking expectation, we obtain the following:
02
B[f(2++)] S f((r) - T1 Σ
M (I
+ 2mh Σ
<1|=; ID) x E
[1|
|
2
i=1
1=1
||0||1
22
fx1) 17111
[+ + 2m
+ 2 |T|||
vb
Summing the above incquality for t to T and using telescoping sum, we have the following inequality:
T h
( + 2n/T aullali ~
([(())))) -1ya 170y
E
vb +
2
Rearranging the terms of the above incquality, and dividing by nTa,, we have:
T
El [[[))(())]] ((x)) - -[[((T++)]
20..1ol11
+ no
T
Tnau
vbai
204
[(x)) - -(xx)
20
+ na [|[|11.
+
Tna
avb
2an
B PROOF OF THEOREM 3
Proof. We analyze the convergence of LAMB for general minibatch size here. Recall that the update of LAMB is the following
2. - (1
Tr
141
|
(
for all i E [h]. For simplicity of notation, we reason the
Since the function o L-smooth, we have the following:
[(x+1) S fxx) + (Va((x)), (1) - (0)) 11:21 - x 10
h
h:
=f(x)) -
(3) 11) x
[x)(())) x
[
+ >
(7)
1|
(8)
==1
n
13

---

Published as a conference paper at ICLR 2020
The above inequality simply follows from the lipschitz continuous nature of the gradient. We bound term Ti the following manner:
h
d.
x (1a.fca))) x
T
S 7
1|r,
de
32
S --1
1/t
G2
(0(1=1²1)
[.o( x si.
(6)
0(11>("1 ||) [.,(())]
T
((sign(V:[(rx)) + sign sign(r(.) (r(2))
1/t
|r:
This follows from the fact that || V 1-32 and voi S G. If B2 = 0, then Ti can be bounded as follows:
Ti T
(0(1=
t=1
==1
de
o(1=2"11 ID) x (ff(( x
1
# sign(r(.))
(0)
The rest of the proof for B2 0 is similar to argument for the case B2 0, which is shown below. Taking expectation, we have the following:
h
d,
E f
7
32
(1 x (V..(a))); x x
G2d
de
1/
E
[V.](x:)]s x
|| 1n(01//
(ssig((V + (sign(())))
B2
[b(11= (I) (.f())))) o(.)]
1/6
E
G2d,
h
1
E au
sign(g())))
1=1
h d
32
(0 x (V. x 9 E (((l=(10))
S
-
G2d
j=1
h d,
+ T
an a. [ P(sign (V,ggn((..fll (z;))) & sign(g(
1=
Using the bound on the probability that the signs differ, we get:
7a(1 - 32)
B(Ti] < -na V
Vf(x)) + m.
G2d
==
Substituting the above bound Ti in equation 7, we have the following bound:
h(1 -
32) [
||āl|
B[f(1++1) < f(x))
+ maa
(9)
G2d
vb
2
Summing the above inequality for = to T and using telescoping sum, we have the following inequality:
h(1 - 82)
|+ll1
E(f(er+1) < f(x1) man
Σ
B[IVf((>))11 + aT
|[]|]
Gdd
vb
14

---

Published as a conference paper at ICLR 2020
Algorithm 4 NNN-AAM
Algorithm 3 N-LAMB
Input: 21 E RR", learning rate (1: 12 parameters
Input: E R learning rate {m.) parame-
0 < B1, 32 < 1, scaling function @. > 0, parame-
ters 0 91, 82 < 1, scaling function 0,
ters 0 [ )
parameters < { /t1 <1
Set mg 0, Uo
Set mm = 0, Up =
for 1 to T
for t = to 1' do
Draw samples S. from P.
Draw samples St from P. Compute 0:
[8.1 E,,s, 8t).
Compute gr
[i L..s, Vi(zi, s:).
31m--11 + $1)ge
m:
8.m m 1 + -1 81 gtt
n
2+ ml
+ (1 (1-if)or
1-m oj
+ (1 (m
T- P
1- Pi
Vt 32 (1 32)92
Vt = 82 + (1 -2)9
0
T-m1
1-11 1-
Compute ratio rt V
Compute ratio rt
(1) "++1
= 1e. +x=
(rz) + Aaa)
(1) t-1
ellr (I)
(re + Xax)
1. + x 1
end for
end for
Rearranging the terms of the above inequality, and dividing by T7,,, we have:
h(1 82) T
Gdd T ((] f(x1) - ([(((T+))]
aallōl12
avvb + " L|12
Tnan
2
f(z1) - f(x')
0u
+ no LII+
Tnjn
avbb
201
C COMPARISON OF CONVERGENCE RATES OF LARS AND SGD
Inspired by the comparison used by (Bernstein et al., 2018) for comparing SIGN SGD with SGD, we define the following quantities:
V fxx))l
IV..((0>1)
h
h
||4||7
1
3b dllo
||0||3
Then LARS convergence rate can written in the following manner:
([llV((>))))* 0o (f((r)) - f(xx))
I. vh u2
| v2
+
T
T
v3
If vh. < viz and ti < 2i then LARS (i.e., gradient more denser than curvature stochasticity). we gain over SGD. Otherwise, SGD'S upper bound on convergence rate is better.
D N-LAMB: NESTEROV MOMENTUM FOR LAMB
Sutskever al. (2013) report that Nesterov's accelerated gradient (NAG) proposed by Nesterov (1983) is conceptually and empirically better than the regular momentum method for convex, non-stochastic objectives. Dozat (2016) incorporated Nesterov's momentum into Adam optimizer and proposed the Nadam optimizer. Specifically, only the first moment of Adam was modified and the second moment of Adam was unchanged. The results several applications (Word2Ved Image Recognition,
15

---

Published as a conference paper at ICLR 2020
ImageNet/ResNett55 (Batch Size=32K, 90 epochs)
77
Top-1 Validation Accuracy
76
75
74
73
72
71
Momentum
Nadam
LAMB
N-LAMB
NN-LAMB
Figure 1: This figure shows N-LAMB and NN-LAMB can achieve comparable accuracy compared to LAMB optimizer. Their performances are much better than momentum solver. The result of momentum optimizer
reported by Goyal et al. (2017). For Nadam, we the learning rate recipe of (Goyal et al., 2017): (1) 5-epoch warmup to stablize the initial stage; and (2) multiply the learning rate by 0.1 at 30th, 60th, and 80th epoch. The target accuracy is around 0. 7.66 (Goyal et al., 2017). also tuned the learning rate Nadam (le-4, 2e-4, 9e-4, le-3, 2e-3, 9e-3, le-2).
and LSTM Language Model) showed that Nadam optimizer improves the speed of convergence and the quality of the learned models. We also tried using Nesterov's momentum replace the regular momentum of LAMB optimizer's 's first moment. In this way, we got new algorithm named N-LAMB (Nesterov LAMB). The complete algorithm is in Algorithm 3. We also Nesterov's momentum to replace the regular momentum of LAMB optimizer's second moment. We refer to this algorithm as NN-LAMB (Nesterov's momentum for both the first moment and the second moment). The details of NN-LAMB were shown in Algorithm
Dozat (2016) suggested the best performance of Nadam was achieved by 31 = 0.975, 32 0.999, and le-8. We used the same settings for N-LAMB and NN-LAMB. We scaled the batch size to 32K for ImageNet training with ResNet-50. Our experimental results show that N-LAMB and NN-LAMB achieve comparable accuracy compared to LAMB optimizer. Their performances much better than momentum solver (Figure 1).
E LAMB WITH LEARNING RATE CORRECTION
There are two operations at each iteration in original Adam optimizer (let us call it adam-correction)):
= mm/(L (1 -
0/ = / B3)
has an impact on the learning rate by 7) == T1 (1- (11--5)(1120 3.. According to our experimental results, adam-correction essentially has the same effect learning rate warmup (see Figure 2). The warmup function often was implemented in the modern deep learning system. Thus, we can remove adam-correction from the LAMB optimizer. We did not observe any drop in the test or validation accuracy for BERT and ImageNet training.
LAMB WITH DIFFERENT NORMS
We need to compute the matrix/tensor norm for each layer when we do the parameter updating in the LAMB optimizer. We tried different norms in LAMB optimizer. However, we did not observe significant difference the validation accuracy of ImageNet training with ResNet-50. our experiments, the difference in validation accuracy is less than 0.1 percent (Figure 3). We use 1.2 norm the default.
16

---

Published as a conference paper at ICLR 2020
08
Learning Rate
02
adam-correct (betal=0.9, beta2=0.999) gradual learning warmup
Iterations
Figure 2: The figure shows that udam-correction has the same effect as learning rate warmup. We removed adam-correction from the LAMB optimizer. We did not observe any drop in the test or validation accuracy for BERT and ImageNet training.
ImageNet/ResNet-50 by LAMB (90 epochs, Batch Size=32K)
Validation Accuracy
0.925
L1 Norm (76.5%/93.2%)
0.900
L2 Norm (76.5%/93.1%)
0.875
Inf Norm (76.4%/93.2%)
0.850
0.825
0.800
0.775
0.750
Top-1 Accuracy
Top-5 Accuracy
Figure 3: We tried different norms in LAMB optimizer. However, we did not observe a significant difference in the validation accuracy of ImageNet training with ResNet-50. We use L2 norm as the default.
G REGULAR BATCH SIZES FOR SMALL DATASETS: MNIST AND CIFAR-10.
According DAWNBench, DavidNet custom 9-layer Residual ConvNet) is the fastest model for CIFAR-10 dataset (as of April 1st, 2019)². The baseline uses the momentum SGD optimizer. Table and Figure 4 show the test accuracy of CIFAR-10 training with DavidNet. The Py PyTorch implementation (momentum SGD optimizer) on GPUs was reported on Standford DAWNBench's website, which achieves 94.06% in 24 epochs. The Tensorflow implementation (momentum SGD optimizer) on TPU achieves 93.72% accuracy in 24 epochs We use the implementation of TensorFlow on TPUs. LAMB optimizer is able to achieve 94.08% % test accuracy in 24 epochs, which better than other adaptive optimizers and momentum SGD. Even the smaller tasks like MNIST training with LeNet, LAMB is able to achieve better accuracy than existing solvers (Table 7).
H IMPLEMENTATION DETAILS AND ADDITIONAL RESULTS
There are several we explain them to help users to have better understanding. 81 is used for decaying the running average of the gradient. 92 is used for decaying the running average of the square of gradient. The default setting for other parameters: weight decay rate A=0.01, =0,0,9 00999, c=le-6. We did not tune 31 and 3g However, our experiments show that tuning them may get a higher accuracy.
LAMB optimizer. Although users do not need to tune them,
http://wwww.c.stttooooooccccchem w/trraa 10/train html http://h......coooooooooooooea/a
cifar 10_ www.dddda_ffi....mm
17

---

Published as a conference paper at ICLR 2020
CIFAR-10 with DavidNet (under 1 minute on TPU, 24 epochs)
Test Accuracy
8935
0438
0925
0.30
0305
LAMB
Momentum
Adagrad
Adam
AdamW
Figure LAMB is better than the existing solvers (batch size = 512). We make sure all the solvers are carefully tuned. The learning rate tuning space of Adam, AdamW, Adagrad and LAMB is (0.0001, 0.0002, 0.0004, 0.0006, 0.0008, 0.001, 0.002, 0.004, 0.006, 0.008, 0.01, 0.02, 0.04, 0.06, 0.08, 0.1, 0.2, 0.4, 0.6, 0.8, 1, 2, 4, 6, 8, 10, 15, 20, 25, 30, 35, 40, 45, 50). The momentum optimizer was tuned by the baseline implementer. The weight decay tenm of AdamW was tuned by (0.0001, 0.001, 0.01, 0.1, 1.0).
Table CIFAR-I0 training with DavidNet (batch size 512). All of them run 24 epochs and finish the training under one minute on one cloud TPU. We make sure all the solvers are carefully tuned. The learning rate tuning space of Adam, AdamW, Adagrad and LAMB is (0.0001, 0.0002, 0.0004, 0.0006, 0.0008, 0.001, 0.002, 0.004, 0.006, 0.008, 0.01, 0.02, 0.04, 0.06, 0.08, 0.1, 0.2, 0.4, 0.6, 0.8, 1. 2, 4, 6, 8, 10, 15, 20, 25, 30, 35, 40, 45, 50). The momentum optimizer was tuned by the baseline implementer. The weight decay term of AdamW was tuned by (0.0001, 0.001, 0.01, 0.1, 1.0).
Optimizer
ADAGRAD
ADAM
ADAMW
momentum
LAMB
Test Accuracy
0.9074
0.9225
0.9271
0.9372
0.9408
Based our experience, learning rate is the most important hyper-parameter that affects the learning efficiency and final accuracy. Bengio (2012) suggests that is often the single most important hyper-parameter and that it always should be tuned. Thus, make sure have solid baseline, carefully tune the learning rate of ADAM, ADAMW, ADAGRAD, and momentum SGD
In our experiments, we found that the validation loss is not reliable for large-batch training. A lower validation loss does not necessarily lead higher validation accuracy (Figure 5). Thus, we use the test/val accuracy or Fl score on dev set to evaluate the optimizers.
H.0.1 BERT
Table 8 shows some of the tuning information from BERT training with ADAMW optimizer. ADAMW stops scaling at the batch size of 16K. The target Fl score is 90.5. LAMB achieves a Fl score of 91.345. The table shows the tuning information of ADAMW. In Table 8, report the best F1 score we observed from our experiments.
The loss curves of BERT training by LAMB for different batch sizes are shown in Figure 6. We observe that the loss curves are almost identical each other, which means our optimizer scales well with the batch size.
The training loss curve of BERT mixed- pre-training with LAMB is shown in Figure 7. This figure shows that LAMB can make the training converge smoothly at the batch size of 64K.
Figure 8 shows that we can achieve 76.8% scaling efficiency by scaling the batch size (49.1 times speedup by 64 times computational resources) and 10.8% scaling efficiency with mixed-batch (65.2 times speedup by 64 times computational resources)
H.0.2 IMAGENET
Figures 9 14 show the LAMB trust ratio different iterations for ImageNet training with ResNet-50. From these figures we can see that these ratios are very different from each other for different layers. LAMB uses the trust ratio to help the slow learners train faster.
18

---

Published as a conference paper at ICLR 2020
Table 7:
Test Accuracy by MNIST training with LeNet (30 epochs for Batch Size 1024). The tuning space of learning rate for all the optimizers is (0.0001, 0.001, 0.01, 0.1). We use the same learning rate warmup and decay schedule for all of them.
Optimizer
Momentum
Addgrad
ADAM
ADAMW
LAMB
Average accuracy over 5 runs
0.9933
0.9928
0.9936
0.9941
0.9945
We can not trust val loss (ImageNet/ResNNet550 Batch=8K)
Validation Accuracy 76.4%
3.5
Validation Accuracy 73.9%
Validation Loss
3.0
2.5
2.0
1.5
10
20
30
40
50
60
70
90
90
Epoch
Figure 5: Our experiments show that even the validation loss is not reliable in the large-scale training. lower validation loss may lead a worse accuracy. Thus, use the test/val accuracy Fl score on dev set to evaluate the optimizers.
H.1 BASELINE TUNING DETAILS FOR IMAGENET TRAINING WITH RESNET-50
If you not interested in the baseline tuning details, please skip this section.
Goyal et al. (2017) suggested proper learning rate warmup and decay scheme may help improve the ImageNet classification accuracy. We included these techniques in Adda//AdmW/Addaaa tuning. Specifically. we use the learning rate recipe Goyal et al. (2017): (1) 5-epoch warmup to stablize the initial stage; and (2) multiply the learning rate by 0.1 at 30th, 60th, and 80th epoch. The target accuracy is around 76.3% (Goyal al., 2017). There techniques help to im- prove the accuracy of Adam/Adam W/AdaGrad to around 73%. However, even with these techniques, stil can not achieve the target validation accuracy.
To make sure our baseline is solid, we carefully tuned the hyper-paraments Table shows the tuning information of standard Adagrad. Table 10 shows the tuning information of adding the learning rate scheme of Goyal al. (2017) standard Adagrad. Table 11 shows the tuning information of standard Adam. Table shows the tuning information of adding the learning rate scheme of Goyal et al. (2017) to standard Adam. It is tricky weight decay have the effect on the performance. Thus we have four tuning sets.
the AdamW optimizer since both the L2 regularization and
The first tuing set is based on Adam W with default 1.2 regularization. We tune the learning rate and weight decay. The tuning information in Figures 13, 14, 15, and 16.
The second tuning set is based on AdamW with disabled L2 regularization. We tune the learning rate and weight decay. The tuning information is in Figures 17, 18, 19, and 20.
Then we add the learning rate scheme of Goyal et al. (2017) to AdamW and refer to it as AdamW+
The third tuning set is based Adam with default L2 regularization. We tune the learning rate and weight decay. The tuning information is Figure 21 and 22.
The fourth tuning set is based on AdamW+ W+ with disabled 1.2 regularization. We tune the learning rate and weight decay. The tuning information is in Figures 23, 24, 25.
19

---

Published as a conference paper at ICLR 2020
Table 8: ADAMW stops scaling at the batch size of 16K. The target Fl score 90.5. LAMB achieves a Fl score of 91.345. The table shows the tuning information of ADAMW. In this table, we report the best Fl score observed from experiments.
Solver
batch size
warmup steps
LR
last step infomation
Fl score on dev set
ADAMW
16K
0.05x31250
0.0001
loss=8.044711 step=28126
diverged
ADAMW
16K
0.05x31250
0.0002
loss=7.8967]) step=28126
diverged
ADAMW
16K
0.05x31250
0.0003
Ioss=8.35102, step=28126
diverged
ADAMW
16K
0.10x31250
0.0001
loss=2.01419, step=31250
86.034
ADAMW
16K
0.10x31250
0.0002
loss=1.04689, step=31250
88.540
ADAMW
16K
0.10x31250
0.0003
loss=8.05845, step=20000
diverged
ADAMW
16K
0.20x31250
0.0001
loss=1.53706, step=31250
85.231
ADAMW
16K
0.20x31250
0.0002
loss=I.15500 step=31250
88.110
ADAMW
16K
0.20x31250
0.0003
loss=1.48799 step=31250
85.653
00
LAMB Optimizrr Batch Sze 32C, fnai trein lons 1142, dew ged score 91 31.475
LAME Optember, Batch Seee 16K, firal train 1.167, dev set soore
U22 8222
Ld0s
Idda
LGMS Optimizer. Bittch
fimal loss 1.425,
91.203
XXX
LAMB Optimicor. Batch Size aK. final tra loss 590 dev sot score 91.137
1,430
91,046
N000
3000
6WV Detinizer Betm time 1KK finel train wes 222 deu sed scoce 7911
100
LAMB Optimoer, Batch Sue 512, firal train 1.115, day
Erain
6//
Berations
Figure 6: This figure shows the training loss curve of LAMB optimizer. We just want to use this figure show that LAMB can make the training converge smoothly. Even if we scale the batch size the extremely large cases, the loss curves are almost identical to each other.
Based on our comprehensive tuning results, we conclude the existing adaptive solvers do not perform well on ImageNet training or least it is hard tune them.
20

---

Published as a conference paper at ICLR 2020
BERT pre-train for 64K/32K batch size
loss 1.33871 at 7037 step
Train Loss
3
2
0
1000
2000
3000
4000
5000
6000
7000
steps
Figure 7: This figure shows the training loss curve of LAMB optimizer. This figure shows that LAMB can make the training converge smoothly at the extremely large batch size (e.g. 64K).
Scaling Efficiency
Perfect Scaling
60
Our Scaling
50
Speedup
40
20
10
0
16
32
64
128
256
512
1024 1024-mixed
TPUs
Figure 8: We achieve 76.8% scaling efficiency (49 times speedup by 64 times computational resources) and 101.8% scaling efficiency with mixed, scaled batch size (65.2 times speedup by 64 times computational resources). 1024-mixed means the mixed- batch training on 1024 TPUs.
ImageNet/ResNet-50 trainng by LAMB optimizer (1st iteration)
Trust Ratio
Layer ID
Figure 9: The LAMB trust ratio.
21

---

Published as a conference paper at ICLR 2020
ImageNet/ResNet-50 training by LAMB optimizer (4th iteration)
Trust Ratio
2.45
1200
Layer ID
Figure 10: The LAMB trust ratio.
I0ageNet/ResNett550 training by LAMB optimizer (10th iteration)
$333
8352
Trust Ratio
2.13
1112
Layer
Figure 11: The LAMB trust ratio.
ImageNet/ResNett50 training LAMB optimizer (50th iteration)
6:36
Trust Ratio
0:34
6.00
Layer ID
Figure 12: The LAMB trust ratio.
ImageNet/ResNet-50 training by LAMB optimizer (100th iteration)
6.--
6.00
Trust Ratio
620
Layer ID
Figure 13: The LAMB trust ratio.
22

---

Published as a conference paper at ICLR 2020
ImageNet/ResNet-50 training LAMB optimizer (200th iteration)
Trust Ratio
0:30
Layer ID
Figure 14: The LAMB trust ratio.
Table The accuracy information of tuning default AdaGrad optimizer for ImageNet training with ResNet-50 (batch size 16384, 90 epochs, 7038 iterations).
Learning Rate
Top-1 Validation Accuracy
0.0001
0.0026855469
0.001
0.015563965
0.002
0.022684732
0.004
0.030924479
0.008
0.04486084
0,010
0.054158527
0.020
0.0758667
0.040
0.1262614
0.080
0.24037679
0.100
0.27357993
0.200
0.458313
0.400
0.553833
0.800
0.54103595
1.000
0.5489095
2.000
0.47680664
4.000
0.5295207
6.000
0.36950684
8.000
0.3108
10.00
0.30670166
12.00
0.3091024
14.00
0.3227946
16.00
0.0063680015
18.00
0.11287435
20.00
0.21602376
30.00
0.08315023
40.00
0.0132039385
50.00
0.0009969076
23

---

Published as a conference paper at ICLR 2020
Table 10: The accuracy information of tuning AdaGrad optimizer for ImageNet training with ResNet- 50 (batch size 16384, 90 epochs. 7038 iterations). We use the learning rate recipe of (Goyal et al., 2017): (1) 5-epoch warmup to stablize the initial stage; and (2) multiply the learning rate by 0. at 30th, 60th, and 80th epoch. The target accuracy is around 0.763 (Goyal et al., 2017).
Learning Rate
Top-1 Validation Accuracy
0.0001
0.0011189779
0.001
0.00793457
0.002
0.012573242
0.004
0.019022623
0.008
0.027079264
0.010
0.029012045
0.020
0.0421346
0.040
0.06618246
0.080
0.10970052
0.100
0.13429768
0,200
0.26550293
0.400
0.41918944
0.800
0.5519816 16
1.000
0.58614093
2.000
0.67252606
4.000
0.70306396
6.000
0.709493
8.000
0.7137858
10.00
0.71797687
12.00
0.7187703
14.00
0.72007245
16.00
0.7194214
18.00
0.7149251
20.00
0.71293133
30.00
0.70458984
40.00
0.69085693
50.00
0.67976886
24

---

Published as a conference paper at ICLR 2020
Table 11: The accuracy information of tuning default Adam optimizer for ImageNet training with ResNet-50 (batch size 16384, 90 epochs, 7038 iterations). The target accuracy is around 0.763 (Goyal et al., 2017).
Learning Rate
Top-1 Validation Accuracy
0.0001
0.5521
0.0002
0.6089
0.0004
0.6432
0.0006
0.6465
0.0008
0.6479
0.001
0.6604
0.002
0.6408
0.004
0.5687
0.006
0.5165
0.008
0.4812
0.010
0.3673
Table 12: The accuracy information of tuning Adam optimizer for ImageNet training with ResNet-50 (batch size = 16384, 90 epochs, 7038 iterations). We use the learning rate recipe of (Goyal et al., 2017): (1) 5-epoch warmup stablize the initial stage; and (2) multiply the learning rate by 0. at 30th, 60th, and 80th epoch. The target accuracy around 0.763 (Goyal et al., 2017).
Learning Rate
Top-1 Validation Accuracy
0.0001
0.410319
0,0002
0.55263263
0.0004
0.6455485
0.0006
0.6774495
0.0008
0.6996867
0.001
0.71010333
0.002
0.73476154
0.004
0.73286945
0.006
0.72648114
0.008
0.72214764
0.010
0.71466064
0.012
0.7081502
0.014
0.6993001
0.016
0.69108075
0.020
0.67997235
0.040
0.58658856
0.060
0.51090497
0.080
0.45174155
0.100
0.40297446
25

---

Published as a conference paper at ICLR 2020
Table 13: The accuracy information of tuning default AdamW optimizer for ImageNet training with ResNet-50 (batch size 16384, 90 epochs, 7038 iterations). The target accuracy is around 0.763 (Goyal et al., 2017).
learning rate
weight decay
L2 regularization
Top-1 Validation Accuracy
0.0001
0.00001
default (0.01)
0.53312176
0.0002
0.00001
default (0.01)
0.5542806
0.0004
0.00001
default (0.01)
0.48769125
0.0006
0,00001
default (0.01)
0.46317545
0.0008
0.00001
default (0.01)
0.40903726
0.001
0.00001
default (0.01)
0.42401 123
0.002
0.00001
default (0.01)
0.33870444
0.004
0.00001
default (0.01)
0.12339274
0.006
0.00001
default (0.01)
0.122924805
0.008
0.00001
default (0.01)
0.08099365
0.010
0.00001
default (0.01)
0.016764322
0.012
0.00001
default (0.01)
0.032714844
0.014
0.00001
default (0.01)
0.018147787
0.016
0.00001
default (0.01)
0.0066731772
0.018
0.00001
default (0.01)
0.010294597
0.020
0.00001
default (0.01)
0.008260091
0.025
0.00001
default (0.01)
0.008870442
0.030
0.00001
default (0.01)
0.0064493814
0.040
0.00001
default (0.01)
0.0018107096
0.050
0.00001
default (0.01)
0.003540039
26

---

Published as a conference paper at ICLR 2020
Table 14: The accuracy information of tuning default AdamW optimizer for ImageNet training with ResNet-50 (batch size 16384, 90 epochs, 7038 iterations). The target accuracy is around 0.763 (Goyal et al., 2017).
learning rate
weight decay
L2 regularization
Top-1 Validation Accuracy
0.0001
0.0001
default (0.01)
0.55489093
0.0002
0.0001
default (0.01)
0.56514484
0.0004
0.0001
default (0.01)
0.4986979
0.0006
0.0001
default (0.01)
0.47595215
0.0008
0.0001
default (0.01)
0.44685873
0.001
0.0001
default (0.01)
0.41029868
0.002
0.0001
default (0.01)
0.2808024
0.004
0.0001
default (0.01)
0.08111572
0.006
0.0001
default (0.01)
0.068115234
0.008
0.0001
default (0.01)
0.057922363
0.010
0.0001
default (0.01)
0.05222575
0.012
0.0001
default (0.01)
0.017313639
0.014
0.0001
default (0.01)
0.029785156
0.016
0.0001
default (0.01)
0.016540527
0.018
0.0001
default (0.01)
0.00575765
0.020
0.0001
default (0.01)
0.0102335615
0.025
0.0001
default (0.01)
0.0060831704
0.030
0.0001
default (0.01)
0.0036417644
0.040
0.0001
default (0.01)
0.0010782877
0.050
0.0001
default (0.01)
0.0037638347
27

---

Published as a conference paper at ICLR 2020
Table 15: The accuracy information of tuning default AdamW optimizer for ImageNet training with ResNet-50 (batch size 16384, 90 epochs, 7038 iterations). The target accuracy is around 0.763 (Goyal et al., 2017).
learning rate
weight decay
L2 regularization
Top-1 Validation Accuracy
0.0001
0.001
default (0.01)
0.21142578
0.0002
0.001
default (0.01)
0.4289144
0.0004
0.001
default (0.01)
0.13537598
0.0006
0.001
default (0.01)
0.33803305
0.0008
0.001
default (0.01)
0.3261 1084
0.001
0.001
default (0.01)
0.22194417
0.002
0.001
default (0.01)
0.1833903
0.004
0.001
default (0.01)
0.08256022
0.006
0.001
default (0.01)
0.020507812
0.008
0.001
default (0.01)
0.018269857
0.010
0.001
default (0.01)
0.007507324
0.012
0.001
default (0.01)
0.020080566
0.014
0.001
default (0.01)
0.010762532
0.016
0.001
default (0.01)
0.0021362305
0.018
0.001
default (0.01)
0.007954915
0.020
0.001
default (0.01)
0.005859375
0.025
0.001
default (0.01)
0.009724935
0.030
0.001
default (0.01)
0.0019124349
0.040
0.001
default (0.01)
0.00390625
0.050
0.001
default (0.01)
0.0009969076
28

---

Published as a conference paper at ICLR 2020
Table 16: The accuracy information of tuning default AdamW optimizer for ImageNet training with ResNet-50 (batch size 16384, 90 epochs, 7038 iterations). The target accuracy is around 0.763 (Goyal et al., 2017).
learning rate
weight decay
L2 regularization
Top-1 Validation Accuracy
0.0001
0.01
default (0.01)
0.0009765625
0.0002
0.01
default (0.01)
0.0009969076
0.0004
0.01
default (0.01)
0.0010172526
0.0006
0.01
default (0.01)
0.0009358724
0.0008
0.01
default (0.01)
0.0022379558
0.001
0.01
default (0.01)
0.001566569
0.002
0.01
default (0.01)
0.009480794
0.004
0.01
default (0.01)
0.0033569336
0.006
0.01
default (0.01)
0.0029907227
0.008
0.01
default (0.01)
0,0018513998
0.010
0.01
default (0.01)
0.009134929
0.012
0.01
default (0.01)
0.0022176106
0.014
0.01
default (0.01)
0.0040690103
0.016
0.01
default (0.01)
0.0017293295
0.018
0.01
default (0.01)
0.00061035156
0.020
0.01
default (0.01)
0.0022379558
0.025
0.01
default (0.01)
0.0017089844
0.030
0.01
default (0.01)
0.0014241537
0.040
0.01
default (0.01)
0.0020345051
0.050
0.01
default (0.01)
0.0012817383
29

---

Published as a conference paper at ICLR 2020
Table 17: The accuracy information of tuning default AdamW optimizer for ImageNet training with ResNet-50 (batch size 16384, 90 epochs, 7038 iterations). The target accuracy is around 0.763 (Goyal et al., 2017).
learning rate
weight decay
L2 regularization
Top-1 Validation Accuracy
0.0001
0.00001
disable
0.48917642
0.0002
0.00001
disable
0.58152264
0.0004
0.00001
disable
0.63460284
0.0006
0,00001
disable
0.64849854
0.0008
0.00001
disable
0.6598918
0.001
0.00001
disable
0.6662801
0.002
0.00001
disable
0.67266846
0.004
0.00001
disable
0.6692708
0.006
0.00001
disable
0.6573079
0.008
0.00001
disable
0.6639404
0.010
0.00001
disable
0.65230304
0.012
0.00001
disable
0.6505534
0.014
0.00001
disable
0.64990234
0.016
0.00001
disable
0.65323895
0.018
0.00001
disable
0.67026776
0.020
0.00001
disable
0.66086835
0.025
0.00001
disable
0.65425617
0.030
0.00001
disable
0.6476237
0.040
0.00001
disable
0.55478925
0.050
0.00001
disable
0.61869305
30

---

Published as a conference paper at ICLR 2020
Table 18: The accuracy information of tuning default AdamW optimizer for ImageNet training with ResNet-50 (batch size 16384, 90 epochs, 7038 iterations). The target accuracy is around 0.763 (Goyal et al., 2017).
learning rate
weight decay
L2 regularization
Top-1 Validation Accuracy
0.0001
0.0001
disable
0.5033366
0.0002
0.0001
disable
0.5949707
0.0004
0.0001
disable
0.62561035
0.0006
0.0001
disable
0.6545207
0.0008
0.0001
disable
0.66326904
0.001
0.0001
disable
0.6677043
0.002
0.0001
disable
0.67244464
0.004
0.0001
disable
0.6702881
0.006
0.0001
disable
0.66033936
0.008
0,0001
disable
0.66426593
0.010
0.0001
disable
0.66151935
0.012
0.0001
disable
0.6545817
0.014
0.0001
disable
0.65509033
0.016
0.0001
disable
0.6529338
0.018
0.0001
disable
0.65651447
0.020
0.0001
disable
0.65334064
0.025
0.0001
disable
0.655009
0.030
0.0001
disable
0.645528 14
0.040
0.0001
disable
0.6425374
0.050
0.0001
disable
0.5988
31

---

Published as a conference paper at ICLR 2020
Table 19: The aceuracy information of tuning default AdamW optimizer for ImageNet training with ResNet-50 (batch size 16384, 90 epochs, 7038 iterations). The target accuracy is around 0.763 (Goyal et al., 2017).
learning rate
weight decay
L2 regularization
Top-1 Validation Accuracy
0.0001
0.001
disable
0.4611206
0.0002
0.001
disable
0.0076293945
0.0004
0.001
disable
0.29233804
0.0006
0.001
disable
0.57295734
0.0008
0.001
disable
0.5574748
0.001
0.001
disable
0.5988566
0.002
0.001
disable
0.586263
0.004
0.001
disable
0.62076825
0.006
0.001
disable
0.61503094
0.008
0.001
disable
0.4697876
0.010
0.001
disable
0.619751
0.012
0.001
disable
0.54243976
0.014
0.001
disable
0.5429077
0.016
0.001
disable
0.55281574
0.018
0.001
disable
0.5819295
0.020
0.001
disable
0.5938924
0.025
0.001
disable
0.541097
0.030
0.001
disable
0.45890298
0.040
0.001
disable
0.56193036
0.050
0.001
disable
0.5279134
32

---

Published as a conference paper at ICLR 2020
Table 20: The accuracy information of tuning default AdamW optimizer for ImageNet training with ResNet-50 (batch size 16384, 90 epochs, 7038 iterations). The target accuracy is around 0.763 (Goyal et al., 2017).
learning rate
weight decay
L2 regularization
Top-1 Validation Accuracy
0.0001
0.01
disable
0.0009969076
0.0002
0.01
disable
0.0008951823
0.0004
0.01
disable
0.00095621747
0.0006
0.01
disable
0.0012817383
0.0008
0.01
disable
0.016886393
0.001
0.01
disable
0.038146973
0.002
0.01
disable
0.0015258789
0.004
0.01
disable
0.0014241537
0.006
0.01
disable
0.08 1444462246
0.008
0.01
disable
0.028116861
0.010
0.01
disable
0.01 0.01 1820476
0.012
0.01
disable
0.08138021
0.014
0.01
disable
0.010111491
0.016
0.01
disable
0.0041910806
0.018
0.01
disable
0.0038248699
0.020
0.01
disable
0.002746582
0.025
0.01
disable
0.011555989
0.030
0.01
disable
0.0065104165
0.040
0.01
disable 0.016438803
0.050
0.01
disable 0.007710775
33

---

Published as a conference paper at ICLR 2020
Table 21: The accuracy information of tuning AdamW optimizer for ImageNet training with ResNet- 50 (batch size 16384, 90 epochs, 7038 iterations). We use the learning rate recipe of (Goyal et al., 2017): (1) 5-epoch warmup to stablize the initial stage; and (2) multiply the learning rate by 0. at 30th, 60th, and 80th epoch. The target accuracy around 0.763 (Goyal et al., 2017).
learning rate
weight decay
L2 regularization
Top-1 Validation Accuracy
0.0001
0.01
default (0.01)
0.0009969076
0.0002
0.01
default (0.01)
0.0009969076
0.0004
0.01
default (0.01)
0.0009969076
0.0006
0.01
default (0.01)
0.0009358724
0.0008
0.01
default (0.01)
0.0009969076
0.001
0.01
default (0.01)
0.0009765625
0.002
0.01
default (0.01)
0.0010172526
0.004
0.01
default (0.01)
0.0010172526
0.006
0.01
default (0.01)
0.0010172526
0.008
0.01
default (0.01)
0.0010172526
0.0001
0.001
default (0.01)
0.0010172526
0.0002
0.001
default (0.01)
0.0010172526
0.0004
0.001
default (0.01)
0.0010172526
0.0006
0.001
default (0.01)
0.0009969076
0.0008
0.001
default (0.01)
0.0010172526
0.001
0.001
default (0.01)
0.0010172526
0.002
0.001
default (0.01)
0.0010172526
0.004
0.001
default (0.01)
0.0038452148
0.006
0.001
default (0.01)
0.011881511
0.008
0.001
default (0.01)
0.0061442056
34

---

Published as a conference paper at ICLR 2020
Table 22: The accuracy information of tuning AdamW optimizer for ImageNet training with ResNet- 50 (batch size 16384, 90 epochs, 7038 iterations). We use the learning rate recipe of (Goyal et al., 2017): (1) 5-epoch warmup to stablize the initial stage; and (2) multiply the learning rate by 0. at 30th, 60th, and 80th epoch. The target accuracy around 0.763 (Goyal et al., 2017).
learning rate
weight decay
L2 regularization
Top-1 Validation Accuracy
0.0001
0.0001
default (0.01)
0.3665975
0.0002
0.0001
default (0.01)
0.5315755
0.0004
0.0001
default (0.01)
0.6369222
0.0006
0.0001
default (0.01)
0.6760457
0.0008
0.0001
default (0.01)
0.695576977
0.001
0.0001
default (0.01)
0.7076009
0.002
0.0001
default (0.01)
0.73065186
0.004
0.0001
default (0.01)
0.72806805
0.006
0.0001
default (0.01)
0.72161865
0.008
0.0001
default (0.01)
0.71816
0.0001
0.00001
default (0.01)
0.49804688
0.0002
0.00001
default (0.01)
0.6287028
0.0004
0.00001
default (0.01)
0.6773885
0.0006
0.00001
default (0.01)
0.67348224
0.0008
0.00001
default (0.01)
0.6622111
0.001
0.00001
default (0.01)
0.6468709
0.002
0.00001
default (0.01)
0.5846761
0.004
0,00001
default (0.01)
0.4868978
0.006
0.00001
default (0.01)
0.34969077
0.008
0.00001
default (0.01)
0.31193033
35

---

Published as a conference paper at ICLR 2020
Table 23: The accuracy information of tuning AdamW optimizer for ImageNet training with ResNet- 50 (batch size 16384, 90 epochs, 7038 iterations). We use the learning rate recipe of (Goyal et al., 2017): (1) 5-epoch warmup to stablize the initial stage; and (2) multiply the learning rate by 0. at 30th, 60th, and 80th epoch. The target accuracy around 0.763 (Goyal et al., 2017).
learning rate
weight decay
L2 regularization
Top-1 Validation Accuracy
0.0001
0.01
disable
0.0010172526
0.0002
0.01
disable
0.0009765625
0.0004
0.01
disable
0.0010172526
0.0006
0.01
disable
0.0009969076
0.0008
0.01
disable
0.0010172526
0.001
0.01
disable
0.0009765625
0.002
0.01
disable
0.0009969076
0.004
0.01
disable
0.0009969076
0.006
0.01
disable
0.0009765625
0.008
0.01
disable
0,0010172526
0.0001
0.001
disable
0.0009765625
0.0002
0.001
disable
0.0010172526
0.0004
0.001
disable
0.0010172526
0.0006
0.001
disable
0.0010172526
0.0008
0.001
disable
0.0010172526
0.001
0.001
disable
0.0009969076
0.002
0.001
disable
0.0010579427
0.004
0.001
disable
0.0016886393
0.006
0.001
disable
0.019714355
0.008
0.001
disable
0.1329956
36

---

Published as a conference paper at ICLR 2020
Table 24: The accuracy information of tuning AdamW optimizer for ImageNet training with ResNet- 50 (batch size 16384, 90 epochs, 7038 iterations). We use the learning rate recipe of (Goyal et al., 2017): (1) 5-epoch warmup to stablize the initial stage; and (2) multiply the learning rate by 0. at 30th, 60th, and 80th epoch. The target accuracy around 0.763 (Goyal et al., 2017).
learning rate
weight decay
L2 regularization
Top-1 Validation Accuracy
0.0001
0.0001
disable
0.28515625
0.0002
0,0001
disable
0.44055176
0.0004
0.0001
disable
0.56815594
0.0006
0.0001
disable
0.6234741
0.0008
0.0001
disable
0.6530762
0.001
0.0001
disable
0.6695964
0.002
0.0001
disable
0.70048016
0.004
0.0001
disable
0.71698
0.006
0.0001
disable
0.72021484
0.008
0.0001
disable
0.7223918
0.010
0.0001
disable
0.72017413
0.012
0.0001
disable
0.72058105
0.014
0.0001
disable
0.7188924
0.016
0.0001
disable
0.71695966
0.018
0.0001
disable
0.7154134
0.020
0.0001
disable
0.71358234
0.025
0.0001
disable
0.7145386
0.030
0.0001
disable
0.7114258
0.040
0.0001
disable
0.7066447
0.050
0.0001
disable
0.70284015
37

---

Published as a conference paper at ICLR 2020
Table 25: The accuracy information of tuning AdamW optimizer for ImageNet training with ResNet- 50 (batch size 16384, 90 epochs, 7038 iterations). We use the learning rate recipe of (Goyal et al., 2017): (1) 5-epoch warmup to stablize the initial stage; and (2) multiply the learning rate by 0. at 30th, 60th, and 80th epoch. The target accuracy around 0.763 (Goyal et al., 2017).
learning rate
weight decay
L2 regularization
Top-1 Validation Accuracy
0.0001
0,00001
disable
0.31247965
0.0002
0.00001
disable
0.4534912
0.0004
0.00001
disable
0.57765704
0.0006
0.00001
disable
0.6277669
0.0008
0.00001
disable
0.65321857
0.001
0.00001
disable
0.6682129
0.002
0.00001
disable
0.69938153
0.004
0,00001
disable
0.7095947
0.006
0.00001
disable
0.710612
0.008
0.00001
disable
0.70857745
0.010
0.00001
disable
0.7094116
0.012
0.00001
disable
0.70717365
0.014
0.00001
disable
0.7109375
0.016
0.00001
disable
0.7058309
0.018
0.00001
disable
0.7052409
0.020
0.00001
disable
0.7064412
0.025
0.00001
disable
0.7035319
0.030
0,00001
disable
0.6994629
0.040
0.00001
disable
0.6972656
0.050
0.00001
disable
0.6971232
38


Reviews:
[{"review_id": "r1g_3CFatr", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes a learning rate adaptation mechanism, called LAMB, for large-batch distributed training. The goal is to stabilize the training as the batch size increases. The idea is simple and straightforward -- there should be a layerwise learning rate adjusted by normalizing the layer weights and gradients at each layer so that layers with larger weights take larger learning steps, and vice versa. The authors perform empirical studies on BERT-large and ResNet to conclude that LAMB can scale up training batch size while still being able to converge in time with comparable accuracy.\n\nStrengths: \n+ Demonstrate the scalability of large-batch training (up to 64K) on BERT-Large with comparable accuracy. \n+ A leap from the prior work LARS that demonstrates the layer-wise learning rate adjustment scheme also works with Adam for NLP tasks.\n+ The re-warmup technique for stabilizing the second phase of mixed sequence training is neat.\n\nWeaknesses:\n- Although the authors' analysis is based on a large set of models and clearly outperforms the prior work LARS, it is still hard to assess the generality of the obtained results. The authors made an effort to show evaluation results on MNIST and CIFAR-10, but they are much less challenging tasks. \n- Technical novelty over LARS seems to be incremental, where a large portion of the work is essentially applying LARS to Adam and demonstrate its effectiveness on BERT and ResNet.\n\nOverall, the LAMB technique seems to be simple to apply yet very useful in practice for large scale training. The work can potentially help the practitioner to scale-out large model training to hundreds or even thousands of GPUs/TPUs with good scalability. Moving forward, the authors are encouraged to report LAMB optimization results on transformer-based models such as GPT, RoBERTa, and XLNet.\n\nQuestion:\nDoes the training take advantage of FP16 half-precision training?\nHow does the training process handle overflow and NaN gradients?  \nWhat is the significance of the range [α_l, α_u] in Theorem 2 and Theorem 3, and how to choose the value for them in practice?", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I did not assess the derivations or theory."}, "extra_text": {}}, {"review_id": "H1ggL7XhKS", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper developed a novel layerwise adaptation strategy, LAMB, that allows training BERT model with large mini-batches (32k vs baseline 512). This significantly speeds up the status quo in training BERT model, and effectively reduces the training time from original 3 days to only 76 minutes. In addition to demonstrating superior results across various tasks in practice, the paper also provides theoretical convergence analysis on LAMB optimizer.  \n\nOverall, this paper has made sound contributions that enables BERT-alike language to be trained with significant speedups, which is not otherwise achievable through LARS. The paper is well written and structured. I recommend acceptance. \n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I did not assess the derivations or theory.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I made a quick assessment of this paper."}, "extra_text": {}}, {"review_id": "H1eSuSc9YS", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\nIn this paper, the authors made a study on large-batch training for the BERT, and successfully trained a BERT model in 76 minutes. The results look quite exciting, however, after looking into the details of the paper, I would say that this is just a kind of RED AI – the results were mostly achieved by putting together a huge number of TPUs, without necessary technical innovation and fundamental contributions.\n\n1)\tThe work used 1024 TPUs to achieve 76-min training. If we compare this with the original BERT training (16 TPU for 81 hours), there is no algorithmic speed up at all (only system speedup). Not to mention that making a single BERT training faster by using more resources does not seem to be a big thing – one can do multiple BERT training experiments in parallel or in pipeline, which will correspond to similar innovation speed.\n\n2)\tThe theoretical analysis is not very impressive, and to certain degree, is not helpful. The theory just says that in certain conditions, both LARS and LAMB converge fasters than SGD. However, LAMB ha no advantage over LARS at all, which cannot well explain the experimental observations. Furthermore, when \\beta_2 > 0, the convergence rate of LAMB is even slower than LARS, which delivers some contradictory message. As we know, \\beta_2>0 is very important, otherwise the optimization algorithm will not be ADAM at all.\n\nOverall speaking, I am afraid that such work do not have sufficient theoretical or algorithmic contributions. And I doubt the true value of adding a huge number of computational resources to achieve speedup. \n\n\n**I read the author rebuttal. Thanks for the clarification on the algorithmic contribution of LAMB. However, my other concerns still remain. I have adjusted my rating by a little, but I can hardly move to the positive side.", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents a range of methods for over-coming the challenges of large-batch training with transformer models.  While one reviewer still questions the utility of training with such large numbers of devices, there is certainly a segment of the community that focuses on large-batch training, and the ideas in this paper will hopefully find a range of uses. 
Decision: Accept (Poster)