[논문 리뷰] Progressively Stacking 2.0: A Multi-stage Layerwise Training Method for BERT Training Speedup

Cheng Yang, Shengnan Wang외 4명|arXiv (Cornell University)|2020. 11. 27.

Topic Modeling참고 문헌 22인용 수 10

한 줄 요약

이 논문은 BERT 사전학습을 가속화하기 위해 점진적으로 얕은 모델을 쌓고 각 단계에서 상단에 있는 레이어들만 미세조정하면서 이전에 학습된 레이어들은 동결하는 다단계 계층별 훈련(MSLT) 방법을 제안한다. 이 방법은 역전파 계산과 통신 오버헤드를 줄여주며, 표준 BERT 훈련과 비교해 성능 저하가 최소한이면서 110% 이상의 훈련 속도 향상을 달성한다.

ABSTRACT

Pre-trained language models, such as BERT, have achieved significant accuracy gain in many natural language processing tasks. Despite its effectiveness, the huge number of parameters makes training a BERT model computationally very challenging. In this paper, we propose an efficient multi-stage layerwise training (MSLT) approach to reduce the training time of BERT. We decompose the whole training process into several stages. The training is started from a small model with only a few encoder layers and we gradually increase the depth of the model by adding new encoder layers. At each stage, we only train the top (near the output layer) few encoder layers which are newly added. The parameters of the other layers which have been trained in the previous stages will not be updated in the current stage. In BERT training, the backward computation is much more time-consuming than the forward computation, especially in the distributed training setting in which the backward computation time further includes the communication time for gradient synchronization. In the proposed training strategy, only top few layers participate in backward computation, while most layers only participate in forward computation. Hence both the computation and communication efficiencies are greatly improved. Experimental results show that the proposed method can achieve more than 110% training speedup without significant performance degradation.

연구 동기 및 목표

• 대규모 BERT 모델의 높은 계산 비용과 긴 훈련 시간 문제를 해결하기 위해.
• 분산 환경에서 역전파 및 통신 오버헤드를 줄여 훈련 효율을 향상시키기 위해.
• 기존 점진적 스택킹 방법의 비효율성을 해결하기 위해, 후반 단계에서 전체 모델 업데이트로 인해 속도가 저하되는 문제를 해결하기 위해.
• 하부 레이어가 조기에 안정화됨을 관찰하여 지속적인 재학습이 불필요하다는 점을 활용하기 위해.
• 훈련 시간을 크게 줄이면서도 모델 성능를 유지하는 훈련 전략을 개발하기 위해.

제안 방법

• 이 방법은 BERT 훈련을 여러 단계로 분해하며, 처음에는 몇 개의 인코더 레이어만 포함하는 얕은 모델로 시작한다.
• 각 단계에서 이전에 학습된 레이어들은 모두 동결되어 전방전파에만 참여하고, 새로 추가된 상단 레이어들만 업데이트된다.
• 모델 깊이가 단계별로 점진적으로 증가하며, 각 단계에서 동일한 사전학습 목표를 유지한다.
• 모든 레이어가 추가된 후, 최종 재학습 단계에서 모든 파라미터를 함께 몇 단계 동안 업데이트하여 모델을 정밀하게 보정한다.
• 이 방법은 교차 최적화 원리를 활용하여, 한 번에 작은 파라미터 부분집합만 업데이트함으로써 수렴 효율을 향상시킨다.
• 학습률 충돌을 방지하기 위해, 처음에는 미학습된 상단 레이어들을 먼저 훈련하고, 이후에 작은 학습률로 모든 레이어를 함께 재학습한다.

실험 결과

연구 질문

• RQ1 깊은 BERT 모델에서 상단 레이어들만 업데이트함으로써 훈련 효율을 크게 향상시킬 수 있는가?
• RQ2 후반 단계에서 이전에 학습된 레이어들을 동결함으로써 성능 저하 없이 훈련 시간을 단축시킬 수 있는가?
• RQ3 속도와 정확도 측면에서 점진적 스택킹 전략은 표준 훈련 방식과 비교해 어떻게 성능을 내는가?
• RQ4 하부 레이어의 어텐션 분포가 조기에 안정화됨을 관찰한 것을 바탕으로 훈련을 가속화할 수 있는가?
• RQ5 다단계 계층별 훈련에서 속도 향상과 성능 저하 사이의 최적의 트레이드오프는 무엇인가?

주요 결과

• 제안된 MSLT 방법은 동일한 훈련 스텝 수로 표준 BERT 훈련 대비 110% 이상의 훈련 속도 향상을 달성한다.
• 동일한 훈련 시간 동안 MSLT는 표준 훈련보다 더 높은 성능를 기록하여, 더 뛰어난 샘플 효율성을 보여준다.
• 모델 성능는 표준 BERT와 유사하며, GLUE 벤치마크에서 미미한 성능 저하만 발생한다 (예: BERT-Base: 평균 점수 79.5 vs. 79.8).
• 최종 재학습 단계는 총 훈련 스텝의 약 20%에 불과하며, 이는 이 단계 이전에 모델이 이미 잘 최적화되어 있음을 시사한다.
• 각 단계에서 업데이트되는 파라미터 수가 제한되어 있어, 특히 분산 훈련 환경에서 역전파 계산과 기울기 동기화 오버헤드를 크게 줄인다.
• 실험 결과, 하부 레이어의 어텐션 분포가 후반 단계에서도 크게 변화하지 않음을 확인하여, 이러한 레이어들을 동결하는 데 기반한 핵심 가정이 타당함을 입증한다.