Large Batch Optimization for Deep Learning: Training BERT in 76 minutes — part 3/5

Cleaned-text code points [31254, 41220). Bibliography omitted here; source.md includes it.

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 rates. Furthermore, Lamb is a general purpose optimizer that works for both small and large batches. We also provided theoretical analysis for the Lamb optimizer, highlighting the cases where it performs better than standard Sgd . Lamb achieves a 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 is also the first large batch adaptive solver that can achieve state-of-the-art 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.

## Appendix

## Appendix A Proof of Theorem 2

###### Proof.

We analyze the convergence of Lars for general minibatch size here. Recall that the update of Lars is the following x t + 1 ( i ) = x t ( i ) − η t ​ ϕ ​ ( ‖ x t ( i ) ‖ ) ​ g t ( i ) ‖ g t ( i ) ‖ , x_{t+1}^{(i)}=x_{t}^{(i)}-\eta_{t}\phi(\|x_{t}^{(i)}\|)\frac{g_{t}^{(i)}}{\|g_{t}^{(i)}\|}, for all i ∈ [ h ] i\in[h] . For simplicity of notation, we reason the

Since the function f f is L L -smooth, we have the following: f ⁡ ( x t + 1 ) ≤ f ⁡ ( x t ) + ⟨ ∇ i f ​ ( x t ) , x t + 1 ( i ) − x t ( i ) ⟩ + ∑ i = 1 h L i 2 ​ ‖ x t + 1 ( i ) − x t ( i ) ‖ 2 \displaystyle f(x_{t+1})\leq f(x_{t})+\langle\nabla_{i}f(x_{t}),x_{t+1}^{(i)}-x_{t}^{(i)}\rangle+\sum_{i=1}^{h}\frac{L_{i}}{2}\|x_{t+1}^{(i)}-x_{t}^{(i)}\|^{2} = f ⁡ ( x t ) − η t ​ ∑ i = 1 h ∑ j = 1 d i ϕ ⁡ ( ‖ x t ( i ) ‖ ) × ( [ ∇ i f ​ ( x t ) ] j × g t , j ( i ) ‖ g t ( i ) ‖ ) + ∑ i = 1 h L i ​ η t 2 ​ ϕ 2 ​ ( ‖ x t ( i ) ‖ ) 2 \displaystyle=f(x_{t})-\eta_{t}\sum_{i=1}^{h}\sum_{j=1}^{d_{i}}\phi(\|x_{t}^{(i)}\|)\times\left([\nabla_{i}f(x_{t})]_{j}\times\frac{g_{t,j}^{(i)}}{\|g_{t}^{(i)}\|}\right)+\sum_{i=1}^{h}\frac{L_{i}\eta_{t}^{2}\phi^{2}(\|x_{t}^{(i)}\|)}{2} ≤ f ⁡ ( x t ) − η t ​ ∑ i = 1 h ∑ j = 1 d i ϕ ⁡ ( ‖ x t ( i ) ‖ ) × ( [ ∇ i f ​ ( x t ) ] j × ( g t , j ( i ) ‖ g t ( i ) ‖ − [ ∇ i f ​ ( x t ) ] j ‖ ∇ i f ​ ( x t ) ‖ + [ ∇ i f ​ ( x t ) ] j ‖ ∇ i f ​ ( x t ) ‖ ) ) + η t 2 ​ α u 2 2 ​ ‖ L ‖ 1 \displaystyle\leq f(x_{t})-\eta_{t}\sum_{i=1}^{h}\sum_{j=1}^{d_{i}}\phi(\|x_{t}^{(i)}\|)\times\left([\nabla_{i}f(x_{t})]_{j}\times\left(\frac{g_{t,j}^{(i)}}{\|g_{t}^{(i)}\|}-\frac{[\nabla_{i}f(x_{t})]_{j}}{\|\nabla_{i}f(x_{t})\|}+\frac{[\nabla_{i}f(x_{t})]_{j}}{\|\nabla_{i}f(x_{t})\|}\right)\right)+\frac{\eta_{t}^{2}\alpha_{u}^{2}}{2}\|L\|_{1} = f ⁡ ( x t ) − η t ​ ∑ i = 1 h ϕ ⁡ ( ‖ x t ( i ) ‖ ) × ‖ ∇ i f ​ ( x t ) ‖ − η t ​ ∑ i = 1 h ∑ j = 1 d i ( [ ∇ i f ​ ( x t ) ] j × ( g t , j ( i ) ‖ g t ( i ) ‖ − [ ∇ i f ​ ( x t ) ] j ‖ ∇ i f ​ ( x t ) ‖ ) ) + η t 2 ​ α u 2 2 ​ ‖ L ‖ 1 \displaystyle=f(x_{t})-\eta_{t}\sum_{i=1}^{h}\phi(\|x_{t}^{(i)}\|)\times\|\nabla_{i}f(x_{t})\|-\eta_{t}\sum_{i=1}^{h}\sum_{j=1}^{d_{i}}\left([\nabla_{i}f(x_{t})]_{j}\times\left(\frac{g_{t,j}^{(i)}}{\|g_{t}^{(i)}\|}-\frac{[\nabla_{i}f(x_{t})]_{j}}{\|\nabla_{i}f(x_{t})\|}\right)\right)+\frac{\eta_{t}^{2}\alpha_{u}^{2}}{2}\|L\|_{1} (4) The first inequality follows from the lipschitz continuous nature of the gradient. Let Δ t ( i ) = g t ( i ) − ∇ i f ​ ( x t ) \Delta_{t}^{(i)}=g_{t}^{(i)}-\nabla_{i}f(x_{t}) . Then the above inequality can be rewritten in the following manner: f ⁡ ( x t + 1 ) ≤ f ⁡ ( x t ) − η t ​ ∑ i = 1 h ϕ ⁡ ( ‖ x t ( i ) ‖ ) ​ ‖ ∇ i f ​ ( x t ) ‖ \displaystyle f(x_{t+1})\leq f(x_{t})-\eta_{t}\sum_{i=1}^{h}\phi(\|x_{t}^{(i)}\|)\|\nabla_{i}f(x_{t})\| − η t ∑ i = 1 h ∑ j = 1 d i ϕ ( ∥ x t ( i ) ∥ ) × ( [ ∇ i f ( x t ) ] j × ( ( Δ t , j ( i ) + [ ∇ i f ​ ( x t ) ] j ) ‖ Δ t ( i ) + ∇ i f ​ ( x t ) ‖ − [ ∇ i f ​ ( x t ) ] j ‖ ∇ i f ​ ( x t ) ‖ ) ) + η t 2 ​ α u 2 2 ∥ L ∥ 1 \displaystyle\qquad\qquad-\eta_{t}\sum_{i=1}^{h}\sum_{j=1}^{d_{i}}\phi(\|x_{t}^{(i)}\|)\times\left([\nabla_{i}f(x_{t})]_{j}\times\left(\frac{(\Delta_{t,j}^{(i)}+[\nabla_{i}f(x_{t})]_{j})}{\|\Delta_{t}^{(i)}+\nabla_{i}f(x_{t})\|}-\frac{[\nabla_{i}f(x_{t})]_{j}}{\|\nabla_{i}f(x_{t})\|}\right)\right)+\frac{\eta_{t}^{2}\alpha_{u}^{2}}{2}\|L\|_{1} = f ⁡ ( x t ) − η t ​ ∑ i = 1 h ϕ ⁡ ( ‖ x t ( i ) ‖ ) ​ ‖ ∇ i f ​ ( x t ) ‖ \displaystyle=f(x_{t})-\eta_{t}\sum_{i=1}^{h}\phi(\|x_{t}^{(i)}\|)\|\nabla_{i}f(x_{t})\| − η t ∑ i = 1 h ϕ ( ∥ x t ( i ) ∥ ) × ( ⟨ Δ t ( i ) + ∇ i f ​ ( x t ) , ∇ i f ​ ( x t ) ⟩ ‖ Δ t ( i ) + ∇ i f ​ ( x t ) ‖ − ∥ ∇ i f ( x t ) ∥ ) + η t 2 ​ α u 2 2 ∥ L ∥ 1 \displaystyle\qquad\qquad-\eta_{t}\sum_{i=1}^{h}\phi(\|x_{t}^{(i)}\|)\times\left(\frac{\langle\Delta_{t}^{(i)}+\nabla_{i}f(x_{t}),\nabla_{i}f(x_{t})\rangle}{\|\Delta_{t}^{(i)}+\nabla_{i}f(x_{t})\|}-\|\nabla_{i}f(x_{t})\|\right)+\frac{\eta_{t}^{2}\alpha_{u}^{2}}{2}\|L\|_{1} = f ⁡ ( x t ) − η t ​ ∑ i = 1 h ϕ ⁡ ( ‖ x t ( i ) ‖ ) ​ ‖ ∇ i f ​ ( x t ) ‖ \displaystyle=f(x_{t})-\eta_{t}\sum_{i=1}^{h}\phi(\|x_{t}^{(i)}\|)\|\nabla_{i}f(x_{t})\| + η t ∑ i = 1 h ϕ ( ∥ x t ( i ) ∥ ) × ( ‖ ∇ i f ​ ( x t ) ‖ ​ ‖ Δ t ( i ) + ∇ i f ​ ( x t ) ‖ − ⟨ Δ t ( i ) + ∇ i f ​ ( x t ) , ∇ i f ​ ( x t ) ⟩ ‖ Δ t ( i ) + ∇ i f ​ ( x t ) ‖ ) + η t 2 ​ α u 2 2 ∥ L ∥ 1 \displaystyle\qquad\qquad+\eta_{t}\sum_{i=1}^{h}\phi(\|x_{t}^{(i)}\|)\times\left(\frac{\|\nabla_{i}f(x_{t})\|\|\Delta_{t}^{(i)}+\nabla_{i}f(x_{t})\|-\langle\Delta_{t}^{(i)}+\nabla_{i}f(x_{t}),\nabla_{i}f(x_{t})\rangle}{\|\Delta_{t}^{(i)}+\nabla_{i}f(x_{t})\|}\right)+\frac{\eta_{t}^{2}\alpha_{u}^{2}}{2}\|L\|_{1} = f ⁡ ( x t ) − η t ​ ∑ i = 1 h ϕ ⁡ ( ‖ x t ( i ) ‖ ) ​ ‖ ∇ i f ​ ( x t ) ‖ + η t 2 ​ α u 2 2 ​ ‖ L ‖ 1 \displaystyle=f(x_{t})-\eta_{t}\sum_{i=1}^{h}\phi(\|x_{t}^{(i)}\|)\|\nabla_{i}f(x_{t})\|+\frac{\eta_{t}^{2}\alpha_{u}^{2}}{2}\|L\|_{1} + η t ∑ i = 1 h ϕ ( ∥ x t ( i ) ∥ ) × ( ‖ ∇ i f ​ ( x t ) ‖ ​ ‖ Δ t ( i ) + ∇ i f ​ ( x t ) ‖ − ‖ Δ t ( i ) + ∇ i f ​ ( x t ) ‖ 2 + ⟨ Δ t ( i ) , Δ t ( i ) + ∇ i f ​ ( x t ) ⟩ ‖ Δ t ( i ) + ∇ i f ​ ( x t ) ‖ ) . \displaystyle\qquad\quad+\eta_{t}\sum_{i=1}^{h}\phi(\|x_{t}^{(i)}\|)\times\left(\frac{\|\nabla_{i}f(x_{t})\|\|\Delta_{t}^{(i)}+\nabla_{i}f(x_{t})\|-\|\Delta_{t}^{(i)}+\nabla_{i}f(x_{t})\|^{2}+\langle\Delta_{t}^{(i)},\Delta_{t}^{(i)}+\nabla_{i}f(x_{t})\rangle}{\|\Delta_{t}^{(i)}+\nabla_{i}f(x_{t})\|}\right). (5)

