[論文レビュー] DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining

Sang Michael Xie, Hieu Pham+8 more|arXiv (Cornell University)|May 17, 2023

Topic Modeling被引用数 14

ひとこと要約

DoReMi は、Group DRO を用いた小さな代理モデルを用いて、事前学習データのドメイン重みを学習し、再重み付けデータで大きなLMを訓練する。これにより、タスク固有の微調整なしでトレーニングを高速化し、ダウンストリームの性能を向上させる。

ABSTRACT

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 set the 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% points 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.

研究の動機と目的

• 事前学習データのドメイン構成がダウンストリームタスク全体でのLM性能にどう影響するかを動機づける。
• ダウンストリームタスクの知識なしに、データ駆動型の手法で自動的にドメイン重みを設定する。
• 小さなモデルでのドメイン再重み付けが、はるかに大きなモデルへ転送可能であることを示す。
• 本手法がドメイン全体の perplexity と標準タスクのダウンストリーム精度を改善することを示す。

提案手法

• 初期のドメイン重みを用いて小さな参照モデルを訓練し、各ドメインの難易度のベースラインを確立する。
• Group DRO で訓練された代理モデルを使用して、基準モデルに対する各ドメインの最悪ケースの過剰損失を最小化することでドメイン重みを最適化する。
• 訓練全体で最適化されたドメイン重みを集約・平均して最終的なドメイン重みを得る。
• 最適化されたドメイン重みを用いて大規模ターゲットモデルの訓練データを再サンプリングし、標準的な訓練手順でフルサイズのモデルを訓練する。
• 任意で DoReMi を複数ラウンドで反復し、1ラウンドの調整済み重みを次ラウンドの基準として使用する。

実験結果

リサーチクエスチョン

• RQ1 Group DRO で最適化された小さな代理モデルは、後に訓練されるはるかに大きな LM の性能を改善するドメイン重みを特定できるか。
• RQ2 ダウンストリームタスクの知識なしで見つけたドメイン重みは、ドメイン横断のダウンストリーム性能に一般化するか。
• RQ3 DoReMi は個々のドメイン別の perplexity および標準の few-shot タスクにおける全体のダウンストリーム精度にどう影響するか。

主な発見

• DoReMi は The Pile において、デフォルトのドメイン重みで訓練したベースラインと比べて、8B モデルの平均的なダウンストリーム精度を 6.5 ポイント改善する。
• 最適化されたドメイン重みは、The Pile のすべてのドメインで perplexity を低下させ、たとえ一部のドメインが低重み付けされてもそうなる。
• DoReMi は The Pile でベースラインのダウンストリーム精度に、訓練ステップで2.6倍速く到達する。
• GLaM データセットでは、反復 DoReMi が、ダウンストリームタスクで調整されたドメイン重みと同等の性能を達成し、最適化にはダウンストリームデータを使用しない。
• DoReMi の利得は、代理モデルのサイズ範囲と主モデルのスケールに対して頑健である。