R&B: 効率的な基盤モデル訓練のためのドメイン再編成とデータ混合バランス調整

R&B: Domain Regrouping and Data Mixture Balancing for Efficient Foundation Model Training

May 1, 2025

著者: Albert Ge, Tzu-Heng Huang, John Cooper, Avi Trost, Ziyi Chu, Satya Sai Srinath Namburi GNVV, Ziyang Cai, Kendall Park, Nicholas Roberts, Frederic Sala

cs.AI

要旨

データ混合戦略は、言語モデルのトレーニングに伴うコストを削減することに成功してきました。しかし、このような手法には2つの欠点があります。第一に、事前に定義されたデータドメイン（例：データソース、タスクタイプ）に依存しており、重要な意味的ニュアンスを捉え損ねる可能性があり、性能を十分に引き出せないことがあります。第二に、これらの手法はドメインの数に応じて計算量が指数的に増加するため、計算コストが高くなります。これらの課題に対処するため、我々はR&Bというフレームワークを提案します。R&Bは、意味的類似性に基づいてトレーニングデータを再分割（Regroup）し、より細かい粒度のドメインを作成し、トレーニング全体で得られるドメイン勾配によって誘導されるグラム行列を活用してデータ構成を効率的に最適化（Balance）します。従来の手法とは異なり、損失や勾配などの評価情報を取得するための追加の計算を必要としません。我々は、標準的な正則条件の下でこの技術を分析し、非適応的な混合手法と比較してR&Bの有効性を理論的に裏付けます。実験的には、自然言語から推論、マルチモーダルタスクに至る5つの多様なデータセットにおいてR&Bの有効性を実証します。わずか0.01%の追加計算オーバーヘッドで、R&Bは最先端のデータ混合戦略の性能に匹敵するか、それを上回る結果を示します。

English

Data mixing strategies have successfully reduced the costs involved in training language models. While promising, such methods suffer from two flaws. First, they rely on predetermined data domains (e.g., data sources, task types), which may fail to capture critical semantic nuances, leaving performance on the table. Second, these methods scale with the number of domains in a computationally prohibitive way. We address these challenges via R&B, a framework that re-partitions training data based on semantic similarity (Regroup) to create finer-grained domains, and efficiently optimizes the data composition (Balance) by leveraging a Gram matrix induced by domain gradients obtained throughout training. Unlike prior works, it removes the need for additional compute to obtain evaluation information such as losses or gradients. We analyze this technique under standard regularity conditions and provide theoretical insights that justify R&B's effectiveness compared to non-adaptive mixing approaches. Empirically, we demonstrate the effectiveness of R&B on five diverse datasets ranging from natural language to reasoning and multimodal tasks. With as little as 0.01% additional compute overhead, R&B matches or exceeds the performance of state-of-the-art data mixing strategies.