Using Cauchy-Schwarz inequality in the above inequality, we have: f ⁡ ( x t + 1 ) \displaystyle f(x_{t+1}) ≤ f ⁡ ( x t ) − η t ​ ∑ i = 1 h ϕ ⁡ ( ‖ x t ( i ) ‖ ) ​ ‖ ∇ i f ​ ( x t ) ‖ \displaystyle\leq f(x_{t})-\eta_{t}\sum_{i=1}^{h}\phi(\|x_{t}^{(i)}\|)\|\nabla_{i}f(x_{t})\| + η t ∑ i = 1 h ϕ ( ∥ x t ( i ) ∥ ) × ( ∥ ∇ i f ( x t ) ∥ − ∥ Δ t ( i ) + ∇ i f ( x t ) ∥ + ∥ Δ t ( i ) ∥ ) + η t 2 ​ α u 2 2 ∥ L ∥ 1 \displaystyle\qquad\qquad+\eta_{t}\sum_{i=1}^{h}\phi(\|x_{t}^{(i)}\|)\times\left(\|\nabla_{i}f(x_{t})\|-\|\Delta_{t}^{(i)}+\nabla_{i}f(x_{t})\|+\|\Delta_{t}^{(i)}\|\right)+\frac{\eta_{t}^{2}\alpha_{u}^{2}}{2}\|L\|_{1} ≤ f ⁡ ( x t ) − η t ​ ∑ i = 1 h ϕ ⁡ ( ‖ x t ( i ) ‖ ) ​ ‖ ∇ i f ​ ( x t ) ‖ + 2 ​ η t ​ ∑ i = 1 h ϕ ⁡ ( ‖ x t ( i ) ‖ ) × ‖ Δ t ( i ) ​ ‖ + η t 2 ​ α u 2 2 ‖ ​ L ‖ 1 \displaystyle\leq f(x_{t})-\eta_{t}\sum_{i=1}^{h}\phi(\|x_{t}^{(i)}\|)\|\nabla_{i}f(x_{t})\|+2\eta_{t}\sum_{i=1}^{h}\phi(\|x_{t}^{(i)}\|)\times\|\Delta_{t}^{(i)}\|+\frac{\eta_{t}^{2}\alpha_{u}^{2}}{2}\|L\|_{1}

