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

TL;DR

R&B introduces a novel framework for dynamic data regrouping and balancing based on semantic similarity and domain gradients, significantly improving foundation model training efficiency and performance without extra compute.

Contribution

It proposes a new data mixing method that adaptively re-partitions and balances training data based on semantic and gradient information, reducing computational costs.

Findings

R&B matches or exceeds state-of-the-art data mixing strategies.

Effective across diverse datasets including language, reasoning, and multimodal tasks.

Achieves these results with only 0.01% additional compute overhead.

Abstract

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.