TL;DR
Existing data mixing methods for foundation models rely on rigid, human-defined task categories and require computationally prohibitive evaluation passes to determine mixture weights. R&B (Regroup & Balance) solves this by first clustering data into semantically coherent groups via text embeddings, then dynamically reweighting domain mixtures on the fly using training gradients accumulated during standard backpropagation. This two-stage approach matches or beats state-of-the-art data mixing performance across text, reasoning, and multimodal tasks while reducing evaluation compute overhead by over 99% (adding as little as 0.01% compute).
What Problem Does This Paper Address?
Optimizing the proportions of different data domains during pretraining and instruction tuning ("data mixture optimization") can dramatically boost model capabilities and sample efficiency. However, prevailing methods hit two major roadblocks:
- Suboptimal Domain Definitions: Datasets are usually partitioned by coarse human labels or data origins (e.g., "Summarization", "Open QA", or "Wikipedia"). These surface-level categories fail to reflect how neural networks actually share knowledge and transfer representations across tasks.
- Computational Inefficiency: State-of-the-art dynamic reweighting methods (like Skill-It, Aioli, and DGA) require periodic auxiliary forward or backward passes over large evaluation benchmarks or per-domain subsets. As the number of domains grows, the compute overhead explodes—sometimes exceeding the cost of training the model itself by orders of magnitude.
Methods
R&B decomposes data mixture optimization into two decoupled stages:
[ Raw Training Data ] ──(ModernBERT Embeddings + k-Means)──> [ Semantically Regrouped Clusters ]
│
(Standard Training)
│
[ Dynamically Balanced Mixture p^(t+1) ] <──(Softmax of Gram Matrix Gp)── [ Accumulate Gradients on-the-fly ]
- Stage 1: Regroup (Semantic Partitioning): Instead of using metadata categories, the authors embed input queries using ModernBERT and apply $k$-means clustering. The optimal number of clusters $k^*$ can be predicted using standard clustering metrics like the silhouette score, eliminating the need to train multiple exploratory models.
- Stage 2: Balance (Gradient-Based Dynamic Reweighting):
- Rather than running separate forward/backward evaluation passes, R&B uses an efficient chain-rule decomposition (Goodfellow, 2015) to track final-layer gradients for each domain during normal mini-batch backpropagation.
- At each reweighting window, it constructs a domain gradient Gram matrix $G$, where $G_{ij} = \frac{1}{|S_i||S_j|} \nabla L_i^\top \nabla L_j$.
- R&B solves a regularized linear objective over the probability simplex, yielding a closed-form softmax update: $$p^{t+1} = \text{softmax}\left(\frac{\lambda}{\|Gp\|_2} Gp\right)$$
- The evaluation cost scales purely with $O(T m^2 N)$ (where $m$ is the number of domains, $T$ is the number of reweighting rounds, and $N$ is parameter count), completely bypassing dependence on evaluation set size $D_e$.
Datasets
The authors evaluate R&B across five diverse setups:
- Natural Language & Instruction Tuning:
- Dolly-15k: 15,000 instruction-following examples across 8 original categories.
- Super-NaturalInstructions (In-Domain): 285k examples from 100 tasks spanning 38 original skill categories.
- Super-NaturalInstructions (Out-of-Domain / Test): 3.56M examples across 60 original skill categories.
- Reasoning:
- s1 / s1-1.1k / s1-59k: Challenging math reasoning traces drawn from 54 distinct sources.
- Multimodal:
- DataComp (small-scale): Filtered down to 3.8M high-quality image-text pairs (top 30% CLIP score) from 10M raw samples, evaluated across 38 downstream tasks (ImageNet variants, VTAB, retrieval).
Key Insights and Ideas
- Content Matters More Than Task Format: Clustering Dolly-15k did not reproduce task types like "closed QA" or "summarization." Instead, it grouped data by subject matter (e.g., biology, sports, pop culture, lifestyle). Training on these topical partitions significantly outperformed training on format-based categories.
- Silhouette Score Predicts Mixture Quality: The compactness and separation of clusters (measured via silhouette scores) strongly correlated with final evaluation loss (e.g., $R^2 = 0.895$ on S1-59k), enabling unsupervised domain selection prior to model training.
- Eval-Free Alignment via Kernel Re-use: Because the inner product between training domain gradients and target domain distributions can be approximated via the Gram matrix $G$, models dynamically favor domains that benefit the overall target distribution without computing target-set gradients at runtime.
Key Results
- Massive Compute Reduction with Superior Loss:
- Super-NatInst In-Domain (GPT-Neo 125M): R&B achieved the lowest evaluation loss (2.381) with 0.009% compute overhead, whereas Skill-It and Aioli incurred 595.5% and 1336.5% overheads with worse losses (2.632 and 2.622).
- Super-NatInst Test: R&B scored 2.859 with 0.1% overhead, outperforming Skill-It ($6 \times 10^7\%$ overhead) and Aioli ($7 \times 10^6\%$ overhead).
- Dolly-15k: Achieved a loss of 2.765 with only 0.0006% overhead.
- Accelerated Convergence: R&B converged to baseline validation loss levels in only 20% of the training steps required by static or alternative dynamic schedulers on NaturalInstructions, and in 70% of steps on math reasoning (Qwen2.5-1.5B on s1).
- Multimodal Scalability: On DataComp CLIP training, R&B consistently outperformed stratified sampling across 38 benchmarks when cluster granularity increased ($m \ge 20$), yielding a 3.27% relative improvement at 50 clusters.
Limitations and Future Work
- Out-of-Domain Alignment Constraints: R&B requires evaluation categories to correspond to training clusters. When evaluation data is strictly out-of-distribution, eval samples must be mapped to their nearest training clusters; unmapped or disjoint domains can cause degenerate gradient norms.
- Fixed Pre-Clustering: Semantic regrouping is currently performed once prior to training using a frozen embedder (ModernBERT) and $k$-means. Dynamically updating cluster assignments as the model's internal representation evolves remains unexplored.
- Implications for Scaling Laws: Existing scaling laws assume static data proportions. R&B's dynamic shifts highlight the need for new theoretical scaling formulations that account for online, adaptive data allocation.