RegMix: Data Mixture as Regression for Language Model Pre-training

Authors: Qian Liu, Xiaosen Zheng, Niklas Muennighoff, Guangtao Zeng, Longxu Dou, Tianyu Pang, Jing Jiang, Min Lin
Primary area: foundation or frontier models, including LLMs
Keywords: language model pre-training, data mixture, regression

Abstract:
The data mixture for large language model pre-training significantly impacts performance, yet how to determine an effective mixture remains unclear. We propose RegMix to automatically identify a high-performing data mixture by formulating it as a regression task. RegMix trains many small models on diverse data mixtures, uses regression to predict performance of unseen mixtures, and applies the best predicted mixture to train a large-scale model with orders of magnitude more compute. To empirically validate RegMix, we train 512 models with 1M parameters for 1B tokens to fit the regression model and predict the best data mixture. Using this mixture we train a 1B parameter model for 25B tokens (i.e. 1000× larger and 25× longer) which we find performs best among 64 candidate 1B parameter models with other mixtures. Furthermore, RegMix consistently outperforms human selection in experiments involving models up to 7B models trained on 100B tokens, while matching or exceeding DoReMi using just 10% of the computational resources. Our experiments also show that (1) Data mixtures significantly impact performance; (2) Web corpora rather than data perceived as high-quality like Wikipedia have the strongest positive correlation with downstream performance; (3) Domains interact in complex ways often contradicting common sense, thus automatic approaches like RegMix are needed; (4) Data mixture effects transcend scaling laws. Our code is available at https://github.com/sail-sg/regmix.

Decision: accept (spotlight)
Reviewer scores: 8, 8, 6, 6, 8

Reviews and discussion:
<thread>
creation date:
1729045484889

thread_type:
official_review

<thread_object>
creator: reviewer_ubls

<content>
summary:
This paper proposes a method that trains a simple regression model over data mixture ratios and the LLM loss on very small models, and then use the trained regressor to predict the best data mixture configuration for training larger scale models. The method is interesting and the experiments are relatively comprehensive. The paper considers two regression methods and verified the predicted mixture on a 1B model.

strengths:
1. the proposed method is novel and has practical impact to LLM training
2. the authors evaluated the predicted mixture on 1B model and compared to a few prior methods based on the downstream performance of the model
3. the paper also has good analysis and insights about regarding optimizing data mixture ratio, the relationship between validation PPL loss and downstream task performance, and the issues regarding scaling laws for data mixture.

weaknesses:
1. there are some details of the regression model that's not clearly explained. It is not clear to me how the authors fit the regression model, how many data points are used to fit the models. It would be nice to have a pseudocode or open-sourced script
2. the mixture weights for the baseline DoReMi is directly taken from the paper. However, it's not clear if the optimal weights learned using the DoReMi method would be different due to model and data processing differences. It's probably better to re-learn the weights using the small model in the current experiment setup.

questions:
1. in table 2, why do you only list the MSE for 1M model but rank correlation for all model sizes? what's the MSE for the other two sizes? Why is rank correlation a better metric here?
2. on line 243, the author mentioned the rank invariance hypothesis. However, it's not super clear to me what exactly this means and how this hypothesis is verified by the experimental results. Could you provide a clear definition of the rank invariance hypothesis and explicitly state how the experimental results support or verify this hypothesis?
3. there are some related work on data selection that falls into the group-level data selection category: https://aclanthology.org/2020.acl-main.754/
</content>
</thread_object>

</thread>
<thread>
creation date:
1730568503661

thread_type:
official_review

<thread_object>
creator: reviewer_wxj7

<content>
summary:
This paper formulates data mixing problems as a regression task. The authors propose to search data mixtures on small models and fit regression models to predict the optimal data mixture, which is then transferred to larger-scale model training. The authors empirically show the rankings of different data mixtures hold consistent between small and large-scale training. And data mixture found to be optimal at small scales can lead to improved performance compared to human heuristics and previous methods at large scales.

strengths:
1. The paper is well-written and easy to follow.
2. The author identifies a helpful assumption, namely ranking invariance regarding training scales, which helps reduce the cost of tuning data mixtures in this paper.
3. The proposed method is simple to implement and efficient, thus appealing to try in practice.
4. The paper contains extensive experiments to show optimizing data mixtures improves model performance.

weaknesses:
1. The feasibility of treating the data mixing problem as a regression problem has been an idea unveiled by previous studies [1,2], as the authors also mentioned in the paper.
2. The methods experiment on as many as 512 training runs. It is unclear whether the regression step is still necessary with so many experimented mixtures. This makes the proposed method actually a grid search with a small-scale proxy.
3. The proposed method highly depends on the assumption of ranking invariance regarding scales. The author only provides limited empirical results on this assumption. However, such an assumption is questionable according to [3]. It would be better if the authors provide more discussion to explain the scope where this assumption holds.

