Accelerated Large Batch Optimization of BERT Pretraining in 54 minutes

For the gradient $g_{t}\in{\mathbb{R}}^{d}$, let $g_{t}=[g_{t,\mathcal{G}_{1}},g_{t,\mathcal{G}_{2}},\dots,g_{t,\mathcal{G}_{B}}]$ be its decomposition into $B$ blocks, where $\mathcal{G}_{b}$ is the set of indices in block $b$, and $g_{t,\mathcal{G}_{b}}$ is the corresponding block of variables. A block can be a parameter tensor/matrix/vector. Recently, You et al. [30] introduced a layer-wise adaptive large batch optimization method, called LAMB (Algorithm 1). LAMB proposed to add a normalization factor to the AdamW (ADAM with weight decay) [16] update that divides the update by its $\ell_{2}$ norm. This ensures that the update for each block has unit $\ell_{2}$-norm. And, the learning rate is rescaled by $\phi(\|x_{t,\mathcal{G}_{b}}\|_{2})$ for some function $\phi:{\mathbb{R}}^{+}\rightarrow{\mathbb{R}}^{+}$. In practice, it is generally set to an identity mapping. In this case, the update preserves the same $\ell_{2}$ norm as the model parameters, and the model parameters change in a smooth trajectory.

It was shown in [30] that LAMB enables a very large mini-batch size of 64K for training BERT while being able to achieve comparable accuracy to the small mini-batch size. With such large mini-batch size, the training time of BERT pretraining reduced from 3 days to 76 minutes on 1024 TPUs.

Momentum methods have been widely used in training deep networks [26]. The classic momentum method, also known as heavy-ball method, introduced in [20] accumulates the past gradients $g_{t}$’s into a momentum vector $m_{t}$ (with $m_{0}=0$), which serves as a smoothing of the velocity:

where $\mu\in[0,1)$ is the momentum parameter. For a twice differentiable strongly convex function, it is known that the classic momentum method can be used to accelerate the gradient descent and improve the convergence rate from $O((1-\kappa)^{t})$ to $O((1-\sqrt{\kappa})^{t})$. where $\kappa$ is the condition number of the functions. For training deep neural network, the momentum method accelerates early optimization and helps gradient descent method escape from the local minimums.

Nesterov’s accelerated gradient (NAG) [18, 26] is another kind of momentum method that is closely related to the classic momentum method in that it can be written as:

Expanding (3) to $x_{t+1}=x_{t}-\eta_{t}(\mu m_{t-1}+g_{t})$, we can see that we get NAG by replacing $m_{t-1}$ with $m_{t}$. Thus, Nesterov’s momentum differs from the classic momentum in that it updates the model parameter using the future momentum vector. One can interpret Nesterov’s momentum as an attempt to add a correction direction to the classic momentum method. NAG is argued to be more effective in the early optimization, and is more tolerant of large values of $\mu$ compared to the classic momentum method [26]. Recently, Adam with Nesterov’s momentum is proposed in [6] and it shows better convergence performance than Adam on some tasks. Inspired by this change, two variants of LAMB using Nesterov’s momentum are proposed in [30]. However, their modifications do not take the normalization factor into account, and the resultant algorithms do not show any improvement over LAMB. In this paper, we propose a different way to modify the momentum component of LAMB to take advantage of the superior performance of Nesterov’s acceleration.

In LAMB, the update is normalized by its $\ell_{2}$ norm. In addition to that, we propose to normalize the gradient in each block:

Then, we use $\tilde{g}_{t,\mathcal{G}_{b}}$ to update first-order and second-order momentums $m_{t}$ and $v_{t}$, respectively. This technique was first introduced in [31] for accelerating Adam in training deep neural networks. Using the per-block gradient normalization, the gradient clipping is no longer necessary. Ignoring the gradient magnitude makes the gradient descent methods more robust to vanishing and exploding gradients.

In order to incorporate Nesterov’s momentum, we first rewrite the step 11 in Algorithm 1 as

