대규모 배치 최적화를 통한 딥 러닝: 76분 만에 BERT 훈련하기

대규모 데이터셋에서 큰 깊은 신경망을 훈련시키는 것은 계산적으로 매우 어려운 문제입니다. 최근에는 이 문제를 해결하기 위해 대형 배치 확률적 최적화 방법의 사용에 대한 관심이 급증하고 있습니다. 이 연구 분야에서 가장 주목받는 알고리즘은 LARS로, 층별 적응 학습률을 사용하여 ImageNet에서 ResNet을 몇 분 만에 훈련시킬 수 있습니다. 그러나 LARS는 BERT와 같은 주의 모델에서는 성능이 저조하여, 그 성능 개선이 작업 간 일관성이 없다는 것을 시사합니다. 본 논문에서는 먼저 대형 미니배치를 사용하여 깊은 신경망의 훈련 속도를 가속화하는 원칙적인 층별 적응 전략을 연구합니다. 이 전략을 바탕으로 새로운 층별 적응 대형 배치 최적화 기법인 LAMB를 개발하였습니다. 또한 LAMB와 LARS의 수렴성을 분석하여 일반 비볼록 설정에서 정상점으로의 수렴을 보여주었습니다. 우리의 경험적 결과는 BERT와 ResNet-50 훈련과 같은 다양한 작업에서 매우 적은 하이퍼파라미터 조정으로 LAMB의 우수한 성능을 입증합니다. 특히 BERT 훈련에서, 우리의 최적화기는 성능 저하 없이 32868와 같은 매우 큰 배치 크기를 사용할 수 있게 합니다. TPUv3 Pod의 메모리 한계까지 배치 크기를 늘림으로써, BERT 훈련 시간은 3일에서 단 76분으로 줄일 수 있습니다 (표 1 참조). LAMB 구현체는 https://github.com/tensorflow/addons/blob/master/tensorflow_addons/optimizers/lamb.py 에서 확인할 수 있습니다.