[1] Data Mixing Made Efficient: A Bivariate Scaling Law for Language Model Pretraining

[2] Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance

[3] Scaling Laws for Data Filtering— Data Curation cannot be Compute Agnostic

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730569680280

thread_type:
official_review

<thread_object>
creator: reviewer_fgaq

<content>
summary:
The paper proposes a method called REGMIX for automatically selecting an effective data mixture to optimize the pre-training of large language models. REGMIX formulates the data mixture selection as a regression task, training a set of small models with diverse data mixtures and fitting a regression model to predict their performance. The fitted regression model is then used to simulate and identify the top-performing data mixture, which is subsequently used to train a large-scale model. The empirical results demonstrate that REGMIX can improve downstream task performance and achieves results comparable to or surpassing the DoReMi method while using only 10% of the compute budget.

strengths:
- The paper presents a novel method, REGMIX, which formulates the data mixture selection problem as a regression task. This is a creative approach that leverages small-scale proxy models to predict optimal data mixtures for large-scale models.
- The authors conducted extensive experiments, training 512 models with 1M parameters on 1B tokens to fit the regression model. They then validated this model by training a 1B parameter model on 25B tokens, showing superior performance compared to human selection and the DoReMi method.
- The method allows for parallel training of small proxy models, making it more scalable than previous approaches that require training a single model for a long time.
- The paper provides several interesting findings, such as the significant impact of data mixtures on performance, the strong positive correlation of web corpora with downstream performance, and the complex interactions between domains.

weaknesses:
- The key assumption of rank invariance of data mixtures across different model sizes and token counts is not thoroughly validated. This assumption might not hold in all cases, especially with significant changes in model scale and data distribution.
- The paper claims stability across different proxy model sizes, but the experiments are limited to models with up to 1B parameters. It remains unclear if the method would be equally effective for much larger models commonly used in practice (e.g., 7B or 70B parameters). If so, the additional computation cost could not be ignored. 
- The authors only trained 25B tokens using the obtained data mixtures. This raises the question of whether the data scale could be enlarged to 50 times or even 100 times. And can LLM sitll benefit from the obtained data mixture?
- Although the method is more efficient than some previous approaches, training 512 small models still requires substantial computational resources. This could be a limitation for teams with limited access to such resources. The trade-off between performance gains and additional costs may not always hold when the model scales up.

questions:
- Can the authors provide more theoretical or empirical evidence to support the rank invariance assumption? How does this assumption hold up with significant changes in model scale and data distribution?
- How does REGMIX perform with proxy models larger than 1B parameters? Can the authors provide any preliminary results or insights on this? Or could the obtained data mixtures guide us to train a better model using much more tokens, e.g., 100B?
- Can the authors provide a detailed comparison of the computational resources required by REGMIX and other methods? This would help in understanding the practical feasibility of the method.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730687538471

thread_type:
official_review

<thread_object>
creator: reviewer_mox3

<content>
summary:
The paper presents RegMix, a new method for optimizing data mixtures in pre-training large language models (LLMs). Recognizing the importance of data composition, the authors frame mixture selection as a regression task. RegMix uses small proxy models trained on various data mixtures to build a predictive model that identifies the optimal mixture for larger models.

The authors conduct extensive experiments to validate the approach, showing that models trained with RegMix-selected data mixtures outperform those trained with mixtures chosen by other methods, including human selection and DoReMi, while utilizing only 10% of the compute budget.

The authors provide insights into the effects of data mixtures, offering empirical evidence that data mixtures can significantly impact performance, with single-task performance variations reaching up to 14.6%. They also emphasize the complex interactions that occur between different data domains.

strengths:
1. RegMix introduces a fresh approach by framing data mixture selection as a regression problem rather than relying on complex optimizations or heuristics, making the process scalable and computationally efficient.

2. The paper’s experimental setup is robust, with 512 small proxy models across diverse data mixtures, creating a solid regression model for data selection.

3. The paper is well-organized, clearly explaining the methodology and experiments. It introduces the hypothesis of rank invariance in data mixtures, supported by visual aids, making the regression model’s role easy to understand.

weaknesses:
The paper conducts a set of small-proxy models trained with small-scale tokens.

The paper only experiments with 1M models with 1B tokens. 
It is unclear how to decide the size of the proxy model parameter and training token.

questions:
How do we decide the size for proxy model and training tokens.?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730761461464