To apply the same trick as in Nesterov’s momentum, first we substitute $m_{t,\mathcal{G}_{b}}$ for $m_{t-1,\mathcal{G}_{b}}$ in (6), and then we modify the normalization factors to ensure unit $\ell_{2}$-norm for both (6) and (6), leading to the following new update rule:

where $a_{t,\mathcal{G}_{b}}=\frac{g_{t,\mathcal{G}_{b}}}{\sqrt{\tilde{v}_{t,\mathcal{G}_{b}}}+\epsilon}$. Note that we remove the factor $1/(1-\beta_{1}^{t})$ in (6) for (7), as this factor leads to a bias towards $g_{t,\mathcal{G}_{b}}$ when the normalization is modified and regularization parameter $\lambda>0$. Interestingly, the resultant update is simply a convex combination between LAMB updates with and without first-order momentum. Combining (7) with (4), we obtain Algorithm 2.

For large mini-batch optimization, warmup is usually used at the start of the training [9]. Goyal et al. [9] proposed to use a linear warmup in the beginning and return to the original learning rate schedule afterwards. For BERT pretraining, LAMB uses a learning rate schedule of form [30]

where $\eta>0$ is the maximum learning rate that the optimization algorithms use throughout the training and $T_{warmup}$ denotes the number of iterations in $warmup$ stage. It can be seen that $\eta_{t}$ gradually increases to $\eta$ when $t$ approaches $T_{warmup}$ and decreases to $0$ when $t\rightarrow T$. In [30], a square root scheduling rule is proposed for increasing mini-batch size: $\eta=\sqrt{k}\tilde{\eta}$, where $k$ is the mini-batch size and $\tilde{\eta}$ is a reference learning rate for a small mini-batch size. To achieve speedup using $\tau$ times larger mini-batch size, the number of training iterations $T$ is reduced by $\tau$ times while the learning rate $\eta$ is increased by $\sqrt{\tau}$ times. Using such scheduler, LAMB successfully scaled BERT pretraining up to a mini-batch size of 32K without any accuracy loss. For a larger mini-batch size, this square root scheduling breaks as it exceeds a maximum rate that does not depend on the mini-batch size. Thus, a smaller learning rate is used for a mini-batch size of 64K with a small degradation of accuracy.

Considering that the learning rate is theoretically upper bounded by the inverse of the Lipschitz constant $L$ [17, 8, 11, 30] up to some small constants (e.g., $1$ or $2$), one cannot scale the learning rate indefinitely. We propose to add a $constant$ transient phase after the $warmup$ stage as shown below

where $T_{const}$ is the number of iterations in which a constant learning rate is used. This scheme allows the training to have sufficient progress even one cannot further increase $\eta$.

Figure 1 shows visualization of (8) with $\eta=0.007,0.01$ and (9) with $\eta=0.007$. $\eta=0.01$ refers to the ideal learning rate that we scaled the mini-batch size from 32K to 128K. However, $0.01$ has exceeded the maximum learning rate and results in divergence. Therefore, we have to use a smaller one such as $0.007$. Nonetheless, this smaller learning rate downgrades performance. In particular, the difference between areas under curve of (8) with $\eta=0.007,0.01$ is $5.28$. Using the proposed schedule (9), we can reduce the difference to $1.91$.

In large-scale mini-batch training, the quality of the mini-batch plays an important role. In order to use large learning rate, one need to have as small gradient variance as possible. For example, random sampling with replacement results in a variance bound of $O(\frac{\sigma^{2}}{k})$ [4] while random sampling without replacement gives a better bound of $O(\frac{n-k}{k(n-1)}\sigma^{2})$ [15], where $\sigma^{2}$ is the upper bound of the gradient variance. It can be seen that the variance only goes to zero when $k\rightarrow\infty$ for random sampling with replacement while the variance is zero when $k=n$ for random sampling without replacement. Thus, random sampling without replacement results in better efficiency of using the same mini-batch size. In distributed training, to make sure that the mini-batch does not have redundant samples, we only grant each worker access to a shard of the dataset. Within each shard, random shuffling is used to construct the mini-batch samples.