Taking expectation, we obtain the following: 𝔼 ⁡ [ f ⁡ ( x t + 1 ) ] \displaystyle\mathbb{E}[f(x_{t+1})] ≤ f ⁡ ( x t ) − η t ​ ∑ i = 1 h ϕ ⁡ ( ‖ x t ( i ) ‖ ) ​ ‖ ∇ i f ​ ( x t ) ‖ + 2 ​ η t ​ ∑ i = 1 h ϕ ⁡ ( ‖ x t ( i ) ‖ ) × 𝔼 ⁡ [ ‖ Δ t ( i ) ‖ ] + η t 2 ​ α u 2 2 ​ ‖ L ‖ 1 \displaystyle\leq f(x_{t})-\eta_{t}\sum_{i=1}^{h}\phi(\|x_{t}^{(i)}\|)\|\nabla_{i}f(x_{t})\|+2\eta_{t}\sum_{i=1}^{h}\phi(\|x_{t}^{(i)}\|)\times\mathbb{E}[\|\Delta_{t}^{(i)}\|]+\frac{\eta_{t}^{2}\alpha_{u}^{2}}{2}\|L\|_{1} ≤ f ⁡ ( x t ) − η t ​ α l ​ ∑ i = 1 h ‖ ∇ i f ​ ( x t ) ‖ + 2 ​ η t ​ α u ​ ‖ σ ‖ 1 b + η t 2 ​ α u 2 2 ​ ‖ L ‖ 1 . \displaystyle\leq f(x_{t})-\eta_{t}\alpha_{l}\sum_{i=1}^{h}\|\nabla_{i}f(x_{t})\|+2\eta_{t}\alpha_{u}\frac{\|\sigma\|_{1}}{\sqrt{b}}+\frac{\eta_{t}^{2}\alpha_{u}^{2}}{2}\|L\|_{1}. (6) Summing the above inequality for t = 1 t=1 to T T and using telescoping sum, we have the following inequality: 𝔼 ⁡ [ f ⁡ ( x T + 1 ) ] \displaystyle\mathbb{E}[f(x_{T+1})] ≤ f ⁡ ( x 1 ) − η ​ α l ​ ∑ t = 1 T ∑ i = 1 h 𝔼 ⁡ [ ‖ ∇ i f ​ ( x t ) ‖ ] + 2 ​ η ​ T ​ α u ​ ‖ σ ‖ 1 b + η 2 ​ α u 2 ​ T 2 ​ ‖ L ‖ 1 . \displaystyle\leq f(x_{1})-\eta\alpha_{l}\sum_{t=1}^{T}\sum_{i=1}^{h}\mathbb{E}[\|\nabla_{i}f(x_{t})\|]+2\eta T\frac{\alpha_{u}\|\sigma\|_{1}}{\sqrt{b}}+\frac{\eta^{2}\alpha_{u}^{2}T}{2}\|L\|_{1}. Rearranging the terms of the above inequality, and dividing by η ​ T ​ α l \eta T\alpha_{l} , we have: 1 T ​ ∑ t = 1 T ∑ i = 1 h 𝔼 ⁡ [ ‖ ∇ i f ​ ( x t ) ‖ ] \displaystyle\frac{1}{T}\sum_{t=1}^{T}\sum_{i=1}^{h}\mathbb{E}[\|\nabla_{i}f(x_{t})\|] ≤ f ⁡ ( x 1 ) − 𝔼 ⁡ [ f ⁡ ( x T + 1 ) ] T ​ η ​ α l + 2 ​ α u ​ ‖ σ ‖ 1 b ​ α l + η ​ α u 2 2 ​ α l ​ ‖ L ‖ 1 \displaystyle\leq\frac{f(x_{1})-\mathbb{E}[f(x_{T+1})]}{T\eta\alpha_{l}}+\frac{2\alpha_{u}\|\sigma\|_{1}}{\sqrt{b}\alpha_{l}}+\frac{\eta\alpha_{u}^{2}}{2\alpha_{l}}\|L\|_{1} ≤ f ⁡ ( x 1 ) − f ⁡ ( x ∗ ) T ​ η ​ α l + 2 ​ α u ​ ‖ σ ‖ 1 α l ​ b + η ​ α u 2 2 ​ α l ​ ‖ L ‖ 1 . \displaystyle\leq\frac{f(x_{1})-f(x^{*})}{T\eta\alpha_{l}}+\frac{2\alpha_{u}\|\sigma\|_{1}}{\alpha_{l}\sqrt{b}}+\frac{\eta\alpha_{u}^{2}}{2\alpha_{l}}\|L\|_{1}.

∎

## Appendix 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 x t + 1 ( i ) = x t ( i ) − η t ​ ϕ ​ ( ‖ x t ( i ) ‖ ) ​ r t ( i ) ‖ r t ( i ) ‖ , x_{t+1}^{(i)}=x_{t}^{(i)}-\eta_{t}\phi(\|x_{t}^{(i)}\|)\frac{r_{t}^{(i)}}{\|r_{t}^{(i)}\|}, for all i ∈ [ h ] i\in[h] . For simplicity of notation, we reason the