DoGE: Domain Reweighting with Generalization Estimation

要約

事前学習データのカバレッジと構成は、大規模言語モデル（LLM）の汎化能力に大きく影響する。その重要性にもかかわらず、最近のLLMは、データ領域の影響を増減するために、ヒューリスティックや試行錯誤に頼っている。我々は、各ドメインからのサンプリング確率（ドメイン重み）を原理的な方法で最適化する、汎化推定によるDOmain reweighting (DoGE)を提案する。我々のアプローチは、(i)2レベル最適化アルゴリズムを用いて、ドメイン重みを得るために代理モデルを訓練すること、(ii)学習したドメイン重みに従って訓練ドメインをサンプリングすることにより、より大きなベースモデルを訓練すること、からなる2段階のプロセスである。我々の実験では、DoGEがどのようなターゲットデータ混合に対するベースモデルの汎化性をどのように向上させるかを広範囲に示す。スリムパジャマ（SlimPajama）データセットにおいて、我々の基本モデルはベースライン手法と比較して、$6$のタスクにわたって、より良い当惑度（perplexity）と少数ショット推論精度（few-shot reasoning accuracy）を得る。さらに、事前学習コーパス（OODドメイン）では未見である、ドメイン外のターゲットタスクへの汎化を目指すことで、DoGEはドメイン間依存性を効果的に同定することができ、一貫してターゲットドメイン上でより良いテスト当惑度を達成する。

要約(オリジナル)

The coverage and composition of the pretraining data significantly impacts the generalization ability of Large Language Models (LLMs). Despite its importance, recent LLMs still rely on heuristics and trial and error to increase or reduce the influence of data-domains. We propose DOmain reweighting with Generalization Estimation (DoGE), which optimizes the probability of sampling from each domain (domain weights) in a principled way. Our approach is a two-stage process consisting of (i) training a proxy model to obtain domain weights using a bi-level optimization algorithm; (ii) training a larger base model by sampling training domains according to the learned domain weights. In our experiments, we extensively show how DoGE improves the generalization of the base model to any target data mixture. On the SlimPajama dataset, our base model gets better perplexity and few-shot reasoning accuracies across $6$ tasks compared to baseline methods. Moreover, aiming to generalize to out-of-domain target tasks, which is unseen in the pretraining corpus (OOD domain), DoGE can effectively identify inter-domain dependencies, and consistently achieves better test perplexity on the target domain.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL