[논문 리뷰] Accelerated Large Batch Optimization of BERT Pretraining in 54 minutes

Shuai Zheng, Haibin Lin외 2명|arXiv (Cornell University)|2020. 06. 24.

Topic Modeling참고 문헌 33인용 수 10

한 줄 요약

이 논문은 BERT 미사전학습을 위한 가속화된 대용량 배치 최적화 방법인 LANS를 제안한다. 이 방법은 블록 별 기울기 정규화와 웜업 이후 일정 학습률 단계를 갖는 새로운 학습률 스케줄러를 결합한다. 96K 및 33K의 미니배치 크기에서 안정적인 학습을 가능하게 함으로써, LANS는 192대의 AWS EC2 P3dn.24xlarge 인스턴스를 사용해 BERT 미사전학습을 54분 만에 완료하여 정확도 손실 없이 클라우드 환경에서 새로운 속도 기록을 수립한다.

ABSTRACT

BERT has recently attracted a lot of attention in natural language understanding (NLU) and achieved state-of-the-art results in various NLU tasks. However, its success requires large deep neural networks and huge amount of data, which result in long training time and impede development progress. Using stochastic gradient methods with large mini-batch has been advocated as an efficient tool to reduce the training time. Along this line of research, LAMB is a prominent example that reduces the training time of BERT from 3 days to 76 minutes on a TPUv3 Pod. In this paper, we propose an accelerated gradient method called LANS to improve the efficiency of using large mini-batches for training. As the learning rate is theoretically upper bounded by the inverse of the Lipschitz constant of the function, one cannot always reduce the number of optimization iterations by selecting a larger learning rate. In order to use larger mini-batch size without accuracy loss, we develop a new learning rate scheduler that overcomes the difficulty of using large learning rate. Using the proposed LANS method and the learning rate scheme, we scaled up the mini-batch sizes to 96K and 33K in phases 1 and 2 of BERT pretraining, respectively. It takes 54 minutes on 192 AWS EC2 P3dn.24xlarge instances to achieve a target F1 score of 90.5 or higher on SQuAD v1.1, achieving the fastest BERT training time in the cloud.

연구 동기 및 목표

• 기존에 수일이 걸리는 BERT 미사전학습의 학습 시간을 단축하기 위해 효율적인 대용량 배치 최적화를 가능하게 하기 위해.
• 기존의 대용량 배치 최적화기인 LAMB와 같이 64K를 초과하는 미니배치 크기를 확장할 경우 발생하는 불안정성과 정확도 저하 문제를 해결하기 위해.
• 발산 없이 더 오랜 기간 동안 높은 학습률을 유지할 수 있도록 설계된 학습률 스케줄러를 통해 빠른 수렴을 가능하게 하기 위해.
• 모델 성능을 희생시키지 않고 클라우드 환경에서 BERT의 최신 기술 수준의 학습 속도를 달성하기 위해.

제안 방법

• LANS는 LAMB의 블록 단위 적응형 학습률을 확장하기 위해 블록 별 기울기 정규화를 도입하여 다양한 레이어에서 안정적인 업데이트를 보장한다.
• 클래식한 모멘텀과 네스테로프 모멘텀 간의 관계를 활용해 모멘텀 항을 수정함으로써 수렴 성능을 향상시킨다.
• 웜업 이후 일정 학습률 단계를 포함하는 새로운 학습률 스케줄을 도입하여 초기 웜업 단계 이후에도 고학습률을 오랫동안 유지한다.
• 학습률는 제곱근 규칙을 사용해 스케일링되지만, 대용량 배치에서의 발산을 방지하기 위해 조정되며, 경험적 튜닝을 통해 안정성을 유지한다.
• 데이터를 워커 간에 분할하여 그래디언트 분산을 최소화하고 중복되지 않는 샘플링을 보장함으로써 분산 학습을 지원한다.
• 클라우드 인프라에 최적화된 구현으로, 고성능 통신을 위한 EFA를 사용하는 192대의 AWS EC2 P3dn.24xlarge 인스턴스를 활용한다.

실험 결과

연구 질문

• RQ1 기존의 대용량 배치 최적화기인 LAMB와 비교해, 수정된 대용량 배치 최적화기인 LANS는 96K 및 33K의 미니배치 크기로 확장해도 모델 정확도를 유지할 수 있는가?
• RQ2 웜업 이후 일정 학습률 단계를 도입함으로써 선형 웜업 스케줄보다 학습 안정성과 수렴 속도가 향상되는가?
• RQ3 정확도 저하 없이 클라우드 기반 GPU 클러스터에서 LAMB를 초월하는 학습 속도를 달성할 수 있는가?
• RQ4 빠른 수렴을 달성하기 위해 최적의 하이퍼파라미터 설정(학습률, 웜업 기간, 일정 단계 비율)은 무엇인가?

주요 결과

• LANS는 1536장의 V100 GPU를 사용해 53.6분 만에 SQuAD v1.1 개발 세트 F1 스코어 90.60을 달성하여 클라우드에서 가장 빠른 BERT 미사전학습 기록을 수립했다.
• LANS는 BERT 미사전학습의 1단계에서 96K, 2단계에서 33K의 미니배치 크기를 성공적으로 확장했으며, 성능 저하 없이 수행했다.
• LAMB는 96K/33K 배치 크기에서 수렴하지 못했지만, LANS는 훨씬 적은 반복 수(4301회 대비 8599회)로 동일한 목표 F1 스코어를 달성했다.
• 제안된 학습률 스케줄은 학습률 0.007과 0.01 사이의 성능 격차를 AUC 기준 5.28 포인트에서 1.91 포인트로 감소시켜 더 높은 안정성을 입증했다.
• 웜업 이후 일정 학습률 단계를 도입함으로써 고학습률에서 지속적인 학습 진전이 가능해져 더 빠른 수렴을 가능하게 했다.
• 192대의 AWS EC2 P3dn.24xlarge 인스턴스를 사용해 학습 시간을 54분으로 단축했으며, TPU를 사용한 이전 방법의 76.2분보다 뛰어난 성능을 보였다.