[Paper Review] Large Batch Training of Convolutional Networks

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

Advanced Neural Network Applications10 references509 citations

TL;DR

The paper introduces Layer-wise Adaptive Rate Scaling (LARS) to enable large-batch synchronous SGD for CNNs, enabling AlexNet and ResNet-50 to train with batch sizes up to 32K without accuracy loss.

ABSTRACT

A common way to speed up training of large convolutional networks is to add computational units. Training is then performed using data-parallel synchronous Stochastic Gradient Descent (SGD) with mini-batch divided between computational units. With an increase in the number of nodes, the batch size grows. But training with large batch size often results in the lower model accuracy. We argue that the current recipe for large batch training (linear learning rate scaling with warm-up) is not general enough and training may diverge. To overcome this optimization difficulties we propose a new training algorithm based on Layer-wise Adaptive Rate Scaling (LARS). Using LARS, we scaled Alexnet up to a batch size of 8K, and Resnet-50 to a batch size of 32K without loss in accuracy.

Motivation & Objective

• Motivate and analyze the challenges of training CNNs with large global batch sizes using data-parallel SGD.
• Investigate why standard linear learning rate scaling with warm-up may fail and cause divergence.
• Propose and evaluate a layer-wise adaptive learning rate scheme (LARS) to stabilize training.
• Demonstrate the effectiveness of LARS on AlexNet, AlexNet-BN, and ResNet-50 up to batch sizes of 32K.

Proposed method

• Analyze the instability of large-batch SGD when using a single global learning rate.
• Introduce a per-layer learning rate inspired by the ratio of weight norms to gradient norms: lambda^l = eta * ||w^l|| / (||∇L(w^l)|| + beta*||w^l||).
• Apply a global momentum and weight decay while updating each layer with its own local learning rate.
• Implement LARS within SGD with momentum and a polynomial decay schedule, enabling large minibatch training without accuracy loss.
• Use a warm-up-like strategy through the global learning rate gamma_t, while adjusting per-layer rates via lambda^l.
• Experimentally partition large batches with iter_size to simulate B up to 32K on GPUs, comparing with baseline small-batch training.

Experimental results

Research questions

• RQ1 Can large-batch CNN training match small-batch accuracy using a per-layer adaptive learning rate?
• RQ2 Does Layer-wise Adaptive Rate Scaling (LARS) stabilize training and prevent divergence for large batch sizes?
• RQ3 What batch sizes (up to 32K) and architectures (AlexNet, AlexNet-BN, ResNet-50) can be trained without accuracy loss using LARS?
• RQ4 How do learning rate strategies (warm-up, per-layer LR, BN vs non-BN) affect convergence and final accuracy?

Key findings

• LARS enables AlexNet and ResNet-50 to train with batch sizes up to 32K without reported accuracy loss in the results reported.
• AlexNet-BN with large batches (16K, 32K) remains near baseline accuracy, with some small losses that can be mitigated by tuning LR and warm-up.
• BN broadens the safe LR range for large batches and improves convergence with large learning rates.
• Without LARS, large batches diverge or incur significant accuracy loss even with warm-up and BN; with LARS, stability is improved across layers.
• Training longer with very large batches can recover accuracy when the batch is extremely large (e.g., 32K for AlexNet-BN).
• ResNet-50 scaling to 32K with LARS achieves comparable top-1 accuracy to the 256-B baseline under their training setup (though not always equal to state-of-the-art with more aggressive augmentation/testing).
• Layer-wise updates reduce sensitivity to initial weight scales and enable per-layer stabilization of training dynamics.