Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

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

TL;DR

R&B addresses the inefficiency of data mixing under coarse, predefined domains by introducing semantic regrouping of data into finer, semantically coherent clusters and a gradient-driven online balancing mechanism. It formulates a bilevel optimization that first partitions data into $m$ domains and then learns time-varying proportions ${\bm{p}}^t \in \Delta^{m-1}$ using a Gram matrix of domain gradients, $G$, to maximize gradient-aligned loss reduction with minimal extra evaluation compute. The approach, supported by theoretical insights and extensive experiments across natural language, reasoning, and multimodal tasks, achieves or surpasses state-of-the-art data-mixing performance with as little as $0.01\%$ additional compute and scales well to large numbers of domains. The results demonstrate substantial efficiency gains and robustness, enabling data-efficient foundation-model training in settings with abundant unstructured data and limited compute.

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.

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

TL;DR

R&B addresses the inefficiency of data mixing under coarse, predefined domains by introducing semantic regrouping of data into finer, semantically coherent clusters and a gradient-driven online balancing mechanism. It formulates a bilevel optimization that first partitions data into domains and then learns time-varying proportions using a Gram matrix of domain gradients, , to maximize gradient-aligned loss reduction with minimal extra evaluation compute. The approach, supported by theoretical insights and extensive experiments across natural language, reasoning, and multimodal tasks, achieves or surpasses state-of-the-art data-mixing performance with as little as additional compute and scales well to large numbers of domains. The results demonstrate substantial efficiency gains and robustness, enabling data-efficient foundation-model training in settings with abundant unstructured data and limited compute.

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.
Paper Structure (31 sections, 5 theorems, 31 equations, 6 figures, 8 tables, 3 algorithms)

This paper contains 31 sections, 5 theorems, 31 equations, 6 figures, 8 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. R&B: Regroup and Balance Data Mixtures
  4. Problem Setup
  5. Defining Domains
  6. Proposed Method
  7. Determining Optimal Proportions
  8. Experiments
  9. Data Mixing on Natural Language Tasks
  10. Beyond Language Modeling
  11. Conclusion
  12. Acknowledgements
  13. Notation
  14. Theoretical Results
  15. Method Derivation
  16. ...and 16 more sections

Key Result

Lemma 1

Define the regret $R_S(i, j)$ under the skill-assigning function $S$ for class $j$ as the difference between the gradient alignments: Let $i, j \in |m|$ and assume $|D_i| = |D_j|$ and $\nabla {\mathcal{L}}(\theta_{t}; D_{i})^\top \nabla {\mathcal{L}}(\theta_{t}; {\mathcal{D}}_{{\bm{p}}}) \geq$$\nabla {\mathcal{L}}(\theta_{t}; D_{j})^\top \nabla {\mathcal{L}}(\theta_{t}; {\mathcal{D}}_{{\bm{p}}})$

Figures (6)

  • Figure 1: Instead of using pre-determined domains (e.g., by task type), we find that it is often better to first repartition the data into finer-grained, semantically related domains. Optimizing the proportions of these new semantic domains can significantly improve training performance.
  • Figure 2: Top Row: Across various data settings, we find that there is a "sweet spot" in the number of domains used for data mixing, indicated by the green star. The optimal number of groups varies significantly with the dataset, which motivates the need for compute-efficient data mixing. Bottom Row: We find that silhouette score often correlates with model performance, suggesting that it is possible to predict data mixing performance based on clustering metrics.
  • Figure 3: Left: Regrouping skills before applying data mixing strategies can yield substantial improvements. Underlined values indicate where regrouping beats the original grouping for that method and dataset. Highlighted values (with brown background) indicate the best overall performance for each dataset. Note that we do not apply Balance to the original categorization of Sup-NatInst test, as we assume that training data and validation data are bucketed into the same $m$ groups. Right: Loss curve on the Sup-NatInst dataset.
  • Figure 4: Left: Training loss curves for Dolly-15k trained for 40,000 steps with different data mixing methods using the original category partitioning. Right: Average test loss on Dolly-15k after 40,000 training steps using original category partitioning. Highlighted values (with brown background) indicate the best overall performance.
  • Figure 5: Domain weight evolution during training. Our method dynamically adjusts the importance of each domain throughout the training process, with Domains 1 and 5 eventually receiving the highest weights while Domains 0, 2, 3, 7, 8, and 9 are downweighted over time.
  • ...and 1 more figures

Theorems & Definitions (11)

  • Definition 1
  • Lemma 1
  • Lemma 2
  • proof
  • Definition 2
  • Lemma 3
  • proof
  • Lemma 4
  • Definition 3
  • Lemma 5
  • ...and 1 more