AC-ODM: アクター・クリティックを用いたオンラインデータ混合によるサンプル効率的なLLM事前学習

要旨

事前学習データの構成最適化はLLMの汎化にとって極めて重要である。動的ミキシングは進化する学習ダイナミクスを捉えることで静的戦略よりも優れているが、現在の手法では計算効率とサンプル効率、多様なパイプラインに対する構造的柔軟性の両立ができていない。我々はActor-Critic Online Data Mixing（AC-ODM）を提案する。これは強化学習の視点からデータミキシングにアプローチし、理論的に勾配の建設的干渉を最大化する動的線形代理として機能することを証明したパラメータ化ポリシーを用いる。実用的な柔軟性を高めるため、AC-ODMは2つの動作モードをサポートする。（i）固定済みの事前準備されたコーパスに対するプロキシモードでは、小さなモデルで学習したポリシーをより大きなターゲットに転送する。（ii）事前知識なしでスクラッチから直接エンドツーエンド学習する非プロキシモードである。実験的に、AC-ODMは様々なアーキテクチャにおいて収束速度と下流タスク精度で従来手法を大きく上回る。Pythia-1Bでは、競合ベースラインと比較して最大66%少ない訓練ステップで最適な検証困惑度に到達し、MMLU精度で27.5%の相対改善、HumanEvalで2.23倍のpass@1を達成する。しかも、1ステップあたりのウォールクロック増加は事実上無視できる0.4%、メモリオーバーヘッドはわずか2%である。コードは https://github.com/DANG-ai/AC-ODM で入手できる。

English

Optimizing pretraining data composition is pivotal for LLM generalization. While dynamic mixing outperforms static strategies by capturing evolving training dynamics, current methods fail to reconcile computational efficiency with sample efficiency and structural flexibility for diverse pipelines.We introduce Actor--Critic Online Data Mixing (AC-ODM), which approaches data mixing from a reinforcement learning perspective with a parameterized policy that we theoretically prove to act as a dynamic linear surrogate maximizing the constructive interference of gradients. To enhance practical flexibility, AC-ODM supports two operational modes: (i) a proxy mode for fixed, pre-prepared corpora, where a policy learned on a small model is transferred to a larger target; and (ii) a non-proxy mode for direct end-to-end training from scratch without priors. Empirically, AC-ODM significantly outperforms prior methods in convergence speed and downstream accuracy across various architectures. On Pythia-1B, it reaches optimal validation perplexity using up to 66% fewer training steps than competitive baselines, delivering a 27.5% relative improvement in MMLU accuracy and a 2.23 x higher pass@1 on HumanEval, all while incurring a virtually negligible (0.4%) per-step wall-clock increase and only 2% additional memory overhead. Code is available at https://github.com/DANG-ai/AC-ODM.