AutoScale: Automatic Prediction of Compute-optimal Data Composition for Training LLMs

要約

さまざまなダウンストリーム タスクのパフォーマンスを確保するために、LLM はさまざまなドメインにわたるデータ混合を通じて事前トレーニングされます。
この研究では、固定のコンピューティング バジェットに対する最適なデータ構成がトレーニング データの規模に応じて変化することを実証しました。これは、小規模な実験を使用して最適な構成を経験的に決定する一般的な手法では、次のような場合に最適なデータ混合が得られないことを示唆しています。
最終モデルまでスケールアップ。
この課題に対処するために、私たちは *AutoScale* を提案します。これは、任意のターゲット スケールでのトレーニングに最適なコンピューティング データ構成を見つける自動ツールです。
AutoScale は、まず新しい 2 レベル最適化フレームワークである Direct Data Optimization (*DDO*) を使用して小規模で最適な構成を決定し、次に予測子を適合させて大規模な最適な構成を推定します。
予測子の設計は、データ構成に関連するスケーリング則の理論的分析からインスピレーションを得ており、これは独立して興味深いものになる可能性があります。
RedPajama データセットで事前トレーニングされた 774M Decoder-only LM (GPT-2 Large) を使用した実証研究では、AutoScale はどのベースラインよりも少なくとも 25% 速く検証の複雑さを減少させ、再重み付けなしと比較して最大 38% の速度向上を実現し、全体的に最高のパフォーマンスを達成しました。
下流タスク全体のパフォーマンス。
マスクされた言語モデリングを使用したエンコーダー専用 LM (BERT) の事前トレーニングでは、DDO はすべてのドメインで損失を削減し、GLUE ベンチマークでの平均タスク パフォーマンスを 8.7%、大規模 QA データセット (SQuAD) で 5.9% 向上させることが示されています。
再ウェイトなしとの比較。
AutoScale により、トレーニングが最大 28% 高速化されます。
私たちのコードはオープンソースです。

要約(オリジナル)

To ensure performance on a diverse set of downstream tasks, LLMs are pretrained via data mixtures over different domains. In this work, we demonstrate that the optimal data composition for a fixed compute budget varies depending on the scale of the training data, suggesting that the common practice of empirically determining an optimal composition using small-scale experiments will not yield the optimal data mixtures when scaling up to the final model. To address this challenge, we propose *AutoScale*, an automated tool that finds a compute-optimal data composition for training at any desired target scale. AutoScale first determines the optimal composition at a small scale using a novel bilevel optimization framework, Direct Data Optimization (*DDO*), and then fits a predictor to estimate the optimal composition at larger scales. The predictor’s design is inspired by our theoretical analysis of scaling laws related to data composition, which could be of independent interest. In empirical studies with pre-training 774M Decoder-only LMs (GPT-2 Large) on RedPajama dataset, AutoScale decreases validation perplexity at least 25% faster than any baseline with up to 38% speed up compared to without reweighting, achieving the best overall performance across downstream tasks. On pre-training Encoder-only LMs (BERT) with masked language modeling, DDO is shown to decrease loss on all domains while visibly improving average task performance on GLUE benchmark by 8.7% and on large-scale QA dataset (SQuAD) by 5.9% compared with without reweighting. AutoScale speeds up training by up to 28%. Our codes are open-sourced.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google