[논문 리뷰] Large Batch Optimization for Deep Learning: Training BERT in 76 minutes

Yang You, Jing Li외 8명|arXiv (Cornell University)|2019. 04. 01.

Advanced Neural Network Applications참고 문헌 36인용 수 164

한 줄 요약

Lamb를 도입한 계층별 적응형 대배치 최적화기, Bert 학습을 최대 32K 배치로 가능하게 하고 TPUv3 포드에서 Bert 사전학습 시간을 3일에서 76분으로 단축하며 Lamb와 Lars의 수렴 보장을 제공합니다.

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

연구 동기 및 목표

• 큰 미니배치로 인해 큰 신경망의 학습 비효율성을 동기부여하고 해결합니다.
• 대배치 SGD에 적합한 계층별 적응형 최적화 방법(Lamb)을 개발합니다.
• 비볼록 설정에서 Lars와 Lamb의 이론적 수렴 분석을 제공합니다.
• 대배치 구간에서 Bert 및 ResNet-50에 대해 Lamb의 우수한 성능을 실증적으로 보여줍니다.

제안 방법

• 업데이트를 단위 L2-노름으로 계층별로 정규화하고, 계층 노름의 함수에 따라 학습률을 스케일링하여 대배치 학습의 안정성을 확보하는 일반 전략을 제시합니다.
• 이 전략 하에서 두 알고리즘을 구체화합니다: Lars(모멘텀 기반)와 Lamb(계층별 정규화가 포함된 Adam 기반).
• Lars와 Lamb 모두에서 비볼록 목적함수에 대한 정지점 수렴 분석을 제공합니다.
• Bert(SQuAD, 2.5B+0.8B 말뭉치)와 ImageNet/ResNet-50에서 매우 큰 배치 크기와 최소한의 하이퍼파라미터 튜닝으로 평가합니다.
• 제곱근 학습률 스케일링과 선형 워밍업을 사용하고, 다단계 Bert 학습에서 혼합 시퀀스 길이에 대한 재-워밍업을 적용합니다.

실험 결과

연구 질문

• RQ1 계층별 업데이트 정규화와 적응 학습률 스케일링이 매우 큰 미니배치에서도 안정적이고 효율적인 학습을 가능하게 할 수 있는가?
• RQ2 Lamb와 Lars가 SGD와 유사한 비볼록 설정에서 이론적 수렴 보장을 제공하는가, 그리고 Bert와 ResNet-50과 같은 과제에서 경험적으로 어떻게 비교되는가?
• RQ3 Bert와 ResNet-50이 정확도 손실 없이 견딜 수 있는 배치 크기는 얼마이며 TPUv3 포드에서 벽시계 시간 단축은 얼마나 달성되는가?
• RQ4 Lamb가 다양한 아키텍처와 데이터 규모(언어 모델 대 비전 모델)에서 Lars보다 더 강건한가?

주요 결과

• Lamb는 Bert의 배치 크기를 최대 32K까지 안정적으로 학습시키고, SQuAD에 대한 F1을 기준선과 비교하여 유지하거나 향상시킵니다.
• 리드-배치 훈련과 재-워밍업을 사용하는 Lamb를 이용하면 Bert 학습 시간이 TPUv3 포드에서 약 3일에서 76분으로 감소합니다(자원 64x 확장).
• Lamb는 ImageNet의 ResNet-50에서 Adam/AdamW를 능가하거나 경쟁력 있는 최고-1 정확도를 제공하며, 대배치에서 모멘텀/Lars 기반보다 우수하거나 동등합니다.
• Lars는 테스트된 배치 크기에서 일반적으로 Lamb보다 Bert에서 성능이 떨어지지만, Lamb는 배치 크기가 증가해도 높은 F1 점수를 유지합니다(표 1 및 표 2).
• 이론적 결과는 비볼록 설정에서 Lars와 Lamb 모두가 정지점으로 수렴함을 보여주며, Lamb는 평균 계층 Lipschitz 상수(L_avg)에 대한 의존도가 L_infty에 비해 유리합니다.