DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining

要約

事前トレーニング データ ドメイン (Wikipedia、書籍、Web テキストなど) の混合比率は、言語モデル (LM) のパフォーマンスに大きく影響します。
この論文では、ミニマックス最適化によるドメイン再重み付け (DoReMi) を提案します。これは、最初に、ドメイン全体にわたってグループ分布的にロバストな最適化 (グループ DRO) を使用して小規模なプロキシ モデルをトレーニングし、下流タスクの知識なしでドメインの重み付け (混合比率) を生成します。
次に、これらのドメインの重みを使用してデータセットをリサンプリングし、より大きなフルサイズのモデルをトレーニングします。
私たちの実験では、280M パラメーターのプロキシ モデルで DoReMi を使用して、8B パラメーター モデル (30 倍大きい) をより効率的にトレーニングするためのドメインの重みを見つけます。
The Pile では、DoReMi は、ドメインの重みを軽減した場合でも、すべてのドメインにわたる混乱を改善します。
DoReMi は、The Pile のデフォルトのドメイン重みを使用してトレーニングされたベースライン モデルと比較して、平均数ショット ダウンストリーム精度を 6.5% 向上させ、2.6 倍少ないトレーニング ステップでベースライン精度に到達します。
GLaM データセットでは、下流タスクの知識を持たない DoReMi は、下流タスクに合わせて調整されたドメイン重みを使用した場合のパフォーマンスにも匹敵します。

要約(オリジナル)

The mixture proportions of pretraining data domains (e.g., Wikipedia, books, web text) greatly affect language model (LM) performance. In this paper, we propose Domain Reweighting with Minimax Optimization (DoReMi), which first trains a small proxy model using group distributionally robust optimization (Group DRO) over domains to produce domain weights (mixture proportions) without knowledge of downstream tasks. We then resample a dataset with these domain weights and train a larger, full-sized model. In our experiments, we use DoReMi on a 280M-parameter proxy model to find domain weights for training an 8B-parameter model (30x larger) more efficiently. On The Pile, DoReMi improves perplexity across all domains, even when it downweights a domain. DoReMi improves average few-shot downstream accuracy by 6.5% over a baseline model trained using The Pile’s default domain weights and reaches the baseline accuracy with 2.6x fewer training steps. On the GLaM dataset, DoReMi, which has no knowledge of downstream tasks, even matches the performance of using domain weights tuned on downstream tasks.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google