thread_type:
official_review

<thread_object>
creator: reviewer_yiah

<content>
summary:
The work introduces REGMIX, a method for optimizing data mixtures to enhance language model training efficiency. REGMIX treats data mixture selection as a regression task, using small proxy models to predict the performance of different mixtures and identify the best one, enabling larger models to be trained with significantly less compute. Key findings include:
- REGMIX’s mixtures perform as well as or better than those selected by human experts and prior methods like DoReMi, with only a fraction of the compute cost.
- Data mixture has a substantial impact on downstream performance, with single-task performance differences reaching up to 14.6%.
- General web corpora (such as CommonCrawl) outperform traditionally high-quality data like Wikipedia in driving downstream performance.
- Domain interactions are complex and often counterintuitive, highlighting the value of automated approaches like REGMIX.
- Data mixture effects go beyond scaling laws, with REGMIX capturing the complexity by jointly optimizing across all domains.

strengths:
- The paper introduces a novel, regression-based approach for selecting data mixtures that reduces computational costs in language model training. The approach offers an efficient alternative to traditional dynamic or heuristic data allocation methods, making a valuable contribution to the field.

- The paper is technically robust and well-structured, with extensive validation across diverse data scenarios. It empirically supports the rank invariance hypothesis and uses clear, well-structured figures to illustrate the method and results, enhancing reader understanding. REGMIX’s ability to match or outperform other DoReMi and other methods with significantly less compute is a compelling outcome for LLM pre-training efficiency.

- The paper tackles a pertinent problem in LLM pre-training. Given the increasing size of training data and models, this approach could have a significant impact on the field, especially in reducing computational costs and environmental impact.

weaknesses:
- To maximize impact, the authors could highlight specific scenarios where the approach enables previously infeasible experiments due to resource constraints. Also, adding a broader discussion on trade-offs of the method (e.g., scenarios where the rank invariance assumption might not hold) would help readers assess its practical relevance and future applicability.

- The work could have used standardized computation metrics, such as FLOPs or GPU hours, to allow clearer comparison of the method efficiency gains relative to baselines.

questions:
1) Can you further explain why the choice of multiplying the e token distribution by a value from 0.1 to 5.0? Is this a standard practice? Were other ranges tested, and if so, what were the results? 
     - Also, could you discuss the rationale behind this range and whether you conducted any sensitivity analyses to determine its impact on the results?

2) Given sufficient computation available, would segmenting domains further (into finer-grained topic-based segments) likely improve model performance or lead to more effective mixture predictions? Do you think that finer segmentation would affect the rank invariance assumption? I suggest the authors to discuss the potential impacts and challenges of a finer-grained domain segmentation in their future work section. This would help address the broader implications and limitations of their approach.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732190269710

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Rebuttal to All Reviewers

comment:
We thank all reviewers for their constructive feedback, and we have learned a lot from the feedback and have responded to each reviewer individually. We are very excited to receive this positive feedback from all of you. Here we want to highlight some new experimental results that may interest all reviewers.

- **New Experiments on 100 Domains**: We have extended RegMix to 100 domains, achieving an impressive correlation of 98.80% between the predicted ranks and true ranks of data mixtures. This demonstrates the scalability of RegMix when applied to a larger number of domains.

- **New Experiments with 280M Models**: We have added new experiments with 280M models and systematically demonstrate the rank invariance hypothesis across four model scales (1M, 60M, 280M, 1B) and three token scales (1B tokens, 5B tokens, 25B tokens).

- **Ongoing Experiment - 7B Model with Optimized Data Mixture**: We are currently applying the optimized data mixture from RegMix to a 7B model trained on 100B tokens. Once the results are available, we will compare its performance against the human baseline.

- **Ongoing Experiment - 1B Level Proxy Models**: We are also investigating using 1B models as proxy models to generate optimized data mixtures for training 7B models. This experiment aims to provide additional insights into the selection of appropriate proxy model sizes. We will report back once the results are available.

- **Paper Revision**: We have made the following updates:
   - $\textrm{\color{blue}Appendix A}$: describe the broader impact of our paper.
   - $\textrm{\color{blue}Appendix B}$: include a detailed discussion on the limitations around the rank invariance hypothesis.
   - $\textrm{\color{blue}Appendix H}$: extend RegMix to 100 domains.
   - $\textrm{\color{blue}Appendix I}$: add the pseudocode for RegMix.
   - $\textrm{\color{blue}Appendix J}$: provide additional empirical evidence supporting the rank invariance hypothesis.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732539747094

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
More Experimental Results by Authors

