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 (Table 1).
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.
Strengths: 
+ Demonstrate the scalability of large-batch training (up to 64K) on BERT-Large with comparable accuracy. 
+ A leap from the prior work LARS that demonstrates the layer-wise learning rate adjustment scheme also works with Adam for NLP tasks.
+ The re-warmup technique for stabilizing the second phase of mixed sequence training is neat.
Weaknesses:
- 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. 
- 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.
Overall, 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.
Question:
Does the training take advantage of FP16 half-precision training?
How does the training process handle overflow and NaN gradients?  
What is the significance of the range [α_l, α_u] in Theorem 2 and Theorem 3, and how to choose the value for them in practice?