[논문 리뷰] Reducing BERT Pre-Training Time from 3 Days to 76 Minutes

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

Advanced Neural Network Applications인용 수 73

한 줄 요약

이 논문은 TPUv3 풀에서 최대 32,868까지의 매우 큰 배치 크기를 허용함으로써 BERT 사전학습을 3일에서 단 76분으로 가속화하는 새로운 대용량 배치 최적화 방법인 LAMB(Layer-wise Adaptive Moments)를 소개한다. LAMB는 계층별 적응형 학습률과 공식적인 수렴 보장을 통해 기존 방법들보다 BERT와 ResNet-50 모두에서 뛰어난 성능을 달성한다.

ABSTRACT

Training large deep neural networks on massive datasets is very challenging. One promising approach to tackle this issue is through the use of large batch stochastic optimization. However, our understanding of this approach in the context of deep learning is still very limited. Furthermore, the current approaches in this direction are heavily hand-tuned. To this end, we first study a general adaptation strategy to accelerate training of deep neural networks using large minibatches. Using this strategy, we develop a new layer-wise adaptive large batch optimization technique called LAMB. We also provide a formal convergence analysis of LAMB as well as the previous published layerwise optimizer LARS, showing convergence to a stationary point in general nonconvex settings. Our empirical results demonstrate the superior performance of LAMB for BERT and ResNet-50 training. In particular, for BERT training, our optimization technique enables use of very large batches sizes of 32868; thereby, requiring just 8599 iterations to train (as opposed to 1 million iterations in the original paper). By increasing the batch size to the memory limit of a TPUv3 pod, BERT training time can be reduced from 3 days to 76 minutes. Finally, we also demonstrate that LAMB outperforms previous large-batch training algorithms for ResNet-50 on ImageNet; obtaining state-of-the-art performance in just a few minutes.

연구 동기 및 목표

• BERT와 같이 대규모 컴퓨팅 자원과 긴 학습 시간을 요구하는 대규모 딥 뉴럴 네트워크의 느린 학습 시간 문제를 해결하기 위해.
• 딥 러닝에서 아직 잘 이해되지 않았고 실제로는 수작업으로 다수 조정되는 대용량 배치 확률적 최적화의 효율성을 향상시키기 위해.
• 매우 큰 미니배치로 안정적이고 빠른 학습을 가능하게 하는 일반화 가능하고 적응형 최적화 전략을 개발하기 위해.
• 비볼록 설정에서 LAMB와 이전의 LARS 최적화기의 공식적인 수렴 분석을 제공하여 이론적 안정성을 확보하기 위해.
• 훈련 시간과 반복 횟수를 크게 줄여 상태 최고 성능을 달성하는 BERT와 ResNet-50에서의 성능을 입증하기 위해.

제안 방법

• 각 계층의 기울기 노름 대 비율에 따라 학습률을 조정함으로써 대용량 배치 최적화 기법인 LAMB를 제안한다.
• Adam 최적화기를 계층별 적응형 학습률을 도입하여 대용량 배치 학습에서의 안정성과 수렴성을 향상시킨다.
• 기울기와 파라미터 크기의 균형을 유지함으로써 학습을 안정화시키는 정규화 메커니즘을 도입한다.
• LAMB와 LARS 모두 일반적인 비볼록 설정에서 정류점으로의 수렴을 보장하는 공식적인 수렴 분석을 수행한다.
• TPUv3 풀의 메모리 용량을 활용하여 배치 크기를 최대 32,868까지 확장함으로써 훈련 반복 횟수를 극적으로 감소시킨다.
• 매우 큰 배치 크기에서도 모델 안정성을 유지할 수 있도록 학습률 스케줄링 전략을 사용하여 기존의 대용량 배치 방법에서 흔히 발생하는 발산을 방지한다.

실험 결과

연구 질문

• RQ1 대용량 배치 학습을 위한 일반적인 적응 전략이 모델 품질을 손상시키지 않고 BERT의 사전학습 시간을 크게 단축시킬 수 있는가?
• RQ2 계층별 적응형 학습률은 딥 네트워크의 대용량 배치 설정에서 최적화 안정성과 수렴성을 어떻게 향상시킬 수 있는가?
• RQ3 LAMB는 이전의 LARS 최적화기와 유사하게 비볼록 최적화에서 이론적 수렴 보장을 달성하는가?
• RQ4 LAMB는 현대 가속기인 TPUv3 풀의 메모리 한계에 가까운 배치 크기로 BERT를 학습시킬 수 있는가?
• RQ5 ImageNet과 GLUE와 같은 표준 벤치마크에서 LAMB는 기존의 대용량 배치 최적화 방법과 비교해 정확도와 학습 속도 측면에서 어떻게 성과를 내는가?

주요 결과

• LAMB는 TPUv3 풀에서 최대 배치 크기 32,868를 허용함으로써 BERT 사전학습 시간을 3일에서 단 76분으로 단축시켰다.
• LAMB를 사용한 BERT 학습은 원래 BERT 논문의 100만 번의 반복 대비 단 8,599회의 반복으로 완료되었다.
• LAMB는 몇 분 내로 ResNet-50/ImageNet에서 최고 성능을 달성하며 이전의 대용량 배치 방법들을 능가했다.
• 매우 큰 배치 크기에서도 안정적인 학습을 보였으며, 일반적인 대용량 배치 최적화에서 흔히 발생하는 발산 문제를 피했다.
• 공식적인 수렴 분석을 통해 LAMB가 일반적인 비볼록 설정에서 정류점으로 수렴한다는 것이 확인되었으며, 이는 이론적 기반을 제공한다.
• LAMB는 BERT와 ResNet-50 모두에서 이전의 대용량 배치 학습 알고리즘보다 뛰어난 성능을 보이며 빠른 속도와 정확도 향상을 지속적으로 달성했다.