DomainPilot: ドメインレベルの損失に基づく2段階データ混合最適化による効率的な言語モデルファインチューニング

大規模言語モデル（LLM）の学習効果は、トレーニングデータの品質と構成に根本的に制約される。既存の動的データスケジューリング手法は、産業規模の事前学習と教師ありファインチューニング（SFT）において重大な限界に直面している。データ選択はテラバイト規模のコーパスで法外なO(N)コストを発生させ、混合最適化方式は深刻なI/Oボトルネックを引き起こすか、補助的な参照モデルの学習を必要とする。サンプルレベルの再重み付け戦略は、ノイズ、困難度、新規性を混同する損失信号に依存している。

本研究では、ドメインレベルの損失に基づく2段階データ混合最適化フレームワーク「DomainPilot」を提案する。DomainPilotはトークンレベルのドメイン損失監視を導入し、データパイプラインを停止することなく各ドメインの学習ダイナミクスを捕捉する。この信号に基づき、スケーリング則に導かれた粗最適化段階でドメイン固有の収束曲線をフィッティングし、混合調整の原理的な事前分布を導出する。続くミキシング則に導かれた微最適化段階では、制御されたスイープ実験を通じてクロスドメイン相互作用効果をモデル化し、混合を洗練する。メカニズム全体はパッチベースのアーキテクチャで実現され、既存のトレーニングフレームワーク（MindSpeed/Megatron-LMなど）に約30行のフレームワーク固有アダプターコードを注入するだけで動作する。

Qwen3-1.7BモデルをSFTで検証した結果、元のデータ混合と比較して、最適化された混合は総データ量やトレーニングコストを増やすことなく、MMLU-Reduxで+2%、AIME24で+1.8%、LiveCodeBench v5で+3.8%、BFCL v3で+3.6%の改善を達成した。これらの結果は、ドメインレベルのトレーニング信号が、高価なデータ選択や補助モデル学習に代わる効果的で軽量な混合最適化手法を提供することを示している。DomainPilotのコードはarXiv:2607.22769で公開されている。

本フレームワークの重要な革新点は、2段階最適化戦略にある。第1段階ではスケーリング則を用いて各ドメインのデータ量増加に伴う学習収束挙動を予測し、合理的な初期混合比率を決定する。第2段階ではミキシング則によりドメイン間の相互影響を捕捉し、例えば数学コード領域のデータが推論能力に相乗効果をもたらすような相互作用をモデル化する。この分割統治アプローチにより、最適化プロセスは効率的かつ理論的根拠に基づくものとなる。実験結果は、1.7Bパラメータのモデルでもドメインレベルの信号が顕著な利益をもたらすことを示しており、より大規模なモデルではさらに大きな改善が期待される。