comment:
Hi all reviewers, thank you once again for your valuable comments and suggestions, which have been very helpful for us! We understand that this is a particularly busy time, and we greatly appreciate it if you could take a moment to review our additional experiments and response and let us know if they adequately address your concerns. Should there be any additional feedback, we will do our best to incorporate it in the next few days.

As previously committed, we have completed the experiments and are now reporting our findings:
- $\textrm{\color{blue}Appendix K}$: Comparing RegMix and Human on a 7B model over 100B tokens, RegMix continues to outperform Human by a significant margin, with the ability to accelerate performance by up to 75% on some benchmarks.
- $\textrm{\color{blue}Appendix L}$: Investigating whether it could bring more gains by training large proxy models (e.g., 1B). The conclusion is that our 1M proxy model is already strong enough, so we recommend researchers to begin with the ultra small models for data mixture optimization.

## Appendix K: Scaling RegMix to 7B Models over 100B Tokens

In order to further validate the effectiveness of our method on larger models, we conducted an experiment using Human baseline and RegMix data mixtures on a 7B model trained on 100B tokens. The results, summarized in ${\textrm{\color{blue}Table 10}}$  (see below), demonstrate that RegMix still outperforms the Human baseline, achieving an average performance boost of 2%:

| Benchmark | Human | Our Method |
| --- | --- | --- |
| Social IQA | 41.6 | 43.3 |
| HellaSwag | 55.0 | 63.3 |
| PiQA | 72.4 | 75.3 |
| OpenBookQA | 34.1 | 36.7 |
| Lambada | 44.9 | 51.0 |
| SciQ | 91.7 | 91.2 |
| ARC Easy | 63.5 | 65.4 |
| COPA | 75.0 | 80.1 |
| RACE | 35.1 | 36.6 |
| LogiQA | 25.7 | 24.1 |
| QQP | 58.9 | 56.1 |
| WinoGrande | 58.6 | 60.7 |
| MultiRC | 52.0 | 51.0 |
| **Average** | **54.5** | **56.5 (+2.0)** |

To provide a more detailed illustration, we benchmarked the downstream performance of RegMix and Human on every dataset at intervals of 25B tokens in $\textrm{\color{blue}Figure 16}$ of $\textrm{\color{blue}Appendix K}$. The results show that RegMix can significantly speed up pre-training, with a **50% acceleration on most benchmarks (e.g., HellaSwag)** and **up to 75% on some benchmarks (e.g., PiQA)**. Notably, the performance boost does not decrease with the amount of training tokens. However, we also observed that both RegMix and Human struggle to improve on certain benchmarks (e.g., MultiRC), even with increased token amounts. These findings suggest that RegMix mainly benefits downstream tasks whose performance increases with the amount of training data, but may not improve tasks that do not follow scaling laws. This observation is intriguing and worthy for further investigation.

---

## Appendix L: Using Larger Proxy Model

We conducted a preliminary study on the impact of proxy model size on effectiveness. Specifically, we compared two configurations: (1) 128 proxy models of 1B parameters each, and (2) 512 proxy models of 1M parameters each (the setting used in our main experiments). Both configurations used 1B training tokens per proxy model. We limited our investigation to these configurations due to computational constraints that prevented us from exploring scenarios with more 1B-parameter proxy models.

To evaluate these configurations, we used their respective optimized data mixtures to train two 7B models on 100B tokens and compared their performance. The results, summarized in $\textrm{\color{blue}Table 11}$ (see below), show that both proxy settings achieved similar average performance across downstream tasks. This suggests that increasing proxy model size, even with fewer proxy models, can maintain competitive performance. However, given that the 1B proxy models do not significantly outperform the 1M proxy models, and considering that they incur over much more computational overhead, we recommend prioritizing a larger number of smaller proxy models over fewer larger ones. Based on our findings, we suggest practitioners begin with ultra-small proxy models (e.g., 1M parameters in our setting) as a starting point to optimize data mixtures for language model pre-training.

| **Benchmark** | **RegMix** (1B as proxy) | **RegMix** (1M as proxy) |
|--------------|------------------------------|------------------------------|
| Social IQA | 43.4 | 43.3 |
| HellaSwag | 62.9 | 63.3 |
| PiQA | 75.1 | 75.3 |
| OpenBookQA | 36.2 | 36.7 |
| Lambada | 50.0 | 51.0 |
| SciQ | 91.2 | 91.2 |
| ARC Easy | 65.9 | 65.4 |
| COPA | 79.6 | 80.1 |
| RACE | 35.6 | 36.6 |
| LogiQA | 23.9 | 24.1 |
| QQP | 56.6 | 56.1 |
| WinoGrande | 60.7 | 60.7 |
| MultiRC | 51.6 | 51.0 |
| **Average** | **56.4** | **56.5** |
</content>
</thread_object>

