[Paper Review] Scaling SGD Batch Size to 32K for ImageNet Training.

Yang You, Igor Gitman+1 more|arXiv (Cornell University)|Aug 13, 2017

Advanced Neural Network Applications7 references225 citations

TL;DR

This paper proposes Layer-wise Adaptive Rate Scaling (LARS), a learning rate scheduling method that adapts the learning rate per layer based on weight and gradient norms, enabling stable large-batch training. Using LARS, the authors successfully scale SGD batch sizes to 32,768 for ResNet-50 and 8,192 for AlexNet on ImageNet-1k, achieving significant speedups—e.g., 3× faster training at batch 4096 vs. batch 512—without accuracy degradation.

ABSTRACT

The most natural way to speed-up the training of large networks is to use data-parallelism on multiple GPUs. To scale Stochastic Gradient (SG) based methods to more processors, one need to increase the batch size to make full use of the computational power of each GPU. However, keeping the accuracy of network with increase of batch size is not trivial. Currently, the state-of-the art method is to increase Learning Rate (LR) proportional to the batch size, and use special learning rate with "warm-up" policy to overcome initial optimization difficulty. <p>By controlling the LR during the training process, one can efficiently use large-batch in ImageNet training. For example, Batch-1024 for AlexNet and Batch-8192 for ResNet-50 are successful applications. However, for ImageNet-1k training, state-of-the-art AlexNet only scales the batch size to 1024 and ResNet50 only scales it to 8192. The reason is that we can not scale the learning rate to a large value. To enable large-batch training to general networks or datasets, we propose Layer-wise Adaptive Rate Scaling (LARS). LARS LR uses different LRs for different layers based on the norm of the weights and the norm of the gradients. By using LARS algoirithm, we can scale the batch size to 32768 for ResNet50 and 8192 for AlexNet. Large batch can make full use of the system's computational power. For example, batch-4096 can achieve 3x speedup over batch-512 for ImageNet training by AlexNet model on a DGX-1 station (8 P100 GPUs).

Motivation & Objective

• To address the challenge of maintaining model accuracy when scaling SGD batch sizes beyond 8,192 in ImageNet training.
• To overcome the instability and accuracy drop typically observed when increasing batch size beyond a certain threshold.
• To enable efficient utilization of multi-GPU systems by scaling batch size to 32,768 while preserving convergence and performance.
• To develop a learning rate scheduling strategy that adapts to layer-specific weight and gradient norms for stable large-batch optimization.

Proposed method

• Proposes Layer-wise Adaptive Rate Scaling (LARS), which computes a layer-specific learning rate as the product of a global base learning rate and a scaling factor based on the ratio of the L2 norm of layer weights to the L2 norm of the gradients.
• Applies a global base learning rate that is scaled proportionally to the batch size, following the common practice of linear scaling rule.
• Introduces a warm-up phase for the learning rate to stabilize training during early epochs, especially for large batches.
• Employs a weight decay component to regularize the model and improve generalization during large-batch training.
• Uses a consistent training protocol across models (AlexNet, ResNet-50) with fixed hyperparameters except for learning rate scheduling.
• Employs data-parallel training across 8 P100 GPUs on a DGX-1 system to evaluate scalability and speedup.

Experimental results

Research questions

• RQ1 Can large-batch SGD training (up to 32,768) be stabilized and made accurate for ImageNet-1k using a novel learning rate scheduling strategy?
• RQ2 How does layer-wise adaptive learning rate scaling (LARS) compare to standard linear scaling and warm-up in terms of convergence and accuracy at large batch sizes?
• RQ3 To what extent can training speedup be achieved by increasing batch size from 512 to 4,096 while maintaining model accuracy?
• RQ4 Does LARS enable stable training for both deeper (ResNet-50) and shallower (AlexNet) networks at extreme batch sizes?

Key findings

• LARS enables stable training with a batch size of 32,768 for ResNet-50 on ImageNet-1k, achieving state-of-the-art performance at scale.
• For AlexNet, LARS successfully scales the batch size to 8,192, significantly exceeding prior state-of-the-art limits of 1,024.
• At batch size 4,096, training with LARS achieves a 3× speedup over batch size 512 on an 8-GPU DGX-1 system using the AlexNet model.
• The combination of linear scaling of the base learning rate and LARS' layer-wise adaptation maintains model accuracy even at extreme batch sizes.
• The warm-up phase in LARS effectively mitigates initial optimization instability during large-batch training.
• LARS enables consistent performance across different network architectures, including both ResNet-50 and AlexNet, at large-batch settings.