MosaicBERT: A Bidirectional Encoder Optimized for Fast Pretraining

要約

BERT スタイルのエンコーダー モデルは NLP 研究で頻繁に使用されていますが、多くの研究者は、トレーニングのコストが高いため、独自の BERT を最初から事前トレーニングすることはありません。
BERT が最初に注目を集めてから過去 5 年の間に、他のトランスフォーマー アーキテクチャやトレーニング構成に関して多くの進歩が見られましたが、これらはまだ BERT に体系的に組み込まれていません。
ここでは、高速事前トレーニング用に経験的に最適化された BERT スタイルのエンコーダー アーキテクチャとトレーニング レシピである MosaicBERT を紹介します。
この効率的なアーキテクチャには、FlashAttendant、Attendee with Linear Biases (ALiBi)、Gated Linear Units (GLU)、パディングされたトークンを動的に削除するモジュール、および低精度の LayerNorm がクラシックなトランスフォーマー エンコーダー ブロックに組み込まれています。
トレーニング レシピには、RoBERTa およびその他のエンコーダー モデルのベスト プラクティスに加えて、Masked Language Modeling (MLM) 目標の 30% のマスキング率、bfloat16 精度、GPU スループットに最適化された語彙サイズが含まれています。
C4 データセットで最初から事前トレーニングした場合、この基本モデルは、約 20 ドルのコストで、8 台の A100 80 GB GPU で 1.13 時間で 79.6 のダウンストリーム平均 GLUE (開発) スコアを達成しました。
広範な精度と事前トレーニング速度のパレート曲線をプロットし、競合する BERT のベースおよびラージと比較した場合、MosaicBERT のベースおよびラージが一貫してパレート最適であることを示します。
この経験的な事前トレーニングの高速化により、研究者やエンジニアは、既存の汎用モデルを微調整する代わりに、カスタム BERT スタイルのモデルを低コストで事前トレーニングできるようになります。
私たちはモデルの重みとコードをオープンソースにしています。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google