</thread>
<thread>
creation date:
1732679959332

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Comment from Authors

comment:
Dear Reviewers,

Since the rebuttal period has been extended to December 2, we are excited to have another week to  welcome any additional questions or clarifications you may have regarding our submission. Based on the previous round of reviews, we have made significant improvements to RegMix:

- Expanding the experiments to 1,000 domains for RegMix and verify its effectiveness
- Verifying performance improvement on 7B-level models over 100B tokens
- Systematically verifying the rank invariance hypothesis on model scales and more token counts
- Investigating proxy model sizes at the 1B level with new insights

We eagerly welcome the opportunity to further discuss the potential of RegMix. If you have any specific inquiries or require additional information, please do not hesitate to share them. We are sincerely grateful to all reviewers for their constructive, helpful, and timely feedback.

Best Regards,

Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces REGMIX, a novel method for optimizing data mixtures to improve the efficiency of pre-training LLMs. The method frames data mixture selection as a regression task, using small proxy models trained on diverse data mixtures to predict the best-performing mixture for larger-scale training. REGMIX presents a significant advance in automating data mixture selection, reducing computational costs while achieving or surpassing the performance of prior methods and human expertise. Detailed key findings and contributions include:
- Methodology: REGMIX uses small models to evaluate various data mixtures, fits a regression model to predict their performance, and identifies the optimal mixture for larger-scale training.
- Efficiency and Effectiveness (showed empirically): REGMIX-selected mixtures perform as well as or better than mixtures chosen by human experts or prior methods like DoReMi while using only 10% of the compute budget. The method also demonstrates improved downstream task performance, with single-task performance differences of up to 14%. These findings were obtained with comprehensive experiments, including testing with a 1B parameter model, and comparing two regression methods for modeling data mixtures.
- Key Insights: Data composition has a substantial impact on LLM performance, often in counterintuitive ways, emphasizing the importance of automated approaches like REGMIX. Authors also found general web corpora (e.g., CommonCrawl) outperform traditionally high-quality datasets like Wikipedia for downstream performance, and REGMIX can capture complex interactions across data domains, revealing insights that go beyond existing scaling laws.


Strength of this paper

- Novel and Effective Approach: The paper introduces a novel regression-based method, REGMIX, for optimizing data mixtures in LLM pre-training, reducing computational costs while maintaining or surpassing performance compared to methods like DoReMi and human selection. By framing data mixture selection as a regression task, the method avoids complex optimization or heuristic approaches, making it computationally efficient and scalable.
- Experimental Rigor: Extensive experiments to cover diverse data scenarios and validate the method, including training 512 small proxy models on diverse data mixtures and testing with a 1B parameter model trained on 25B tokens. The authors empirically support the rank invariance hypothesis, demonstrating consistency in mixture rankings between small and large models.
- Scalability and Practicality: The method enables parallel training of small proxy models, significantly reducing the time and compute required compared to traditional approaches. It is simple to implement and efficient, making it highly practical for real-world applications.
- Impact and Relevance: The approach addresses a critical problem in LLM pre-training, offering a solution that reduces computational and environmental costs, which is increasingly important as models and datasets grow larger.
- The paper is well-written, clearly explaining the methodology, experiments, and findings. 

Weakness of this paper

Several reviewers raised few concerns/limitations of this paper. By addressing these limitations, the paper could strengthen its experiment and expand impact.

Weaknesses of the Paper:
- Experimental scale and settings: The experiments are limited to models with 1M parameters trained on 1B tokens and validated with 1B parameter models trained on 25B tokens. The scalability to much larger datasets or model sizes is untested. It is also unclear how to determine the optimal size of the proxy model parameters or the training tokens, which limits the reproducibility and generalization of the method. The paper does not use standardized computation metrics like FLOPs or GPU hours, making it harder to quantify efficiency gains relative to baselines.
- Impact, applicability, and reproducibility: The authors could better emphasize specific scenarios where the approach enables previously infeasible experiments due to resource constraints, or discuss on trade-offs, such as when the rank invariance assumption may not hold, to help readers better understand the method's limitations and future applicability. Key details about the regression model, such as the number of data points used for fitting, are not clearly explained. Adding pseudocode or open-sourced scripts would improve clarity and reproducibility.
- Baseline Comparisons: The mixture weights for DoReMi were directly taken from the original paper, but differences in model or data processing could affect the results. Re-learning these weights in the current setup would provide a fairer comparison.