R&B: Domain Regrouping and Data Mixture Balancing for Efficient Foundation Model Training
This addresses computational bottlenecks in training large language models, offering a more efficient and effective data mixing approach for AI researchers and practitioners.
The paper tackles inefficiencies in data mixing strategies for foundation model training by introducing R&B, which regroups data based on semantic similarity and balances composition using domain gradients, achieving performance matching or exceeding state-of-the-art methods with only 0.01% additional compute overhead.
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.