Large Batch Optimization for Deep Learning: Training BERT in 76 minutes

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 (Table 1). The LAMB implementation is available at https://github.com/tensorflow/addons/blob/master/tensorflow_addons/optimizers/lamb.py

Published as a conference paper at ICLR 2020

References in corpus (21)

• Adam: A Method for Stochastic Optimization
• Decoupled Weight Decay Regularization
• Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour
• One weird trick for parallelizing convolutional neural networks
• Don't Decay the Learning Rate, Increase the Batch Size
• On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima
• Large Batch Training of Convolutional Networks
• Optimizing Neural Networks with Kronecker-factored Approximate Curvature
• Train longer, generalize better: closing the generalization gap in large batch training of neural networks
• Highly Scalable Deep Learning Training System with Mixed-Precision: Training ImageNet in Four Minutes
• Aligning Books and Movies: Towards Story-like Visual Explanations by Watching Movies and Reading Books
• Extremely Large Minibatch SGD: Training ResNet-50 on ImageNet in 15 Minutes
• Measuring the Effects of Data Parallelism on Neural Network Training
• AdaBatch: Adaptive Batch Sizes for Training Deep Neural Networks
• Image Classification at Supercomputer Scale
• signSGD: Compressed Optimisation for Non-Convex Problems
• Massively Distributed SGD: ImageNet/ResNet-50 Training in a Flash
• Yet Another Accelerated SGD: ResNet-50 Training on ImageNet in 74.7 seconds
• Scale out for large minibatch SGD: Residual network training on ImageNet-1K with improved accuracy and reduced time to train
• Large-Scale Distributed Second-Order Optimization Using Kronecker-Factored Approximate Curvature for Deep Convolutional Neural Networks
• Large-Batch Training for LSTM and Beyond

Cited by in corpus (60)

• Knowledge Distillation: A Survey
• ALBERT: A Lite BERT for Self-supervised Learning of Language Representations
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• On the Cross-lingual Transferability of Monolingual Representations
• Explainable COVID-19 Detection Using Chest CT Scans and Deep Learning
• CANINE: Pre-training an Efficient Tokenization-Free Encoder for Language Representation
• Optimization for deep learning: theory and algorithms
• DeCLUTR: Deep Contrastive Learning for Unsupervised Textual Representations
• MobileBERT: a Compact Task-Agnostic BERT for Resource-Limited Devices
• NEZHA: Neural Contextualized Representation for Chinese Language Understanding
• Learning to Decode the Surface Code with a Recurrent, Transformer-Based Neural Network
• AdamP: Slowing Down the Slowdown for Momentum Optimizers on Scale-invariant Weights
• Single Headed Attention RNN: Stop Thinking With Your Head
• On the Effectiveness of Transfer Learning for Code Search
• Audio ALBERT: A Lite BERT for Self-supervised Learning of Audio Representation
• Federated Learning with Buffered Asynchronous Aggregation
• SEED RL: Scalable and Efficient Deep-RL with Accelerated Central Inference
• Lost in Pruning: The Effects of Pruning Neural Networks beyond Test Accuracy
• Learning image representations for anomaly detection: application to discovery of histological alterations in drug development
• Physics-inspired deep learning to characterize the signal manifold of quasi-circular, spinning, non-precessing binary black hole mergers
• Yuan 1.0: Large-Scale Pre-trained Language Model in Zero-Shot and Few-Shot Learning
• A community-powered search of machine learning strategy space to find NMR property prediction models
• Training Large Neural Networks with Constant Memory using a New Execution Algorithm
• Local AdaAlter: Communication-Efficient Stochastic Gradient Descent with Adaptive Learning Rates
• MTOP: A Comprehensive Multilingual Task-Oriented Semantic Parsing Benchmark
• AraGPT2: Pre-Trained Transformer for Arabic Language Generation
• Towards Practical Adam: Non-Convexity, Convergence Theory, and Mini-Batch Acceleration
• Low-Rank Bottleneck in Multi-head Attention Models
• Learning To Retrieve: How to Train a Dense Retrieval Model Effectively and Efficiently
• Distributed Learning of Deep Neural Networks using Independent Subnet Training
• HPC AI500: The Methodology, Tools, Roofline Performance Models, and Metrics for Benchmarking HPC AI Systems
• A Second look at Exponential and Cosine Step Sizes: Simplicity, Adaptivity, and Performance
• Adam${}^{+}$: A Stochastic Method with Adaptive Variance Reduction
• FSMoE: A Flexible and Scalable Training System for Sparse Mixture-of-Experts Models
• Optimizing Dense Retrieval Model Training with Hard Negatives
• MC-BERT: Efficient Language Pre-Training via a Meta Controller
• Progressively Stacking 2.0: A Multi-stage Layerwise Training Method for BERT Training Speedup
• KAISA: An Adaptive Second-Order Optimizer Framework for Deep Neural Networks
• Exploring the limits of Concurrency in ML Training on Google TPUs
• Large-Scale Deep Learning Optimizations: A Comprehensive Survey
• Variance-reduced Language Pretraining via a Mask Proposal Network
• CBAG: Conditional Biomedical Abstract Generation
• Forming Ensembles at Runtime: A Machine Learning Approach
• Morphological Disambiguation from Stemming Data
• Go Wider Instead of Deeper
• Demystifying BERT: Implications for Accelerator Design
• Benchmark Performance of Machine And Deep Learning Based Methodologies for Urdu Text Document Classification
• Pointer Value Retrieval: A new benchmark for understanding the limits of neural network generalization
• Deep-learning-based decomposition of overlapping-sparse images: application at the vertex of neutrino interactions
• Extrapolation for Large-batch Training in Deep Learning
• Acceleration via Fractal Learning Rate Schedules
• Drill the Cork of Information Bottleneck by Inputting the Most Important Data
• AGATHA: Automatic Graph-mining And Transformer based Hypothesis generation Approach
• Image Captions are Natural Prompts for Text-to-Image Models
• Stochastic Normalized Gradient Descent with Momentum for Large-Batch Training
• Online Evolutionary Batch Size Orchestration for Scheduling Deep Learning Workloads in GPU Clusters
• Contrastive Weight Regularization for Large Minibatch SGD
• ResIST: Layer-Wise Decomposition of ResNets for Distributed Training
• LightSeq2: Accelerated Training for Transformer-based Models on GPUs
• Improving Layer-wise Adaptive Rate Methods using Trust Ratio Clipping