[논문 리뷰] MosaicBERT: A Bidirectional Encoder Optimized for Fast Pretraining

Jacob Portes, Alex Trott외 7명|arXiv (Cornell University)|2023. 12. 29.

Topic Modeling인용 수 4

한 줄 요약

MosaicBERT는 FlashAttention, ALiBi, 게이팅형 선형 유닛(Gated Linear Units, GLU), 동적 패딩 해제(dynamic unpadding), 그리고 저정밀도 레이어 정규화(low-precision LayerNorm)를 통합하여 빠른 사전 훈련을 최적화한 BERT 스타일의 양방향 인코더이다. 8개의 A100 80GB GPU에서 약 $20의 비용으로 1.13시간 만에 79.6의 평균 GLUE 점수를 기록하여 표준 BERT 모델 대비 정확도 대비 훈련 시간에서 파레토 최적성을 입증하였다.

ABSTRACT

Although BERT-style encoder models are heavily used in NLP research, many researchers do not pretrain their own BERTs from scratch due to the high cost of training. In the past half-decade since BERT first rose to prominence, many advances have been made with other transformer architectures and training configurations that have yet to be systematically incorporated into BERT. Here, we introduce MosaicBERT, a BERT-style encoder architecture and training recipe that is empirically optimized for fast pretraining. This efficient architecture incorporates FlashAttention, Attention with Linear Biases (ALiBi), Gated Linear Units (GLU), a module to dynamically remove padded tokens, and low precision LayerNorm into the classic transformer encoder block. The training recipe includes a 30% masking ratio for the Masked Language Modeling (MLM) objective, bfloat16 precision, and vocabulary size optimized for GPU throughput, in addition to best-practices from RoBERTa and other encoder models. When pretrained from scratch on the C4 dataset, this base model achieves a downstream average GLUE (dev) score of 79.6 in 1.13 hours on 8 A100 80 GB GPUs at a cost of roughly $20. We plot extensive accuracy vs. pretraining speed Pareto curves and show that MosaicBERT base and large are consistently Pareto optimal when compared to a competitive BERT base and large. This empirical speed up in pretraining enables researchers and engineers to pretrain custom BERT-style models at low cost instead of finetune on existing generic models. We open source our model weights and code.

연구 동기 및 목표

• 기본적으로 BERT 스타일 모델을 새로 훈련할 때 발생하는 높은 계산 비용과 긴 훈련 시간 문제를 해결하기 위해.
• 최근의 트랜스포머 기술 발전—예를 들어 FlashAttention, ALiBi, GLU—을 고전적인 BERT 아키텍처에 체계적으로 통합하여 효율성을 향상시키기 위해.
• 다양한 모델 크기에서 정확도 대비 사전 훈련 속도 향상 측면에서 파레토 개선을 실증적으로 입증하기 위해.
• 연구자들과 엔지니어들이 일반적인 사전 훈련된 모델에 의존하는 대신, 저비용으로 맞춤형 BERT 스타일 모델을 미세조정할 수 있도록 지원하기 위해.
• 재현 가능성과 광범위한 채택을 촉진하기 위해 모델 가중치와 코드를 오픈소스로 공개하기 위해.

제안 방법

• 최적화된 커널 연산을 통해 자기주의 계산을 가속화하기 위해 FlashAttention을 통합한다.
• 위치 임베딩 없이도 상대적 위치 인코딩을 가능하게 하기 위해 Attention with Linear Biases(ALiBi)를 적용한다.
• 표현 능력을 향상시키기 위해 피드포워드 레이어에서 게이팅형 선형 유닛(Gated Linear Units, GLU)을 사용하고, 융합된 구현 방식을 적용한다.
• 패딩 토큰에 대한 계산을 건너뛰기 위해 동적 패딩 해제(dynamic unpadding) 모듈을 통합하여 가변 길이 시퀀스에서의 처리량을 향상시킨다.
• 메모리 대역폭과 계산 비용을 줄이기 위해 저정밀도 레이어 정규화(bfloat16)를 사용한다.
• 마스크된 언어 모델링(Masked Language Modeling, MLM) 목적함수에 대해 30%의 마스킹 비율을 적용하고, GPU 처리량을 최적화하기 위해 어휘 크기를 조정한다.

Figure 1: (A) Schematic of MosaicBERT architecture (B) Pareto curves of average GLUE (dev) scores for MosaicBERT-Base and the standard BERT-Base. Error bars indicate 95% confidence interval over n=5 pretraining seeds. All training was on $8\times$ A100-80GB GPUs. FlashAttention schematic adapted fro

실험 결과

연구 질문

• RQ1 최근 트랜스포머 변종에서 유래한 아키텍처 혁신들이 고전적인 BERT 아키텍처에 효과적으로 통합되어 정확도를 유지하면서도 사전 훈련 속도를 향상시킬 수 있는가?
• RQ2 BERT 스타일 모델에 여러 효율성 기법을 조합할 경우 정확도와 사전 훈련 속도 사이의 상충 관계는 어떻게 되는가?
• RQ3 다양한 사전 훈련 지속 시간 동안 MosaicBERT는 표준 BERT-Base 및 BERT-Large 모델 대비 파레토 최적인가?
• RQ4 개별 아키텍처 구성 요소(예: FlashAttention, GLU, 패딩 해제)는 처리량과 수렴 속도에 어떻게 기여하는가?
• RQ5 BERT-Base 모델을 BERT-Large 성능 수준까지 향상시키기 위해서는 상당히 긴 훈련 시간이 필요한가? 이 문제는 아키텍처 개선을 통해 완화될 수 있는가?

주요 결과

• MosaicBERT-Base는 8개의 A100 80GB GPU에서 1.13시간 만에 79.6의 하류 평균 GLUE(검증) 점수를 기록했으며, 표준 클라우드 인프라에서 약 $20의 비용이 들었다.
• 8× A100 80GB GPU를 사용하여 $22.60의 훈련 비용으로, 30%의 MLM 마스킹 비율과 bfloat16 정밀도를 적용한 상태에서 이 성능에 도달했다.
• 여러 사전 훈련 지속 시간 동안 MosaicBERT-Base 및 MosaicBERT-Large는 표준 BERT-Base 및 BERT-Large 모델 대비 실증적으로 파레토 최적성을 확보했다.
• 처리량 분석 실험 결과, FlashAttention, ALiBi, 동적 패딩 해제 각각이 훈련 속도 향상에 크게 기여하는 것으로 나타났다.
• 융합된 GLU 구현은 일부 환경에서는 표준 GLU보다 약간 느리지만, 표현 능력 향상과 고배치 크기 훈련과의 호환성 덕분에 사용되었다.
• BERT-Large는 장기간의 훈련 이후에야 BERT-Base를 초월하는 성능을 보이며, 아키텍처 최적화의 효율성 향상 효과를 입증한다.

Figure 2: Performance on individual GLUE (dev) finetuning tasks. Our MosaicBERT-Base consistently outperforms BERT-Base on MNLI-m, QNLI, QQP and RTE, and has comparable performance on CoLA, SST-2, MRPC and STSB. Wall clock time is for $8\times$ A100-80GB GPUs, and does not include finetuning time. E