DataFlex：大規模言語モデルのデータ中心の動的トレーニングのための統一フレームワーク

要旨: データ中心の学習は、大規模言語モデル（LLM）を改善する有望な方向性として注目されており、最適化の際にモデルパラメータだけでなく、学習データの選択、構成、重み付けも同時に最適化することが鍵となっています。しかし、データ選択、データ混合の最適化、データの再重み付けに関する既存手法は、多くの場合、インターフェースが一貫していない独立したコードベースの中で開発されており、その結果として、再現性の欠如、公平な比較の困難さ、そして実運用への統合の妨げとなっています。本論文では、LLaMA-Factory を基盤とした統一的なデータ中心の動的学習フレームワークである DataFlex を提案します。DataFlex は、動的データ最適化の主要な 3 つのパラダイム、すなわちサンプル選択、ドメイン混合の調整、サンプルの再重み付けをサポートしつつ、元の学習ワークフローと完全に互換性を維持します。拡張可能なトレーナ抽象化とモジュール化されたコンポーネントを提供することで、標準的な LLM 学習のドロップイン置換を可能にし、埋め込み抽出、推論、勾配計算といった重要なモデル依存操作を、DeepSpeed ZeRO-3 を含む大規模設定にも対応しながら統一します。複数のデータ中心手法に対して包括的な実験を行いました。動的データ選択は、Mistral-7B と Llama-3.2-3B の両方において、MMLU で一貫して静的な全データ学習を上回ります。データ混合については、Qwen2.5-1.5B を SlimPajama 上で 6B および 30B トークン規模で事前学習する際、DoReMi と ODM により、デフォルトの割合に比べて MMLU の精度とコーパスレベルのパープレキシティの両方が改善されます。さらに DataFlex は、元の実装に対して一貫した実行時の改善も達成しています。これらの結果は、DataFlex が LLM のデータ中心動的学習に対して、有効で効率的かつ再現可能な基盤インフラを提供することを示しています。