[論文レビュー] MosaicBERT: A Bidirectional Encoder Optimized for Fast Pretraining

Jacob Portes, Alex Trott+7 more|arXiv (Cornell University)|Dec 29, 2023

Topic Modeling被引用数 4

ひとこと要約

MosaicBERT は、FlashAttention、ALiBi、ゲート付き線形ユニット（GLU）、動的パディング解除、低精度の LayerNorm を統合することで、高速な事前学習を最適化した BERT スタイルの双方向エンコーダーである。8 枚の A100 80GB GPU を用いた 1.13 時間で、コスト約 20 ドルで 79.6 の平均 GLUE スコアを達成し、標準的な BERT モデルと比較して精度と学習時間の両面で Pareto 最適性を示している。

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 アーキテクチャに体系的に統合すること。
• 複数のモデルサイズにおいて、精度と事前学習速度の両面で Pareto 改善を実証的に示すこと。
• 研究者やエンジニアが汎用の事前学習済みモデルに依存せず、低コストでカスタム BERT スタイルのモデルを微調整できるようにすること。
• モデル重みとコードをオープンソース化して、再現性の向上と広範な採用を促進すること。

提案手法

• 最適化されたカーネル演算を用いた FlashAttention を統合し、自己注意計算を高速化する。
• 位置埋め込みを用いずに相対的位置エンコーディングを可能にする Attention with Linear Biases (ALiBi) を適用する。
• 表現能力の向上を図るために、融合実装を用いた Gated Linear Units (GLU) を前向き層に使用する。
• パディング済みトークンの計算をスキップする動的パディングモジュールを統合し、可変長シーケンスにおけるスループットを向上させる。
• メモリ帯域幅と計算コストを削減するために、低精度の LayerNorm（bfloat16）を採用する。
• マスク言語モデル（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 と比較して Pareto 最適か？
• RQ4 個々のアーキテクチャ的要素（例：FlashAttention、GLU、パディング解除）は、スループットと収束速度にどのように寄与しているか？
• RQ5 BERT-Base モデルを BERT-Large の性能水準まで引き上げるには著しく長い学習が必要であり、この問題はアーキテクチャ的改善によって軽減可能か？

主な発見

• MosaicBERT-Base は、8 枚の A100 80GB GPU を用いて 1.13 時間で、標準的なクラウドインfraストラクチャ上でのコスト約 20 ドルで、下流タスクの平均 GLUE（開発セット）スコア 79.6 を達成した。
• 8× A100 80GB GPU を用いた場合、30% の MLM マスキング率と bfloat16 精度を用いることで、学習コストは 22.60 ドルに達した。
• MosaicBERT-Base および MosaicBERT-Large は、複数の事前学習時間において、標準的な BERT-Base および BERT-Large と比較して、実証的に Pareto 最適である。
• スループットのアブレーションスタディから、FlashAttention、ALiBi、動的パディング解除はいずれも学習速度の向上に顕著な寄与をしていることが示された。
• Fused 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