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
Venue: ICLR 2025 Spotlight (2025)
Keywords: language model pre-training
data mixture
regression
Primary area: foundation or frontier models, including LLMs
TL;DR: We introduce RegMix, an automated data mixture method that formulates data mixture as a regression problem. RegMix achieves a 6.3% improvement over human selection on the HellaSwag benchmark, with only a 2% extra training FLOPs.

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.

Published as a conference paper at ICLR 2025
REGMIX: DATA MIXTURE AS REGRESSION FOR LANGUAGE MODEL PRE-TRAINING
Qian Liu'
Xiaosen Zheng Niklas Muennighoff 3.4
Guangtao Zeng"
Longxu Dou'
Tianyu Pang'
Jing Jiang
Min Lin-
'Sea AI Lab
SMU CContettual *Stanford University
SSUTD
1 liugian. scal gmaal..maaa......mmmmmmmm....... com; xs xezheng. 2020@phdcs. smu, edu.sg
ABSTRACT
The data mixture for large language model pre-training significantly impacts per- formance, yet how to determine an effective mixture remains unclear. We propose REGMIX to automatically identify a perforrir data mixture by formulating it as regression task. REGMIX trains many small models on diverse data mixtures, uses regression to predict performance of unseen mixtures, and applies the best pre- dicted mixture to train large-scale model with orders of magnitude more compute. To empirically validate REGMIX, we train 512 models with 1M parameters for IB tokens to fit the regression model and predict the best data mixture. Using this mixture we train a IB parameter model for 25B tokens (i.c. 1000 x larger and x 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: / gi thub. com/ sail- sg/ regmix.
Winogrande
QQP
RACE
OpenBookkA
62
38
60
95
34g
34
50
50
326
RagMix
Homan
RegMix Homon
32
Roghlix
304
RegMix
Homan
25B
75B
1008
25R
50R
25B
1008
50B
75B
Social IQA
HellaSwag
Lanbuda
PIQA
Performance
44
42-
55
ReyMix
50
+ RegMix Hummn
RegMix
70
RegMix
Human
45
68g
Haman
25B
50B
75B
1008
25B
508
75H
1008
25B
1008
25B
50B
758
1008
SdO
MuttiCC:
ARC-Faay
Performance (96)
95
07
60
68
854
55
72
RagMix Homan
75
Reghis Homan
45
RegMix
Homan
RegMix
Homon
25B
50B
75B
1008
25B
50B
75B
1008
25B
S0B
75B
100B
25B3
508
758
100B
Figure 1: Performance comparison 7B parameter models between Human selection and REGMIX across various tasks and training token quantities (e.g., up to 100B). REGMIX demonstrates consistent performance improvements over Human across varying data scales for the majority of tasks evaluated. In particular, for tasks exhibiting positive scaling behavior, the performance benefit is preserved with increased training data, suggesting robust scaling properties of REGMIX.
"The first authors contributed equally.

---

Published as a conference paper at ICLR 2025
INTRODUCTION
The availability of large-scale public datasets has been a key factor enabling the creation of large language models (LLMs). Most data is available on the Internet and includes academic papers (e.g. arXiv), books (e.g. Project Gutenberg), and code (e.g. GitHub). For the creation of one of the first LLMs, GPT-3 (Brown et al., 2020), the authors had already recognized the importance of selecting the best data for training, and thus they decided to upsample Wikipedia due to its perceived high quality. However, such manual data selection is not scalable and may lead to suboptimal selection (Albalak et al., 2024). As the size and diversity of data used for LLM pre-training continue to grow, determining the optimal data mixture becomes increasingly challenging. It gives rise to the critical research question: How can we select the optimal data mixture in scalable and efficient manner?
Prior work (Xie al., 2023a; Fan al., 2023; Albalak al., 2023) employs small-scale models ("proxy models") to predict the domain weights for large-scale language models. These works train proxy models with a substantial number of tokens (e.g.. 100B), sometimes even the same number as used for training LLMs, and dynamically adjust the data allocation strategy by monitoring the training dynamics. However, these approaches become inefficient as the training data used for pre-training LLMs continues to grow. Training proxy model for current models, such as Llama-3, would require using up to 15T tokens (AI, 2024), which is likely too expensive and too slow to make worthwhile.
In this work, we argue that training small models a limited set of tokens is sufficient to predict an effective data mixture for LLM training. Our key assumption is the rank invariance of data mixtures, which posits that the relative ranking of data mixtures in terms of their impact on model performance is consistent across different model sizes and numbers of training tokens. Under this assumption, the key challenge lies in discovering the top-ranked data mixture from the infinite number of potential data mixtures. To do so, we treat the data mixture selection as a regression task. Rather than exhaustively training small models with every possible mixture, we train only set of small models, each with unique data mixture. Based on the performance of these models and their mixtures, we fit regression model predict the performance of other data mixtures. Our approach is significantly more scalable than prior work, as it allows for parallel training of small proxy models rather than training a single model for a long time. Further, the regression model provides insights into domain interactions that can facilitate understanding and data curation.
To validate REGMIX, we train models with IM and IB parameters' with different data mixtures. By training 512 models with IM parameters on 1B tokens, we are able to predict the optimal data mixture among 64 models that are 1000 larger (1B parameters) and trained 25 longer (25B tokens) as depicted in Figure 2. Moreover, the optimized data mixture using REGMIX yields a better model than human selection, and achieves performance on par with the flagship DoReMi method (Xie et al., 2023a) despite it requiring less total compute and allowing for parallel training. We also find that (1) Data mixture significantly impacts downstream performance, resulting in substantial differences of up to 14.6% in single-task performance; (2) General web corpora (e.g., CommonCrawl), rather than Wikipedia, exhibit the strongest positive correlation with improved performance across downstream tasks; (3) The interactions between domains complex and often contradict intuition, highlighting the need for automated approaches like REGMIX; (4) Data mixture effects transcend scaling laws, and REGMIX captures the complexity by considering all domains together.
2 RELATED WORK
Data selection and mixture is concerned with curating data to optimize some goals, usually model performance (Koh & Liang, 2017: Albalak et al., 2024). Prior methods can be categorized into: (1) Token-level selection is the most fine- ane-grained level of selection dealing with the filtering of tokens (Lin et al., 2024). (2) Sample-level selection is about choosing individual training examples. is commonly employed for selecting fine-tuning data (Thakkar et al., 2023; Das & Khetan, 2023; Xie et al., 2023b; Engstrom et al., 2024; Xia et al., 2024; Liu et al., 2024; Bukharin & Zhao, 2023; Kang et al., 2024; Mekala et al., 2024; Sachin Parkar et al., 2024; Yang et al., 2024b). For the pre-training of LLMs, most methods rely on heuristics (Rae et al., 2021; Sharma et al., 2024: Soldaini et al., 2024;
'These approaches often suffer from instability issues. Details can be found in Appendix F.
2 ?Our model sizes mentioned in this paper refer to the numher non-embedding parameters, as embedding parameters account for a disproportionately large portion in smaller models. 'The estimated FLOPss for training 512x 1M models is nearly 2% of the FLOPS required for one 1B model.
2

---

Published as a conference paper at ICLR 2025
Hypothesis: Rank
Invariance of Data Mixtures
Validation Loss
3.2
Small Models
Large Models
3.0
Few Tokens
More Tokens
2.8
Best Data Mixture
8
16
24 32
40
48
56
64
Predicted Rank
Figure 2: Left: We hypothesize the rank invariance of data mixtures across model sizes and numbers of training tokens. Leveraging this hypothesis, we use small models trained on fewer tokens to predict the effective data mixture for training large models with substantially more tokens. Right: By training 512x x 1M models, our method identifies the best data mixture prior to training 64x 1B models. The predicted best data mixture, denoted by the red star, achieves the lowest validation loss.
Li et al., 2024), but there have been some learned approaches using optimization algorithms (Chen al., 2024; Mindermann et al., 2022; Shao al., 2024; Yu et al., 2024), model perplexity (Marion al., 2023; Muennighoff et al., 2023; Ankner al., 2024). LLMs to inform the sample selection process (Wettig et al., 2024; Sachdeva et al., 2024; Zhang et al., 2024b). (3) Group-level selection assumes the data be grouped into pools that then optimally mixed. While early work again relies on manual mixtures (Gao et al., 2021; Brown et al., 2020), learned mixtures have become more common (Albalak al., 2024). Learned approaches either leverage proxy models determine fixed weights for each group ("offline selection") (Rae et al., 2021; Xie et al., 2023a; Fan et al., 2023) or dynamically adjust the weights during training of the final model ("online selection") (Wang et al., 2020; Chen et al., 2023; Albalak et al., 2023). Our approach, REGMIX, is an offline group-level selection method. Different from the flagship algorithm in this category, DoReMi (Xie al., 2023a), REGMIX does not require training single model for hundreds of thousands of steps, but instead a few small models for short durations. these can be trained parallel, our approach more scalable, while also yielding better weights leading to a more performant final model.
Data scaling laws explore interactions data quantity, quality, and mixing proportions, as LLMs are scaled up. Muennighoff et al. (2023) introduce scaling laws for data-constrained scenarios and Goyal et al. (2024) try to extend this approach to deal with multiple data pools. Prior research has confirmed that different datasets require different scaling (Hoffmann et al., 2022; Pandey, 2024), thus Ye et al. (2024) and Ge et al. (2024) propose functional relationships to predict the impact of mixtures on language modeling loss. Some work has investigated getett optimal mixtures during continued pre-training rather than from seratch training (Que et al., 2024; Dou et al., 2024). While most of these works focus on validation loss, others investigate downstream performance and develop predictive relations with loss (Gadre et al., 2024; Yang et al., 2024a; Xia et al., 2022). Different from data scaling works that attempt to find an analytical scaling function (Hoffmann et al., 2022), REGMIX directly optimizes the target metric using regression models. REGMIX is designed for from-scratch pre-training. In line with previous research (Huang et al., 2024), we also find strong correlations between loss and downstream performance, especially for loss web corpora.
3 REGMIX: DATA MIXTURE AS REGRESSION
As illustrated in Figure 3, our method involves four key steps: (1) Generate random data mixtures and train small-scale proxy models these mixtures. (2) Fit linear regression model using the mixtures as features and the target value as the label. (3) Simulate the data mixture space on a larger scale and leverage the regression model to identify the best mixture for the target value. (4) Train large-scale model using the simulated best data mixture.
3.1 TRAIN SMALL-SCALE PROXY MODELS
The first step is to train a set of small-scale proxy models on multiple different data mixtures. To reduce the required runs, aim to select diverse range of data mixtures that extreme weights from 0% to 100% for each domain. We achieve this by using a Dirichlet distribution based on the token distribution, which allows sample wide range of values and expose the regression
3

---

Published as a conference paper at ICLR 2025
Train small-scale proxy models
Fit regression model
using data mixture as source
Hacker News
Github
Philpapers
Target +
Regression Model
9.5%
35.9%
54.6%
5.46
step
Linear Model
Tree Model
87.7%
12.0%
0.3%
5.57
step
24.4%
1.4%
74.2%
6.07
Simulate new data mixtures and predict Target
Train a large-scale model on the best mixture
Simulated Best Data Mixture
Hacker News
Github
Philpapers
Prediction
6.0
22.8%
67.0%
10.2%
5,6
Predicted Best Target
Prediction (Lowest)
20 40 60 Hacker News (%)
60 40 20 Github (%)
5.34
The illustration of our method using Hacker News, GitHub, and Philpapers as training domains, with the loss on the StackExchange domain as the Target (where indicates lower is better). regression model is fitted using small-scale proxy model training logs and employed to predict the best data mixture within the simulation space, enabling direct prediction of the data mixture for large-scale language model pre-traee pre-training. Note that the Philpapers domain is omitted in the simulation plot (3) for simplicity.
Figure 3:
Table 1: Overview of the Pile dataset (Gao al., 2021) with datasets that are longer available due to copyright issues marked in gray. In our experiments, we use the available domains to study the data mixture for language model pre-training.
Component
Effective Size
Component
Effective Size
Pile-CC
227.12 GiB
OpenSubtitles Wikipedia (en)
19.47 GiB
PubMed Central
180.55 GiB
19.13 GiB
Books3
151.44 GiB
DM Mathematics
15.49 GiB
Open WexTexr2
125.54 GiB
Uhuntu IRC
11.03 GiB
ArXiv
112.42 GiB
BookCorpus2
9.45 GiB
Github
95.16 GiB
EuroParl
9.17 GiB
FreeLaw
76.73 GiB
HackerNews
7.80 GiB
Stack Exchange
64.39 GiB
YoutubeSubtitles
7.47 GiB
USPTO Backgrounds
45.81 GiB
PhilPapers
4.76 GiB
PubMed Abstracts
38.53 GiB
NIH ExPorter
3.79 GiB
Gutenberg (PG-19) 19)
27.19 GiB
Enron Emails
1.76 GiB
models various extremes. Simultaneously, basing the distribution on the token distribution ensures that the overall data mixture statistically reflects the availability of data. For example, this prevents any single domain with token count below 1 % from being overly emphasized, which is not feasible for lcrge-scale training since there are not enough available tokens from that domain. In practice, we multiply the token distribution by a value from 0.1 to 5.0 to construct various sparse and near-uniform distributions, then use these distribution vectors as the Dirichlet distribution hyperparameter a.
After training small-scale proxy models for few steps, we can obtain several well-trained small models. For example, in our main experiment, each proxy model contains 1M parameters and is trained on IB tokens. We can then choose to evaluate these trained models on domains or benchmarks to get the target value we want to optimize. Generally, the target value can be the loss domain, as shown in Figure 3 for the StackExchange domain. Once we have obtained these target values, we can use the data mixture as features and the target values labels to fit regression model.
3.2
FIT A REGRESSION MODEL
The second step is to fit regression model using the data mixture features, and the target value as labels. The regression task is a conventional supervised learning task that involves predicting a continuous target variable y based input features X (21, (x1, 22: The goal is find
4

---

Published as a conference paper at ICLR 2025
function f that best maps the input features to the target variable, such that y y= f(X) + where represents the error or noise in the data. In the context of this paper, the input features X correspond to the domain weights of the data mixture, and the target variable y is the value we want to optimize. Using this data, we train regression models that learn a function to predict the target value based on arbitrary data mixtures without requiring further training.
Linear regression. The linear regression model is widely used in regression. It assumes a linear relationship between the input features and the target variable, which can be represented
y = wo + w +... + Wn +
(I)
where wi the intercept, and w (wi,.... . Wn) are the coefficients associated with the respective input features T1, ....... n- The coefficients w are typically estimated using techniques such as ordinary least squares, aiming to minimize the sum of squared residuals between the predicted and
actual target values. In practice, we employ linear regression with L2 regularization, also known as ridge regression, which applies a penalty to the magnitude of w to prevent overfitting.
LightGBM regression. LightGBM (Ke et al., 2017) is powerful gradient- that can be used for both regression and classification tasks. In the context of regression, LightGBM learns an ensemble of decision trees to predict the target variable. The process is guided by a gradient- based optimization algorithm, which minimizes a specified loss function (e.g. mean squared error). Moreover, LightGBM is designed to be efficient and scalable, making it suitable for large datasets.
algorithm
3.3 SIMULATE AND PREDICT
Once we have trained the regression model, we can efficiently explore the entire space of possible data mixtures. By using the trained model to predict the target value for each potential data mixture, we can quickly identify the input that yields the best target value. This simulation based optimization is relatively cheap, both the simulation and the regression prediction are computationally fast. For example, running prediction for 1, 1,000,000 data mixtures takes less than 10 CPU seconds.
3.4 LARGE-SCALE MODEL TRAINING
After identifying the best data mixture with simulation, we generalize the top-ranked data mixture to large-scale model training with many more tokens. As shown in Figure 3, we directly the best data mixture for training the larger model. In practice, to increase the robustness of our regression pre- diction, we select the top 100 mixtures and average them as the data mixture for large-scale training.
4 EVALUATING ON REGRESSION PREDICTION
In this section, we evaluate the ability of REGMIX to predict the effect of unseen data mixtures. First, we fit the regression model using training artifacts of small (i.e., 1M parameter) models and evaluate the loss prediction performance on small models. Then, to verify our rank invariance hypothesis, we test the learned regression on predicting the rank across model sizes and the number of tokens.
4.1 EXPERIMENTAL SETUP
Datasets and models. We conduct our experiments using the domains of the Pile dataset (Gao et al., 2021) depicted in Table 1. Due copyright concerns, we utilize the 17 subsets available on HuggingFace 4 that do not violate copyright issues. We consider both linear and LightGBM regression models, where the target y set to be the validation loss of the Pile-CC domain.
Training and evaluation. The regression model is fitted using the training artifacts of 512x IM models with IB tokens, and evaluated on 256x unseen data mixtures for 1M, 60M models (each trained with IB tokens) and 64x unseen data mixtures for IB models (each trained with 25B tokens).
Evaluation metrics. We use two metrics benchmark our regression models: (1) Spearman Rank Correlation (p) is nan-paramt measure of the strength and direction of the association between two ranked variables. (2) Mean Squared Error (MSE) is common metric used evaluate the model by measuring the average squared differences between the predicted and actual values.
htw.:///wooo.ooooooeeoeoorrrrrro uuucopyrrghted
5

---

Published as a conference paper at ICLR 2025
Table 2: We fit the regression model based the results of the 512x 1M parameter models trained on IB tokens, and evaluate it on unseen data mixtures for IM, 60M, and IB parameter models depicted below. The Spearman correlation p compares the predicted and actual ranks, while MSE measures the loss prediction performance. /M refers to models with IM parameters trained on 1B tokens, 60M refers to models with 60M parameters trained on 1B tokens, and /B refers to models with IB parameters trained on 25B tokens. Due to scale differences, MSE values for 60M and 1B models are not directly comparable and are omitted from analysis.
Test On
IM
60M
1B
Method
(t)
MSE (1)
P (1)
(t)
Linear
90.08
0.13
89.26
88.01
LightGBM
98.45
0.04
98.64
97.12
100
LightGBM (%)
100
Linear p (%)
512
90
90
128
85
85
64
80
0.4
0.6 0.8
0.2
0.6 0.8
Training Tokens (B)
Training Tokens (B)
Figure 4: The plot of Spearman Rank Correlation P between the predicted ranks and true ranks of Linear regression (Left) and LightGBM regression (Right) across different training tokens and different number of proxy models. As shown, increasing the number of proxy models significantly boosts p. while adding more training tokens has diminishing returns.
4.2 EXPERIMENTAL RESULTS
High correlation across model sizes. As shown in Table 2, the LightGBM model demonstrates superior performance over linear regression models across all three metries, with its advantage becom- ing increasingly pronounced when evaluating on larger models with more training tokens. Meanwhile, the fact that 1M models trained with IB tokens can achieve such a high correlation of 97.12% on unseen mixtures of IB models with 25B tokens directly validates our rank invariance hypothesis.
Proxy model count outweighs training token count. Given the same FLOPs budget for small- scale training, we can either increase the token count (i.e., the number of training tokens) or the number of proxy models. Therefore, we study which approach would yield better performance. As shown in Figure 4, increasing the training tokens of the proxy models saturates after approximately 0.25B tokens. In contrast, increasing the number of proxy models consistently enhances performance, particularly for the LightGBM model. Notably, the performance of 512 models trained on 0.2B tokens surpasses that of 128 models trained on 0.8B tokens, indicating that increasing the number of proxy models is more effective than increasing the training token count beyond certain token threshold.
EVALUATING ON DOWNSTREAM TASKS
In this section, we apply our method to demonstrate its effectiveness on realistic downstream tasks. For evaluation, we exclude specific benchmarks that exhibit large performance variance (e.g.. RTE) according to the performance traces reported previous work (Mehta et al., 2024) and our observations during pre- pre-trnin Ultimately, we select the following benchmarks tasks: Social IQA (Sap et al., 2019), HellaSwag (Zellers et al., 2019), PiQA (Bisk et al., 2020), OpenBookQA (Mihaylov et al., 2018), Lambada (Paperno et al., 2016), SciQ (Welbl al., 2017), ARC Easy (Clark et al., 8), COPA (Sarlin et al., 2020), RACE (Lai et al., 2017), LogiQA (Liuu al., 2020), QQP (Wang al., 2018), WinoGrandd (Sakaguchi al., 2021), and MultiRC (Khashabi et al., 2018). These benchmarks cover diverse range of tasks, enabling a comprehensive evaluation of the real-world impact of REGMIX. For each benchmark, we use normalized accuracy as the evaluation metric if provided by haleess- (Gao et al., 2023) else we use regular accuracy.
downstream

---

Published as a conference paper at ICLR 2025
Table 3: We experiment with 64 models, each with 1B parameters trained different data mixtures, and evaluate their performance across various benchmarks. The reported performance on each task the average score from O-shot 5-shot settings, following Muennighoff' et (2023). Here, present the worst and best model performances each task, and detailed experimental results for individual models can be found in Appendix G.
Benchmark
Worst Model
Best Model
Social IQA (Sap et al., 2019)
32.4
33.9
1.5
HellaSwag (Zellers et al., 2019)
33.0
43.4
10.4
PiQA (Bisk al., 2020)
60.2
69.0
8.8
OpenBookQA (Mihaylov al.. 2018)
25.8
31.2
5.4
Lambada (Paperno al., 2016)
18.9
33.5
SeiQ (Welhl ct al., 2017)
76.7
82.9
6.2
ARC Easy (Clark et al., 2018)
44.9
52.2
7.3
COPA (Sarlin al., 2020)
61.5
70.5
9.0
RACE (Lai et al., 2017)
27.9
32.5
LogiQA (Liu et al., 2020)
23,2
27.7
4.5
QQP (Wang ct al., 2018)
48.0
59.7
11.7
WinoGrande (Sakaguchi al., 2021)
50.3
53.2
2.9
MultiRC (Khashabi al., 2018)
47.6
55.7
8.1
Average Performance
43.7
47.9
4.2
Downstream Performance
Downstream Performance
OpenBeoiOO
WaoGGrande
Open
WisoGrande
Helaswag
Lanbols
ARC Basy
Social Social 1ox
LopiQA MulaRC
HellsSing
Lumbals
ARC Bay
Social 1QA
LogiQA
MuiRC
PIçA
RACE
COPA
QOP
PiQA
RACH
CORA
ssiQ
QOP
Validation Loss
ArXiv
Validation Loss
link mokk.spninge..com
Freel.w
mail-archives, spache org
PuhMed Central
palents googla
Wiliieddie
youtub..com
Gibbhh
www Latimes
Stack Exchange
www. qut. comn
coplo.com
(a) Correlation between validation loss by domains of the Pile and downstream performance.
(b) Correlation between validation loss by URL domain within the Pile-CC subset and downstream performance.
Figure 5: The correlation between validation losses across domains and downstream performance for the 64x IB models. Note that we take the negative of the loss value when calculating the correlation, this makes the visualization more intuitive. The same applies for Figure 7.
5.1 DATA MIXTURE SIGNIFICANTLY IMPACTS DOWNSTREAM PERFORMANCE
Initially, we train 64 models, each with 1B parameters, using different data mixtures. Every model is trained on 25B tokens from the Pile dataset (Gao et al., 2021), with tokens allocated based on their corresponding domain weights. Table 3 presents the performance of the worst and best models each downstream task. The reported performance is the average from 0-shot to 5-shot evaluations, scored using the Im-eval- -aaa-hhneeess evaluation framework (Gao et al., 2023; Biderman et al., 2024). We find that the data mixture significantly impacts downstream performances, with the largest performance A reaching 14.6 on the Lambada task. This underscores the importance of studying the optimal data mixture.
5.2
WEB CORPORA BENEFITS DOWNSTREAM PERFORMANCE THE MOST
Next, we visualize the correlation between the validation losses of our 64 IB models across different domains and their performance on various downstream tasks in Figure 5 (a). Prior visualization, hypothesized that the validation loss on the Wikipedia (en) subset would exhibit a strong correlation with most downstream tasks, as it high-quality dataset, and many downstream tasks are derived
'e the token quantity such that is compute-optimal according Chinchilla (Hoffmann al., 2022).
7

---

Published as conference paper at ICLR 2025
Table 4: Performance comparison of different data selection methods. Human refers the weights put forth The Pile (Gao et al., 2021). Pile-CC to only training on the Pile-CC component, PPL to using the perplexity filtering methods from Ankner et al. (2024), ODM the weights from Albalak et al. (2023) and DoReMi to the weights from Xie et al. (2023a). The reported performance on each task is the average score from O-shot to 5-shot settings, while the highest score on each task is highlighted in bold. We estimate the compute (measured in FLOPs) required to arrive at the training data mixture.
Benchmark
Human
DoReMi
PPL
ODM
Pile-CC
REGMIX
Social IQA (Sap al., 2019)
33.6
33.4
33.3
33.7
33.2
33.8
HellaSwag (Zellers al., 2019)
37.4
43.4
43.1
37.2
44.1
44.2
PiQA (Bisk et al., 2020)
65.0
68.3
68.5
64.4
69.2
69.3
OpenBookQA (Mihaylov et al., 2018)
28.2
30.3
30.3
30.0
31.1
30.3
Lambada (Paperno et al., 2016)
29.8
32.1
35,4
29.6
33.2
34.2
SciQ (Welbl al., 2017)
80.1
81.6
78.6
79.8
81.8
82.8
ARC Easy (Clark et al., 2018)
49.4
50.6
50.5
47.9
51.8
51.7
ARC Challenge (Clark et al., 2018)
26.3
26.1
25.9
25.6
26.7
25.7
COPA (Sarlin et al.. 2020)
66.7
68.5
69.2
68.2
65.8
70.2
RACE (Lai al., 2017)
29.0
31.3
31.5
29.7
31.8
31.3
LogiQA (Liu al., 2020)
25.5
26.4
27.5
25.6
27.6
QQP (Wang al., 2018)
52.4
56.6
50.0
53.1
57.0
58.3
WinoGrande (Sakaguchi et al., 2021)
53.1
52.2
52.8
51.8
52.1
53.1
MultiRC (Khashabi et al., 2018)
54.3
53.8
50.4
53.3
50.3
51.7
Estimated FLOPs
0
3.7c19
1.8c19
0
0
3.5c18
Average Performance
45.1
46.8
46.2
45.0
46.8
47.3
Best On
2/ 14
0/14
1/14
0/ 14
5 14
7/ 14
from Wikipedia text. Similarly, previous work often takes WikiText (Merity et al., 2016) as standard benchmark to indicate language model performance.
However, surprisingly, the validation loss on the Pile-CC dataset shows the strongest correlation with most downstream tasks. For instance, the correlation coefficient between the HellaSwag task and the Pile-CC validation loss is remarkably close to 1.0. This unexpected result challenges the conventional assumption that WikiText the most representative dataset for evaluating LLMs. Furthermore, this result aligns with the findings of previous studies (Gadre et al., 2024; Huang et al., 2024), which discovered that the validation loss on the web dataset closely relates to downstream performance.
Moreover, analyze the correlation between the loss of models on the C4100Domain validation set (Magnusson et al., 2023), which taken from the C4 dataset (Raffel et al., 2019) and supposed to share similar distribution Pile-CC since they all derived from the CommonCrawl corpus. Since CommonCraw a collection of diverse domains, we would expect the correlation between the loss of each domain and the downstream tasks vary. However, surprisingly more than 85% of the domains exhibit a very strong correlation with Pile-CC (full correlation graph in Appendix D). This is exempli- fied by the www ign. com domain. which closely mirrors the overall correlation graph of Pile-CC, as illustrated in Figure 5 (b). It also suggests that the high correlation between Pile-CC and downstream task performance may be attributed to diverse coverage across various topics and domains.
5.3 DATA MIXTURE BY REGMIX IMPROVES DOWNSTREAM PERFORMANCE
Previous work has shown that the data mixture method can accelerate LLM pre-training by achieving smaller validation loss (or perplexity) using less training tokens (Xie et al., 2023a). However, critical challenge is determining which validation loss to optimize. While the intuitive approach is to minimize loss across all domains, our analysis of IM training logs reveals significant practical challenges. Therefore, instead of pursuing a broad optimization strategy, we strategically focus on minimizing validation loss Pile-CC, which allows for meaningful progress by leveraging Pile-CC which most consistently predicts overall model performance on downstream tasks.
We implement two approaches to determine the data mixture. The first approach relies on human intuition. Since Pile-CC and own distribution should be the closest match, we hypothesized that pre-training solely Pile-CC might yield better performance than baselines. The second approach leverages REGMIX, using the Pile-CC validation loss as the target variable. We employed LightGBM predict the data mixture which can minimize the Pile-CC validation loss.
8

---

Published as a conference paper at ICLR 2025
1.00
4.50
Validation Loss
3.75
Validation Loss
4.25
3.50
4.00
Human DoReMi
3.75
3.25
3.50
REGMIX
3.00
3.25
2.75
3.00
5
9
13
17
21
25
5
9
13
17
21
25
Training Tokens (B)
Training Tokens (B)
Figure 6: Left: The validation loss on Pile-CC of different methods with Pile-CC in the pre-training corpus. Right: The validation loss on Pile-CC excluding Pile-CC in the pre-training.
Domain Weight
Domain Weight
1a Aoman 13770
gott redl
Soil bedeape
aesembly
Callam
Validation Loss
A-dr
Validation Loss
IMAMad Chemn
commtio--lsp
spp
Ds Mattemaics
Btak Budange
haskel
0
Javascrips
Cherr RCC.
jupyict- acnilr-cccppaa
kolin
Paskied aberace
python
Figure 7: The visualization of correlations between different target domain validation losses and training domain weights using the linear regression model. Left is the Pile dataset, and Right is on the Stack dataset. A high correlation indicates that increasing the training domain weight has a positive impact reducing the target domain validation loss.
We evaluate our proposed approaches against robust benchmarks, including human-curated selections for the Pile corpus (Gao et al., 2021), a sample-level perplexity- based filtering method (PPL) (Ankner al., 2024), an online group-level method Online Data Mixing (ODM) (Albalak et al., 2023), and the flagship group-level method DoReMi (Xie et al., 2023a). For ODM and DoReMi, we obtain the data mixture directly from their reported best domain weights and re-normalize across the available 17 domains. This may result in sub-optimal performance for them compared to the originally reported results. As shown in Table 4, both Pile-CC Only and REGMIX demonstrate strong performance compared to the baselines. On the widely used HellaS wag benchmark, REGMIX shows an improvement of 6.8 over Human selection. Additionally, REGMIX beats all other three methods the task performance in 7 out of 14 cases and yields the highest average score. The surprisingly strong performance of Pile-CC Only reinforces the conclusion from our previous section: web corpora benefits on downstream performance. Finally, REGMIX surpasses the Best Model in Table 3, demonstrating that our automatic data mixture approach is more efficient than random search.
While the Pile-CC validation loss is informative indicator for downstream performance, it may not generalize to every task of interest. Sometimes we may not be able to assume that the validation set stems from similar data distribution as the training set, but rather face an out-of-distributio scenario. To verify the effectiveness of our method in distt--distiiuution scenarios, we fully exclude the Pile-CC domain from the pre- pre-training corpus and use the remaining domains to find the optimal data mixture that minimizes Pile-CC validation loss. As illustrated in Figure 6 (right), our proposed method still outperforms baseline approaches. This demonstrates that REGMIX is robust regardless of whether the target domain is in- or distriiiiiiii..nnn.. We additionally provide the results of regression evaluation under this setting in Figure 6.
5.4 DOMAIN INTERACTIONS ARE CHALLENGING FOR HUMANS TO UNDERSTAND
To understand the impact of different domains on each other, we visualize the coefficients (w) of the linear regression model in Figure 7. The visualization provides insights into how the various data domains contribute to the others, revealing complex interactions among them. We also display code correlation diagrams for each IM code model trained The Stack dataset (Kocetkov et al.,

---

Published as a conference paper at ICLR 2025
Table 5: Performance Comparison 100 Domains for 1M and 60M Models. This table compares the rank correlation (p, higher is better) and mean squared error (MSE, lower is better) for linear and LightGBM regression models. We train 1,000 runs for IM models, fit the regression model, and verify the rank correlation between the ranking predicted by the regression model and the ground-truth ranking of 64 unseen data mixtures on both IM and 60M models.
Test On
IM
60M
Method
P (1)
MSE (L)
p (1)
Linear
90.33
0.12
88.64
LightGBM
99.53
0.02
98.80
2022). Surprisingly, both the domain interaction visualization and the code correlation diagrams display complex relationships that are difficult for human experts to fully comprehend. For example, the PhilPapers domain in the Pile dataset appears to provide gains for all other domains under the linear regression modeling, which a non-obvious finding that challenges intuitive human understanding. These visualizations highlight the inherent complexity in determining the optimal data mixture, underscoring the value of our automated REGMIX approach in efficiently identifying nigh-performing mixtures, rather than relying solely on human intuition.
5.5
DATA MIXTURE EFFECTS TRANSCEND SCALING LAWS
Recent research (Ye et al., 2024; Ge et al., 2024) has demonstrated the feasibility of scaling laws for data mixture. However, our findings in Section 5.4 suggest that the relationship between domain weights and validation loss is more complex than scaling laws might imply. To visualize this complexity, we plotted all experimental points of our 1M training logs in Figure 9. If the scaling law of data mixture held true, we would expect domains. However, our results reveal a more nuanced picture, For example, the DM Mathematics domain, possibly due its distinct distribution compared to other domains, exhibits a near log-log linear relationship between loss and domain weight. In contrast, for most domains like Pile-CC show more complex patterss, where predicting validation loss is non-trivial. As shown, domain interactions appear to be intricate, making it challenging to predict the validation loss for domain based solely on its weight the mixture. These findings suggest that while scaling laws provide valuable insights, they may not fully capture the intricacies of data mixture dynamics. Our approach addresses the challenge by modeling the entire data mixture as input for the regression model, providing a more comprehensive framework for predicting the validation loss while simultaneously accounting for all domain weights.
clear log-log linear relationship across all
5.6 EXTEND REGMIX TO 100 DOMAINS
To demonstrate REGMIX's 's scalability, we conducted preliminary experiments with 100 finer-grained domains. One key challenge lies clustering web content into meaningful domain representations. Intuitively, we define domains by base URLs from the Web dataset (Penedo et al., 2024a), chosen based on token availability. Example domains include articles. latimes.com blogs.wsj.com coj. en eniiiiiedia..or, everything2.com and techerunch.com.
We train 1,000 small-scale models (1M parameters) across different data mixtures, use the training runs fit regression model and then predict the data mixture for models with 1M and 60M parameters. Rank correlation (p) and mean squared error (MSE) were evaluated for both linear and LightGBM regression models as shown in Table 5. The results demonstrate the effectiveness of REGMIX when extending to 100 domains.
CONCLUSION
In this paper, we present REGMIX, novel approach for automatically selecting high-performin data mixtures for pre-training large language models. By formulating the data mixture problem as regression task, REGMIX trains small models to predict the impact of different mixtures, enabling efficient identification of the optimal combination. We demonstrate REGMIX's 's effectiveness by predicting the best data mixture among 64 IB 1B-parameter models. Our large-scal study provides insights into data mixture impacts, the relationship between loss and downstream performance, and the challenges faced by human experts in determining optimal data mixtures.
10

---

Published as a conference paper at ICLR 2025
ETHICS STATEMENT
Optimizing the data mixture for LLM pre-training raises several ethical issues. First, the optimized data mixture might be biased toward certain domains, which is good for achieving better performance. However, certain domains might be underrepresented or misrepresented, leading the trained models to perform poorly or produce biased results for these domains. Second, though our method aims to optimize the data mixture efficiently, searching for the optimal data mixture still requires computa- tional resources, leading to high energy consumption and environmental impact. worthwhile to explore how to further reduce the computation cost.
REFERENCES
Meta AI. Introducing meta llama 3: The most capable openly available Ilm date. //ai.meet wwwpb/ogoo////eeeeeeoeoocoocooce a 1-lma-aaa--- 3/, April 2024.
https:
Alon Albalak, Liangming Pan, Colin Raffel, and William Yang Wang. Efficient online data mixing for language model pre-training. arXiv preprint arXiv:2312.02406 2023.
Alon Albalak, Yanai Elazar, Sang Michael Xie, Shayne Longpre, Nathan Lambert, Xinyi Wang, Niklas Muennighoff, Bairu Hou, Liangming Pan, Haewon Jeong, al. A survey on data selection for language models. arXiv preprint arXiv:2402. 16827, 2024.
Zachary Ankner, Cody Blakeney, Kartik Sreenivasan, Max Marion, Matthew L. Leavitt, and Mansheej Paul. Perplexed by perplexity: Perplexity-based data pruning with small reference models. CoRR, abs/2405 20544,, 2024. doi: 10.48550/ARXI..445520000 48550/arXiv. wwt.rxiii......... 2tt.il......5........... 20541.
URL https: //toii///coo.oee org/10.
Stella Biderman, Hailey Schoelkopf, Lintang Sutawika, Leo Gao, Jonathan Tow, Baber Abbasi, Alham Fikri Aji, Pawan Sasanka Ammanamanchi, Sidney Black, Jordan Clive, et al. Lessons from the trenches on reproducible evaluation of language models. arXiv preprint arXiv: 14782, 2024.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. Piqa: Reasoning about physical commonsense natural language. In Proceedings of the AAAI conference on artificial intelligence, 2020.
Sid Black, Stella Biderman, Eric Hallahan, Quentin Anthony, Leo Gao, Laurence Golding, Horace He, Connor Leahy, Kyle McDonell, Jason Phang, et al. Gpt-neox-20b: An open-source autoregressive language model. arXiv preprint arXiv:2204.06745 2022.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agar- wal, Ariel Herbert-Vore Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Ma- teusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCan- dlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Ad- vances in Neural Information Processing Systems, volume 33, pp. 1877-1901. Curran Asso- ciates, Inc., 2020. URL httpp::::ppcceeepiittttt http://pooeeeiiig..ere.o...... ce cc/paper files/papez wtt.f/eeee/5500dbbbbbbbeeeee 3aw.beFebbbb44444848bb3eec www Pat.bbaacc.22f444-apr..r..
Alexander Bukharin and Tuo Zhao. Data diversity matters for robust instruction tuning. arXiv preprint (tt.:23111.1473:30 2023.
Mayee Chen, Nicholas Roberts, Kush Bhatia, Jue Wang, Ce Zhang, Frederic Sala, and Christopher Ré. Skill-it! data-driven skills framework for understanding and training language models. arXiv preprint arXiv:2307. 14430, 2023.
Xuxi Chen, Zhendong Wang, Daouda Sow, Junjie Yang, Tianlong Chen, Yingbin Liang, Mingyuan Zhou, and Zhangyang Wang. Take the bull by the horns: Hard reweihhreweighted continual training improves IIm generalization. arXiv preprint arXiv:2402. 14270, 2024.
11

---

Published as a conference paper at ICLR 2025
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. 2018.
arXiv preprint arXh::1803.0550
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers solve math word problems. CoRR, abs/2110.14165 2021. URL attps:://axiiss wwwpariwii....ioaoooooooooooooot alss 2110 14168.
Devleena Das and Vivek Khetan. Deft: Data efficient fine-tuning for large language models via unsupervised core-set selection. arXiv preprint arXiv:2310. 2023.
Longxu Dou, Qian Liu, Guangtao Zeng. Jia Guo, Jiahui Zhou, Wei Lu, and Min Lin. Sailor: Open language models for south-east asia. CoRR, abs/2404.03600, 2024. doi: 10:48550/ARXIV.4404.
03608, URL https: /tt:: org 110 411.1..555/00111111. http::///......r.........//5//// 2404.03608.
Logan Engstrom, Axel Feldmann, and Aleksander Madry. Dsdm: Model-aware dataset selection
with datamodels. arXiv preprint arXiv:2401.129220 2024.
Simin Fan, Matteo Pagliardini, and Martin Jaggi. Doge: Domain reweighting with generalization estimation. arXiv preprint 0rXiv:2310.15300 2023.
Samir Yitzhak Gadre, Georgios Smyriss, Vaishaal Shankar, Suchin Gururangan, Mitchell Wortsman, Rulin Shao, Jean Mercat, Alex Fang, Jeffrey Li, Sedrick Keh, Rui Xin, Marianna Nezhurina, Igor Vasiljevic, Jenia Jitsev, Alexandros G. Dimakis, Gabriel Ilharco, Shuran Song, Thomas Kollar, Yair Carmon, Achal Dave, Reinhard Heckel, Niklas Muennighoff, and Ludwig Schmidt. Language models scale reliably with over over- training and on downstream tasks. CoRR, 0bss/2003.085400 2024. doi: 0044850//RXXI1.240.00.8 085400. URL https: wwwpdoooo.ogo.oooooooooeoooooooe 48550/ar2310 2403. 08540.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The pile: An 800gb dataset of diverse text for language modeling. CoRR, abs/2101.0002 2021. URL https: //rxivi///xvvv.///////a org/ abs/ 2101. 00027.
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman. Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, h. Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang. Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang. Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 12 2023. URL https: zewodo. o org/ http///eordd//////
Ce Ge. Zhijian Ma, Daoyuan Chen, Yaliang Li, and Bolin Ding. Data mixing made efficient: A bivariate scaling law for language model pretraining. arXiv preprint arXiv:2405. 2024.
Sachin Goyal, Pratyush Maini, Zachary C. Lipton, Aditi Raghunathan, and J. Zico Kolter. Scaling laws for data filtering data curation cannot be compute agnostic. CoRR, abs/2404.07177, 2024. doi: 10.48550/AAXX1.24444000 07177.
URL https: /tt.//ooiiio// oro/ 10. 48w50/arrii..22.. 2404.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL https:://// /ttp:://///peeeeeieeeeo net forum? didim? jm jmI 3GmQ.
Jordan Hoffimann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, al. Training compute-optimat large language models. arXiv preprint 22.3:::2203.155500 2022.
Yuzhen Huang, Jinghan Zhang, Zifei Shan, and Junxian He. Compression represents intelligence linearly, arXiv preprint 4404v:2404. 09937, 2024.
Feiyang Kang, Hoang Anh Just, Yifan Sun. Himanshu Jahagirdar, Yuanzhi Zhang, Rongxing Du, Anit Kumar Sahu, and Ruoxi Jia. Get more for less: Principled data selection for warming up fine-tuning in lms.. arXiv preprint arXiv:2405.02774 2024.
12

---

Published as a conference paper at ICLR 2025
Guolin Ke, Qi Meng, Thomas Finley, Taifeng Wang, Wei Chen, Weidong Ma, Qiwei Ye, and Tie- Yan Liu. Lightgbm: highly efficient gradient boosting decision tree. Advances neural information processing systems, 30, 2017.
Daniel Khashabi, Snigdha Chaturvedi, Michael Roth, Shyam Upadhyay, and Dan Roth. Looking beyond the surface: A challenge set for reading comprehension over multiple sentences. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Com- putational Linguistics: Human Language Technologies, Volume (Long Papers), pp. 252-262, 2018.
Denis Kocetkov, Raymond Li, Loubna Ben Allal, Jia Li, Chenghao Mou, Carlos Muñoz Ferrandis, Yacine Jernite, Margaret Mitchell, Sean Hughes, Thomas Wolf, Dzmitry Bahdanau, Leandro von Werra, and Harm de Vries. The stack: TB of permissively licensed source code. CoRR, abs/2211.15533, 2022. doi: 10.4850/ARXI.2211.5555 URL https://dee org/10. 48550/arXio 2211 112211115555 15533.
Pang Wei Koh and Percy Liang. Understanding black-box predictions via influence functions. In International conference on machine learning, pp. 1885-1894, PMLR, 2017.
Guokun Lai, Qizhe Xie, Hanxiao Liu, Yiming Yang, and Eduard Hovy. Race: Large-scale reading comprehension dataset from examinations. arXiv preprint arXiv: 0ttl::/704.400000 2017.
Jeffrey Li, Alex Fang, Georgios Smyrnis, Maor Ivgi, Matt Jordan, Samir Gadre, Hritik Bansal, Etash Guha, Sedrick Keh, Kushal Arora, Saurabh Garg, Rui Xin, Niklas Muennighoff, Reinhard Heckel, Jean Mercat, Mayee Chen, Suchin Gururangan, Mitchell Wortsman, Alon Albalak, Yonatan Bitton, Marianna Nezhurina, Amro Abbas, Cheng- Yu Hsich, Dhruba Ghosh, Josh Gardner, Maciej Kilian, Hanlin Zhang, Rulin Shao, Sarah Pratt, Sunny Sanyal, Gabriel Ilharco, Giannis Daras, Kalyani Marathe, Aaron Gokaslan, Jieyu Zhang, Khyathi Chandu, Thao Nguyen, Igor Vasiljevic, Sham Kakade, Shuran Song, Sujay Sanghavi, Fartash Faghri, Sewoong Oh, Luke Zettlemoyer, Kyle Lo, Alaaeldin El-Nouby. Hadi Pouransari, Alexander Toshev, Stephanie Wang. Dirk Groeneveld, Luca Soldaini, Pang Wei Koh, Jenia Jitsev, Thomas Kollar, Alexandros G. Dimakis, Yair Carmon, Achal Dave, Ludwig Schmidt, and Vaishaal Shankar. Datacomp-Im: search of the next generation of training sets for language models, 2024. URL https: / aarivvvivv. org/ abs 2406 11794.
Zhenghao Lin, Zhibin Gou, Yeyun Gong, Xiao Liu, Yelong Shen, Ruochen Xu, Chen Lin, Yujuu Yang, Jian Jiao, Nan Duan, and Weizhu Chen. Rho-1: Not all tokens are what you need. arXiv preprint ar arXiv:2404.07965 2024.
Jian Liu, Leyang Cui, Hanmeng Liu, Dandan Huang, Yile Wang, and Yue Zhang. Logiqa: A challenge dataset for machine reading comprehension with logical reasoning. arXiv preprint arXiv:2007.08/22 2020.
Wei Liu, Weihao Zeng, Keqing He, Yong Jiang, and Junxian He. What makes good data for alignment? a comprehensive study of automatic data selection in instruction tuning. In The International Conference on Learning Representations, 2024.
lan Magnusson, Akshita Bhagia, Valentin Hofmann, Luca Soldaini, Ananya Harsh Jha, Oyvind Tafjord, Dustin Schwenk, Evan Pete Walsh, Yanai Elazar, Kyle Lo, Dirk Groeneveld, Iz Beltagy, Hannaneh Hajishirzi, Noah A. Smith, Kyle Richardson, and lesse Dodge. Paloma: A benchmark for evaluating language model fit. CoRR, abs/2312.10523,, 2023. doi: 1044550/ARXXI22112..0022 URL https:://// http//oi./oo////////// /10. 48554 http/.iiviiiveee. 2tt..i....11.112.....
Max Marion, Ahmet Üstün, Luiza Pozzobon, Alex Wang. Marzieh Fadaee, and Sara Hooker. When less is more: Investigating data pruning for pretraining lims at scale, 2023.
Sachin Mehta, Mohammad Hossein Sekhavat, Qingqing Cao, Maxwell Horton, Yanzi Jin, Chenfan Sun, Iman Mirzadeh, Mahyar Najibi, Dmitry Belenko, Peter Zatloukal, and Mohammad Rastegari. Openelm: An efficient language model family with open training and inference framework. CoRR, abs/2404. 14619, 2024. doi: 408550/ARXXV.2404, 14619. URL https: / //o... .org/10. 48550/a www.. 2404.4. 14619.
13

---

Published as a conference paper at ICLR 2025
Dheeraj Mekala, Alex Nguyen, and Jingbo Shang. Smaller language models capable of selecting instruction-tuma titiionnttnnnggggo training data for larger language models. arXiv preprint 3rXiv:2402.10430 2024.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering. arXiv preprint wrXiv:1809002789 2018.
Sören Mindermann, Jan M Brauner, Muhammed T Razzak, Mrinank Sharma, Andreas Kirsch, Winnie Xu, Benedikt Höltgen, Aidan N Gomez, Adrien Morisot, Sebastian Farquhar, et al. Prioritized training on points that are learnable, worth learning, and not yet learnt. International Conference on Machine Learning. pp. 15630-15649, PMLR, 2022.
Niklas Muennighoff, Alexander M Rush, Boaz Barak, Teven Le Scao, Nouamane Tazi, Aleksandra Piktus, Sampo Pyysalo, Thomas Wolf, and Colin Raffel. Scaling data-constrained language models. In Teirty-seventhh Conference on Neural Information Processing Systems, 2023, URL https: http://ppeeeeiieeeeee net/ forum? ?iceummmiddipe -iddddidddjrsraae J5BuTrE335.
Rohan Pandey. gzip predicts data-dependent scaling laws. arXiv preprint arXiv:2405.16688 2024.
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring broad discourse context. arXiv preprint arXiv:1606.0603: 2016.
Guilherme Penedo, Hynek Kydlíček, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, Thomas Wolf, al. The fineweb datasets: Decanting the web for the finest text data at scale. arXiv preprint arXiv:2406. /7577, 2024a.
Guilherme Penedo, Hynek Kydlíčck, Loubna Ben allal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, and Thomas Wolf. The fineweb datasets: Decanting the web for the finest text data at scale, 2024b. URL https: arxivv..vvvv wwweiiiirr///iieeeeoeeeo.eoooore .rgg abs/ 2406. 17577.
Haoran Que, Jiaheng Liu, Ge Zhang, Chenchen Zhang, Xingwei Qu, Yinghao Ma, Feiyu Duan, Zhiqi Bai, Jiakai Wang, Yuanxing Zhang, Xu Tan, Jie Fu, Wenbo Su, Jiamang Wang, Lin Qu, and Bo Zheng. D-ept law: Domain-specific continual pre-training scaling law for large language models, 2024.
Jack W. Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, H. Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, Eliza Rutherford, Tom Hennigan, Jacob Menick, Albin Cassirer, Richard Powell, George van den Driessche, Lisa Anne Hendricks, Maribeth Rauh, Po-Sen Huang. Amelia Glaese, Johannes Welbl, Sumanth Dathathri, Saffron Huang, Jonathan Uesato, John Mellor, Irina Higgins, Antonia Creswell, Nat McAleese, Amy Wu, Erich Elsen, Siddhant M. Jayakumar, Elena Buchatskaya, David Budden, Esme Sutherland, Karen Simonyan, Michela Paganini, Laurent Sifre, Lena Martens, Xiang Lorraine Li, Adhiguna Kuncoro, Aida Nematzadeh, Elena Gribovskaya, Domenic Donato, Angeliki Lazaridou, Arthur Mensch, Jean- Baptiste Lespiau, Maria Tsimpoukelli, Nikolai Grigorev, Doug Fritz, Thibault Sottiaux, Mantas Pajarskas, Toby Pohlen, Zhitao Gong, Daniel Toyama, Cyprien de Masson d' Autume, Yujia Li, Tayfun Terzi, Vladimir Mikulik, Igor Babuschkin, Aidan Clark, Diego de Las Casas, Aurclia Guy, Chris Jones, James Bradbury, Matthew J. Johnson, Blake A. Hechtman, Laura Weidinger, lason Gabriel, William Isaac, Edward Lockhart, Simon Osindero, Laura Rimell, Chris Dyer, Oriol Vinyals, Kareem Ayoub, Jeff Stanway, Lorrayne Bennett, Demis Hassabis, Koray Kavukcuoglu, and Geoffrey Irving. Scaling language models: Methods, analysis & insights from training gopher. CoRR, abs/2112.1 1446, 2021. URL https: /atxxvxiivvv.///// wwrxii.oor///eeo abs 221222221111 11446.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv e-prints, 2019.
14

---

Published as a conference paper at ICLR 2025
Noveen Sachdeva, Benjamin Coleman, Wang-Cheng Kang, Jianmo Ni, Lichan Hong, Ed H Chi, James Caverlee, Julian McAuley, and Derek Zhiyuan Cheng. How to train data-efficient 1lms. arXiv preprint arX1v:2402.09668, 2024.
Ritik Sachin Parkar, Jachyung Kim, Jong Inn Park, and Dongyeop Kang. Selectllm: Can Ilms select important instructions to annotate? arXiv e-prints, pp- arXiv-2401, 2024.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99-106, 2021.
Maarten Sap, Hannah Rashkin, Derek Chen, Ronan LeBras, and Yejin Choi. SocialIQA: Common- sense reasoning about social interactions. In EMNLP, 2019.
Paul-Edouard Sarlin, Daniel DeTone, Tomasz Malisicwicz, and Andrew Rabinovich. Superglue: Learning feature matching with graph neural networks. In Proceedings the IEEE/CVF conference on computer vision and pattern recognition, pp. 4938-4947 2020.
Yunfan Shao, Linyang Li, Zhaoye Fei, Hang Yan, Dahua Lin, and Xipeng Qiu. Balanced data sampling for language model training with clustering. arXiv preprint arXiv:2402. 14526, 2024.
Vasu Sharma, Karthik Padthe, Newsha Ardalani, Kushal Tirumala, Russell Howes, Hu Xu, Yao Huang, Shang- Wen Li, Armen Aghajanyan, and Gargi Ghosh. Text quality- pruning for efficient training of language models. arXiv preprint arXiv:2405.01582. 2024.
Luca Soldaini, Rodney Kinney, Akshita Bhagia, Dustin Schwenk, David Atkinson, Russell Authur, Ben Bogin, Khyathi Chandu, Jennifer Dumas, Yanai Elazar, et al. Dolma: An open corpus of three trillion tokens for language model pretraining research. arXiv preprini arXiv:2402.001550 2024.
Alon Talmor, Jonathan Herzig, Nicholas Lourie, and Jonathan Berant. CommonsenseQA: A question answering challenge targeting commonsense knowledge. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pp. 4149-4158, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10.18653/v 1421. https : / aa/ aclooohooogy......o wwwplllyyoooooooooo... org/N1 9- 421.
Megh Thakkar, Tolga Bolukbasi. Sriram Ganapathy, Shikhar Vashishth, Sarath Chandar, and Partha Talukdar. Self-influence guided data reweighting for language model pre-training. arXiv preprint arXiv:2311.000 1.00:93,, 2023.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton-Feereer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams. Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurélien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models. CoRR, abs/2307.092888 2023. doi: A485500ARRXIV.2377.02222 URL https: /ooi..i... org/ 10. 48544855//20 http///avvv//..///////.aa/.aaava 2307. 09288.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv.1804.07461. 2018.
Xinyi Wang, Yulia Tsvetkov, and Graham Neubig. Balancing training for multilingual neural machine translation. arXiv preprint arXiv:2004.067493 2020.
15

---

Published as a conference paper at ICLR 2025
Johannes Welbi, Nelson F Liu, and Matt Gardner. Crowdsourcing multiple choice science questions. arXiv preprint (rxXiv:707.062220 2017.
Alexander Wettig, Aatmik Gupta, Saumya Malik, and Danqi Chen. Qurating: Selecting high-quality data for training language models, 2024.
Mengzhou Xia, Mikel Artetxe, Chunting Zhou, Xi Victoria Lin, Ramakanth Pasunuru, Danqi Chen, Luke Zettlemoyer, and Ves Stoyanov. Training trajectories of language models across scales. arXiv preprint arXiv:2212.0980) 02.0903, 2022.
Mengzhou Xia, Sadhika Malladi, Suchin Gururangan, Sanjeev Arora, and Danqi Chen. Less: Selecting influential data for targeted instruction tuning. arXiv preprint arXi::440..0333. 2024.
Sang Michael Xie, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu, Percy Liang, Quoc V Le. Tengyu Ma, and Adams Wei Yu. Doremi: Optimizing data mixtures speeds up language model pretraining. arXiv preprint 2rXiv:2305.1042 2023a.
Sang Michael Xie, Shibani Santurkar, Tengyu Ma, and Perey Liang. Data selection for language models via importance resampling. arXiv preprint 23X.:23302.03666 2023b.
Chen Yang, Junzhuo Li, Xinyao Niu, Xinrun Du, Songyang Gao, Haoran Zhang, Zhaoliang Chen, Xingwei Qu, Ruibin Yuan, Yizhi Li, et al. The fine line: Navigating large language model pretraining with down-streaming capability amalysis. arXiv preprint arXiv:2404.01204, 2024a.
Yu Yang, Siddhartha Mishra, Jeffrey N Chiang, and Baharan Mirzasoleiman. Smalttlaagge (s21); Scalable data selection for fine-tuning large language models by summarizing training trajectories of small models. arXiv preprint wrXhv:2403.07384 2024b.
Jiashnng Ye, Peiju Liu, Tianxiang Sun, Yunhua Zhou, Jun Zhan, and Xipeng Qiu. Data mixing laws: Optimizing data mixtures by predicting language modeling performance. CoRR, 2ttss//40..1652... 2024. doi: 10-4550/ARXIV.240...000 2403. 16952.
URL https:/ /oii. orrpiiorro11..00 48550 /0/85550///iiv../
Zichun Yu, Spandan Das, and Chenyan Xiong. Mates: Model-aware data selection for efficient pretraining with data influence models, 2024.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? arXiv preprint arXiv: 905.078300 2019.
Peiyuan Zhang, Guangtao Zeng, Tianduo Wang, and Wei Lu. Tinyllama: An open-source small language model. arXiv preprint arXiv:2401.02338 2024a.
Yifan Zhang, Yifan Luo, Yang Yuan, and Andrew C Yao. Autonomous data selection with language models for mathematical texts. In ICLR 2024 Workshop on Navigating and Addressing Data Problems for Foundation Models, 2024b.
16

---

Published as a conference paper at ICLR 2025
A BROADER IMPACT
In this paper, we propose REGMIX, novel method for optimizing data mixture for language model pre-training. A unique contribution of REGMIX is its novel use of ultra-small proxy models (i.e., IM parameters) to optimize data mixtures for language model pre-training, an approach previously unexplored in the field. This novel use of ultra-small proxy models reduces computational overhead less than 2% of final training costs, dramatically lowering barriers for data mixture research within academic budgets. Through the focus on computational efficiency and our commitment to open science (i.e., with all datasets and trained models publicly available), we believe REGMIX represents a significant advancement toward democratizing research in language model pre-training.
Beyond the methodology contribution, our work delivers several novel empirical insights in data mixture research. We provide the first comprehensive demonstration of significant performance varia- tions across different data mixtures, supported by extensive experiments with 1B IB-parameter models trained on 64 distinct mixtures and rigorously evaluation across 12 benchmarks. Our results establish the superiority of automatic data mixture optimization over human intuition- based approaches, with PhilPapers serving interesting in-depth case study illustrating how domain interactions follow complex patterns that transcend human intuition.
B LIMITATIONS
Despite making progress in understanding and optimizing data mixtures for better performance, our method still has several limitations.
The rank invaraince assumption. Our investigation of the rank invariance assumption currently focuses on model scales from 1M 1B parameters. While we aimed to verify the hypothesis at larger scales, establishing statistically meaningful correlations for 3B models would require training 64 different models with 50B tokens each, equivalent to training one 3B model on 3.2T tokens, which significantly exceeds our computational resources.
The maximum model parameters. We have verified that small models can used to predict the optimal data mixture for large-scale runs with up to 1B parameters. However, much larger models are commonly trained with 7B or 70B parameters (Touvron et al., 2023). Due to compute constraints we leave the verification of REGMIX at larger scales to future work.
The benchmark coverage. Owing to the scarcity of relevant data in the Pile corpus and the relatively small size of our model at 1B scale, their performance on the MMLU benchmark (Hendrycks et al., 2021) is nearly random and negligible on GSMSK (Cobbe et al., 2021). Consequently, we do not compute the correlation between the validation loss and scores these challenging benchmarks.
The infinite data assumption. Most existing data mixing methods assume the availability of unlimited data for each domain. Although we consider this issue in our no Pile-CC experiments in Section 5.3, systematically incorporating the effect of available data into the method remains challenging. Combining our method with the decay coefficient of data proposed in Muennighoff et al. (2023) could be an interesting future work to explore, potentially addressing the limited data availability scenario.
The domain assumption. A common assumption of existing data mixture methods (including ours) is that the domain each example belongs to is known. However, this may not always be the case and the domain needs to be obtained first. Assigning examples to domains is a hard task, which may make it challenging to apply our methods when the domain boundaries unclear.
The tokenizer assumption. All existing data mixture methods require the use of proxy models to obtain domain weights However, fundamental assumption of these methods is that the proxy model uses the same tokenizer and vocabulary size as the large model. Generalizing weights across different tokenizers poses significant challenges.
17

---

Published as a conference paper at ICLR 2025
6.4
6.4
Correlation: 0.
Correlation:
6.2
6.2
6.0
6.0
True Loss
Truc Loss
5.8
5.8
5.6
5.4
5.4
5.2
5.2
5.0
5.0
4.8
4.8
5.0
6.2
4.8
5.2
5.8 6.0 Pred Loss
Pred Loss
Figure 8: The visualization of loss prediction on small models (e.g., 1M parameters). Left: The scatter plot of predicted and true loss pairs of Linear model. Right: The scatter plot of predicted and true loss pairs LightGBM model.
C ADDITIONAL RESULTS
C.1 THE REGRESSION PREDICTION VISUALIZATION
As shown in Figure 8, we visualize the predicted and true loss pairs of the linear model and LightGBM model on the IM models. The LightGBM model performs better than the linear model, achieving near 100% Spearman Rank Correlation p.
C.2 LOSS AND RANK PREDICTION ON SMALL MODELS FOR OUT-OF-DISTRIBBTTIO SETTING
In Section 5, we verify the effectiveness of our method in out-of-distribution scenarios where fully exclude the Pile-CC domain from the pre-training corpus and use the remaining domains to find the optimal data mixture that minimizes Pile-CC validation loss. We also provide the results of regression evaluation under this setting in Figure 6. Similarly, LightGBM model outperforms the linear model and achieves nearly 100% Spearman Rank Correlation p.
Table 6: The regression model is fitted using the training artifacts of 512x 1M models trained with 1B tokens excluding the Pile-CC domain, and evaluated on unseen data mixtures for 1M parameter models. Pearson's and MSE measure the loss prediction performance, while p compares the predicted and actual ranks.
Test
models with tokens
Method
p
MSE (L)
Linear
83.00
0.08
LightGBM
95.47
0.04
C.3 THE DERIVED DATA MIXTURES
Table 7 presents the derived data mixture weights for different methods. As illustrated, REGMIX assigns high weight of 0.87 to the Pile-CC dataset, aligning with human intuition.
18

---

Published as a conference paper at ICLR 2025
Table 7: The domain weights of different methods. In experiments, DoReMi refers the reported best reference model with 280M parameters and its corresponding domain weights. 'Note that the domain weights of Human, DoReMi and Online are re-normalized from the weights reported in Xie et al. (2023a) to adapt them to the available domains. The DoReMi weight are derived from the hest-performing configuration obtained using a 280M parameter model. The Online weight is derived from the final domain weights obtained by the method.
Domain Weights
Human
DoReMit
Online
Pile-CC
REGMIX
ArXiv
0.134
0.004
0.0267
0.0
0.001
FreeLaw
0.049
0.005
0.0346
0.0
0.001
NIH ExPorter
0.007
0.008
0.0466
0.0
0.001
PubMed Central
0.136
0.006
0.0316
0.0
0.003
Wikipedia (en)
0.117
0.086
0.0504
0.0
0.016
DM Mathematics
0.025
0.002
0.0168
0.0
0.0
Github
0,054
0.022
0.0155
0.0
0.0
PhilPapers
0,003
0.034
0,0451
0.0
0.0
Stack Exchange
0.118
0.019
0.0353
0.0
0.0
Enron Emails
0.004
0.009
0,0228
0.0
0.002
Gutenberg (PG-19)
0.025
0.009
0.0669
0.0
0.002
Pile-CC
0.142
0.743
0.0894
1.0
0.87
Ubuntu IRC
0.009
0.011
0.0363
0.0
0.064
EuroParl
0.005
0.008
0.0315
0.0
0.0
HackerNews
0.01
0.016
0.0604
0.0
0.012
PubMed Abstracts
0.107
0.014
0.0467
0.0
0.024
USPTO Backgrounds
0.053
0.004
0.0403
0.0
0.002
Table 8: Performance comparison of different data selection methods using LightEval following previous work (Penedo et al., 2024b). Human refers to the weights put forth in The Pile (Gao et al., 2021), Pile-CC only training on the Pile-CC component, and DoReMi the weights from Xie et al. (2023a). The reported performance for each task is the average zero-shol task performance across five different runs, and the standard deviation. We estimate the compute (measured FLOPs) required to arrive at the training data mixture. Scores significantly outperforming the Human baseline for each task are highlighted in bold, with significance determined using Cohen's d.
Benchmark
Human
DoReMi
Pile-CC
REGMIX
ARC Easy (Clark et al., 2018)
45.3 0.4
46.6 4.77
47.1 +06
47.2 0.9
ARC Challenge (Clark et al.. 2018)
25.5 +08
25.9 +08
25.6
25.6
CommonsenseQA (Talmor al., 2019)
11.8
34.1
40.
35,0
HellaSwag (Zellers al., 2019)
36.5
41.5
30.7
42.1
OpenBookQA (Mihaylov et al., 2018)
29.8
31.0
40.4
08
PiQA (Bisk et al., 2020)
65.4406
68.7 +..
69.0
6..4
Social IQA (Sap al.. 2019)
41.7
42.0 0.2
42.7 0.03
42,6 0.7
WinoGrande (Sakaguchi et al., 2021)
51.1
51.2
50.7
$0.9
MMLU (Henndrycks al., 2021)
28.6
28.9
28.5
28.7
Average Performance
39.5 10.3
41.1
4..3
41.5 1.2
Beat Human on
5/9
6/9
6/9
Estimated FLOPs
0
3.7e19
0
3.5e18
C.4 THE EVALUATION RESULTS USING LIGHTEVAL
Following the approach of Fine Web (Penedo et al., 2024b), employ the LightEval 'ibbaary to evaluate our models using a suite of benchmarks selected for their stability and suitability. The chosen benchmarks exhibit three key characteristics: low score variance across different data samples, monotonic score improvement during training, and above-random baseline scores for models in the IB parameter range. Table presents the evaluation results. Our method, REGMIX, consistently
'https: pettss://iiii http://gggttuuuuue thub. www.:hw.......gggooooeoocceeeeee / 1ighteeeeaahttt g
19

---

Published as a conference paper at ICLR 2025
ArXiv
PabMed Central
Logt Validattion Los))
Log Validastion Los)
80
Log Validation
1.6
2.0
1.5
10
10
10
107
10
10
102
10
10
Domain Weight
Domain Weight
Domain Weight
Witipedia (en)
Mathematics
Gishrab
Log LV aliihation Lows)
Log(Validation Laws)
2.0
1.5
Log(V Log[Validation alldattion
2.8
L25
2.5
1.00
1.0
B75
10
10
10
10
10
10
107
10
10
Domain Weight
Domain Weight
Domain Wright
Exchongge
Guiemberz (PG-19)
Pi--0C
Log Log Valisiation Lons)
Lag Vaaiiiation Loss)
1.
Log( Log(Validatte Validation Loss)
2.5
7.7
10
20
10
107
10
10
10
10
10
10
Domsin Weight
Domain Weight
Domain Weight
Ubarta IRC
HackerNews
PulMed Abstracts
Log(V Log(Valddatinn Validatioon Lons)
Log Validation Loss)
Log V Validation Loes)
23
1.9
1.7
10
10
10
10 Weight
10
10
10
17
Domain Weight
Domain
Domain
Weight
Figure 9: The visualization of 1M training logs across various data mixtures. The x-axis represents the weight of each domain in the data mixture and the y-axis shows the log value of validation loss for that domain. As seen, predicting the validation loss solely based on the domain weight challenging.
outperforms the Human baseline on 6 benchmarks. Moreover, REGMIX demonstrates superior average performance compared to the DoReMi and the Pile-CC Only methods.
C.5 DATA MIXTURE EFFECTS TRANSCEND SCALING LAWS
As shown in Figure 9, validation loss does follow simple log-log linear trend with domain weight, as scaling laws suggest. While DM Mathematics shows near-linear behavior, most domains, like Pile-CC, exhibit complex interactions that challenge straightforward predictions. This highlights the role of domain interactions in shaping loss, making individual weights insufficient for accurate modeling. Our approach addresses this by using the entire data mixture as input a regression model, capturing cross-domain effects for better loss prediction.
20

---

Published as a conference paper at ICLR 2025
D URL DOMAIN CORRELATION GRAPH
Downstream Performance
OpenBook
WinoGrande
HellaSwag
ARC Easy
Social IQA
Lambada
LogiQA
MultiC
PIQA
RACE
COPA
SeiQ
QQr
www.nytimes.com
en.wikipedia.org
www.latimes.co com
com
www. www.theguardian.com
www.huflpost.com
patents google.com com
www.businessissddrr.comm
www.forbes.com
www.eventbrilc.com
link. sonksspringer.cmm
Validatioo Loss
w www.chhiagoorrbumc.com
www www.ffxxneve www.foxnews.com.
www.aljazeera.co com
www.dailmmail.comm
www.ncbi.nlm.ot
www. www.expres www.exprass.coukk
en.m. wn.m.iikipodia..rg
www.enet.com
www.telegraph.co co.uk
www.theat www.theatlanticc com
forums macrumars.comm
www.oreilly.coo com
www.waslingtonpost.com
www.zdnet.com
www.foxbusiness.cmm
Figure 10: The visualization of correlations between different URL domains within the C4 subsets and the downstream performance (Part 1).
21

---

Published as a conference paper at ICLR 2025
Downstream Performance
OpenBookQA
WinoGrande
HellaSwag wag
ARC Easy
Social IQA
Lambada
LogiQA
MultiRC
RACE
COPA
PiQA
SciQ
QQP
www.reulets.com
www.ibrimes.co.uk
www.rt.com
www.prweb.com
www.devianlart.com
www.si.com
www.bbc.com.
github.com
nypost.com
Validation Loss
itunes itunes.applc.com
www.instructables.com
www.youtube.com
www.booking.com
www.etsy.com
www.marketwired.com
sitcs
www.haltimoresan.com
www.agrealertown..cmm
www.npr.org
www.fool.com
www.tripadvisoo.com
www.bbc.co.nk
lists.w3.org
mashablc.com
Figure 11: The visualization of correlations between different URL domains within the C4 subsets
and the downstream performance (Part 2).
22

---

Published as a conference paper at ICLR 2025
Downstream Performance
Open
WinoGrande
Selllawwag
ARC. Easy
Social IQA
Lambada
LogiQA
MultR C
PiQA
RACE
COPA
SciQ
QQP
www.cnbc.com
answers.sap.com
homestars,
com com
www.hindustamtimes.comm
www.reference.coom com
www.city-data.ccm
modium. com
app-wwringgdiggram.erokkk com
www.csmonitor, com
Validation Loss
www.adweek.co com
com
docs wocs.microsoft.com com
www.yahoo.com
www.thesun.co.. co.uk
www.nydailynews. com com
www.dailystar.co.uk
fineastamerica.com com
www.kickstarler.com
uk rewtere..coo com
www.insiderpages.co com
www. www.inguisitr.com lists, debian.org www.straitssimes com
www.cbsnews.co com simple, wikipedia.org deadline.com
Figure 12: The visualization correlations between different URL domains within the C4 subsets and the downstream performance (Part 3).
23

---

Published as a conference paper at ICLR 2025
Downstream Performance
OpenBookQA
WinoGrande
HellaSwag
ARC Easy
Social IQA
Lambada
LogiQA
MulaRC
PIQA
RACE
COPA
SeiQ
QQP
www.androidheadlinssccom
www.wired.com
www.bustle.com
premium. wrmuudev.rg
www.lihraytthing.com
aaii-arrcives.aacch..oo
scholars. duke.odu
www.glasddoor.cm
www.peworld.coo com
Validation Loss
www.shutterstock.co com
myemail com
www. www.eventbritc.co..r
www. www.fastcomppnny.com com
www.frsspost.com
www.entrepreeneur.com
www.breitbart.com
techerunch.com
www.ndrv.com
financc.yahoo.com
archives wwt.ssllb.......oooooooooooooooe ma.us
www.gsmarena.com com
www.lonelyplamet.com
www.ign.com
Figure 13: The visualization correlations between different URL domains within the C4 subsets and the downstream performance (Part 4).
24

---

Published as a conference paper at ICLR 2025
Weight
10
IM
60M
10-3
ArXiv
FreeLas
NIH Ex Portor
PulsMed Central
Wilpedii (eu)
Mathenation
Github
PhilPapon apers
Stack Exclange
Enron Emalls
Gutenberg (PG. 19)
Pile- CC
Ubanta IRC
EuroPaa
Hark HarkerNewr KerNeess
PubMed Abstracts
USP Backgrous oundo
Domain
Figure 14: REGMIX yields similar data mixture distributions when using the 1M model and the 60M model proxy models, demonstrating the stability of our method. Note that the y-axis is in log-scale for visualization purpose.
E IMPLEMENTATION DETAILS
We utilize the model architecture proposed by Zhang al. (2024a) and create various model variants by modifying the number of layers, the number attention heads, and the dimensions of token embeddings and hidden states, illustrated in Figure 9. For tokenization, employ the GPTNeoX tokenizer (Black et al., 2022), which has a vocabulary size of 50,432.
For models with 1M and 60M parameters, 1M tokens, which means the training budget 1B tokens. Similarly, we train the larger model with IB parameters with 25000 training iterations and the same batch size thus consuming 25B tokens in total. We set the learning rate as 4e-4 and use the cosine learning rate scheduler.
the training iterations as 1000 and the batch size as
For linear regression, employ 5-fold cross- valss-vali with ridge regression to determine the optimal & regularization weight from the set [le-3, le-2, le-1, le0, lel, le2, Ie3J. For LightGBM, we manually set the number of iterations to 1000 and the learning rate le-2. leaving all other hyperparameters at their default values.
Table 9: The detailed model configuration for different model sizes.
Model
1M
60M
1B
7B
Vocabulary Size
50432
50432
50432
50432
Rlyers
2
10
22
32
Theads
8
16
16
demhedding
256
768
2048
4096
dmodel
512
1536
5632
12288
F THE STABILITY OF OUR METHOD
Previous research (Xie al., 2023a; Fan et al., 2023; Albalak et al., 2023) has employed small-scale proxy models, trained on substantial volumes of tokens, to predict optimal data mixtures for large language models. However, these approaches often suffer from instability issues. For example, DoReMi (Xie et al., 2023a) reported that different proxy model sizes can result in significantly different predicted data mixtures. Their findings (Figure Appendix) show that using 280M proxy model resulted in a Pile-CC weight of 0.67, while a 1B proxy model yielded a Pile-CC weight below 0.20. The large discrepancy highlights potential instabilities in previous approaches. To evaluate the robustness of REGMIX against such instabilities, we conducted comparative experiments using two distinct model scales: a 1M proxy model and a 60M proxy model. We used their respective training logs fit regression models and subsequently simulated the top 1024 predictions. The resulting distributions are plotted Figure 14. Our results demonstrate that while the prediction distributions for the IM and 60M models not identical, they exhibit remarkably similar patterns.
25

---

Published as a conference paper at ICLR 2025
This consistency suggests that REGMIX achieves improved stability compared previous approaches, even when varying the scale of proxy training models.
G DETAILED EXPERIMENTAL RESULTS
To facilitate future research, we share all the data mixtures and the corresponding downstream performances of the 64 trained models with IB parameters.
Model Index
3
4
5
6
7
rre-trainig Domain Weights
ArXiv
0.123
0.066
0.055
0,059
0.201
0.036
0.042
0.126
FreeLaw
0.065
0.071
0.052
0.083
0.004
0.212
0.113
0.21
NIH ExPorter
0.0
0.0
0.004
0.0
0.014
0.0
0.0
0.0
PubMed Central
0.126
0.211
0.177
0.174
0.243
0.153
0.089
0.123
Wikipedia (en)
0.036
0.013
0.02
0.177
0.01
0.005
0.022
0.055
DM Mathematics
0.0
0.0
0.011
0.0
0.03
0.047
0.007
0.008
Github
0.034
0.153
0.095
0.194
0.017
0.205
0.028
0.008
PhilPaperrs
0.0
0.033
0.0
0.0
0.0
0.0
0.0
0.0
Stack Exchange
0.039
0.097
0.18
0.0
0.103
0.075
0.011
0.129
Enron Emails
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
Gutenberg (PG-19)
0.0
0.0
0.016
0.0
0.002
0.0
0.217
0.035
Pile-CC
0.27
0.101
0.381
0.192
0.359
0.209
0.232
0.288
Ubuntu IRC
0.0
0.0
0.001
0.005
0.0
0.08
EuroParl
0.0
0.0
0.0
0.109
0.0
0.001
0.117
0.0
HackerNews
0.0
0.011
0.005
0.0
0.0
0.0
0.018
0.0
PubMed Abstracts
0.0
0.136
0.0
0.005
0,014
0.002
0.011
0.016
USPTO Backgrounds
0.307
0.106
0.003
0.0
0.002
0.055
0.011
0.0
Downstream
Performance (%)
Social IQA
33.27
33.33
33.62
33.53
33.49
33.56
33.62
33.55
HellaSwag wag
40.58
36.86
40.58
36.06
40.07
37.85
37.93
39.59
PiQA
67.29
65.14
67.97
64.66
67.03
65.36
66.0
66.55
OpenBookQA
28.63
27.87
29.33
29.1
29.23
28.33
29.13
28.73
Lambada
29.17
26.86
31.55
27.11
29.16
28.92
31.53
30.92
SciQ
80.68
79.98
81.05
80.8
82.4
79.88
78.67
79.7
CОРА
70.5
63.83
69.17
65.0
67.5
66.0
66.67
68.67
RACE
29.47
30.0
32.11
28.82
31.13
30.06
29.9
30.75
ARC Easy
50.03
48.72
50.01
46.64
51.06
47.46
46.75
48.39
LogiQA
23.76
24.17
25.29
25.29
24.55
25.96
25.45
26.32
QQP
55.71
55.9
54.84
56.52
54.01
56.34
52.35
54.2
WinoGrande
51.54
51.59
51.39
50.91
53.13
52.26
51.26
51.45
MultiRC
52.65
53.39
51.89
50.92
49.03
53.09
53.64
50.23
Avg
47.18
45.97
47.60
45.80
47.06
46.54
46.38
46.85
H PESUDOCODE OF REGMIX
To improve the clarity, provide additional details regarding the regression model fitting process and the data used. Specifically, the regression model is trained using N data points generated by evaluating proxy models randomly sampled data mixtures. A pseudocode representation of the Algorithm is included below to outline the procedure.
26

---

Published as a conference paper at ICLR 2025
Model Index
10
11
12
13
14
15
16
Pre- Pre-training Domain Weights
ArXiv
0.184
0.226
0.107
0.139
0.101
0.099
0.251
0.147
FreeLaw
0.009
0.046
0.276
0.048
0.047
0.002
0.024
0.046
NIH ExPorter
0.0
0.0
0.0
0.0
0.001
0.022
0.0
0.0
PubMed Central
0.094
0.261
0.157
0.184
0.119
0.501
0.101
0.196
Wikipedia (en)
0.035
0.001
0.009
0.032
0.049
0.003
0.17
0.14
DM Mathematics
0.007
0.001
0.0
0.001
0.092
0.0
0.0
0.008
Github
0.106
0.189
0.024
0.055
0.078
0.017
0.048
0.237
PhilPapers
0.0
0.0
0.0
0.0
0.0
0.043
0.019
0.0
Stack Exchange
0.142
0.077
0.051
0.109
0.002
0.065
0.007
0.06
Enron Emails
0.0
0.0
0.0
0.0
0.0
0.0
Gutenberg (PG-19)
0.0
0.01
0.001
0.0
0.051
0.091
0.0
0.012
Pile-CC
0.341
0.114
0.273
0.354
0.283
0.055
0.339
0.111
Ubuntu IRC
0.0
0.003
0.0
0.0
0.057
0.0
0.017
0.0
EuroParl
0.0
0.0
0.003
0.003
0.0
0.006
0.0
0.0
HackerNews
0.002
0.0
0.034
0.0
0.0
0.0
0.0
0.001
PubMed Abstracts
0.005
0.039
0.009
0.075
0.061
0.007
0.0
0.01
USPTO Backgrounds
0.075
0.033
0.056
0.0
0.057
0.088
0.024
0.032
Downstream
Performance (%)
Social IQA
33.43
33.21
33.31
33.17
33.28
32.43
33.57
33.7
HellaSwag
40.05
35.89
39.55
39.89
38.63
36.18
39.52
35.94
PiQA
66.6
64.74
66.29
66.27
66.9
64.05
66.7
64.51
OpenBookQA
28.87
26.6
29.33
28.73
29.4
27.87
29.67
27.83
Lambada
31.39
27.37
30.32
30.31
31.38
26.25
29.86
26.95
SciQ
81.1
79.12
79.97
82.85
79.42
81.4
81.38
81.23
CОРА
67.0
64.5
66.83
69.5
67.33
65.83
69.5
66.33
RACE
30.57
29.63
30.49
30.85
30.35
28.66
31.21
29.57
ARC Easy
50.66
47.74
47.47
50.18
49.92
49.52
50.73
48.65
LogiQA
23.6
25.65
26.37
23.81
25.58
26.29
25.86
25.12
QQP
54.89
54.79
54.2
55.23
53.69
57.09
53.95
54.24
WinoGrande
50.83
51.84
51.05
51.83
52.12
52.0
51.01
51.82
MultiRC
54.18
54.48
50.17
52.12
51.42
52.69
51.87
53.48
Avg
47.17
45.81
46.57
47.29
46.88
46.17
47.30
46.11
Algorithm 1 REGMIX: Data Mixture Regression
1: Input: Token mixtures for n domains x" = {2 .2....,,.h number of proxy models N, target
metric y, and regression model
2: Output: Optimal data mixture x* = {-i, 2 ...,,
3: Step 1: Train Proxy Models
4: Generate N random data mixtures {x; J.N Jiz where each = {-i ii.... 2f is sampled from
Dirichlet distribution Dir(a A x x"), with X [0.1,5.0] to ensure diversity.
5: for each mixture x; do
6: Train small-scale proxy model using x; for a fixed number of tokens.
7: Evaluate the proxy model compute the target metric yi (e.g., validation loss).
8: end for
9: Step 2: Fit Regression Model
10: Use ((x.)))) to train the regression model fxx to predict
Step 3: Simulate Data Mixtures and Predict Performance
12: Generate a large set of candidate mixtures {x, }MM )-
13: Predict the target metric for each mixture: yj = f(x)).
Identify the mixture x" that minimizes y.:
arg ming, f(xx).
15: Step 4: Train Large-Scale Model
16: Use the identified optimal mixture x* to train a large-scale model with significantly more tokens. Optionally, average top-pertin op-rerforinng mixtures for robustness.
17: Return: Optimal data mixture x
27

---

Published as a conference paper at ICLR 2025
Model Index
20
22
23
24
Pre-training Domain Weights
ArXiv
0.228
0.0
0.501
0.101
0.047
0.031
0.078
0.068
FreeLaw
0.016
0.019
0.005
0.03
0.014
0.073
0.024
0.181
NIH ExPorter
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
PubMed Central
0.204
0.084
0.156
0.272
0.163
0.053
0.302
0.126
Wikipedia (en)
0.02
0.159
0.17
0.021
0.218
0.129
0.027
0.07
DM Mathematics
0.036
0.009
0.0
0.099
0.0
0.0
0.0
0.001
Github
0.02
0.012
0.022
0.124
0.137
0.066
0.04
0.195
PhilPapers
0.004
0.0
0.017
0.0
0.0
0.0
0.0
0.0
Stack Exchange
0.002
0.052
0.062
0.113
0.173
0.12
0.007
0.24
Enron Emails
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
Gutenberg (PG-19)
0.0
0.001
0.002
0.054
0.001
0.089
0.002
0.0
Pile-CC
0.244
0.361
0.061
0.154
0.19
0.057
0.499
0.023
Ubuntu IRC
0.0
0.296
0.002
0.0
0.029
0.001
0.0
0.0
EuroParl
0.004
0.0
0.0
0.001
0,007
0.0
0.0
0.0
HackerNews
0.0
0.0
0.0
0.0
0,011
0,031
0.0
0.0
PubMed Abstracts
0.196
0.001
0.0
0.011
0.008
0.351
0.0
0.059
USPTO Backgrounds
0.026
0.007
0.002
0.02
0.001
0.001
0.021
0.036
Downstream
Performance (%)
Social IQA
33.89
33.31
33.53
33.38
33.75
33.24
33.56
33.71
HcllaSwag
38.68
39.9
34.67
37.12
37.44
36.07
42.15
34.67
PiQA
66.83
67.39
63.33
64.83
65.0
63.68
67.8
62.99
OpenBookQA
28.13
30.67
28.03
29.4
27.67
27.77
29.37
25.83
Lambada
28.78
28.56
24.13
29.41
27.67
28.03
33.47
24.04
SciQ
79.6
78.83
77.42
78.98
78.95
78.72
81.83
79.12
CОРА
65.17
68.17
65.33
67.33
67.67
62.67
69.83
65.83
RACE
28.74
30.03
29.76
29.49
30.77
29.76
31.21
27.91
ARC Easy
48.86
49.42
47.9
48.3
47.88
46.68
50.92
45.24
LogiQA
25.91
26.34
26.24
25.76
26.11
26.24
24.17
25.91
QQP
53.35
53.18
50.61
51.49
54.27
54.99
52.77
55.19
WinoGrande
52.54
51.17
52.01
51.09
52.13
52.03
52.5
50.28
MultiRC
51.49
52.45
55,4
54.87
51.73
49.49
50.61
50.29
Avg
46.30
46.88
45.26
46.27
46.23
45.34
47.71
44.69
28

---

Published as a conference paper at ICLR 2025
Model Index
25
28
29
30
31
32
Pre-training Domain Weights
ArXiv
0.074
0.076
0.05
0.067
0.244
0.073
0.234
0.08
FreeLaw
0.214
0.085
0.039
0.052
0.023
0.087
0.015
0.134
NIH ExPorter
0.0
0.0
0.0
0.0
0.0
0.026
0.0
0.0
PubMed Central
0.135
0.214
0.049
0.221
0.064
0.175
0.086
0.255
Wikipedia (en)
0.011
0.005
0.068
0.052
0.151
0.017
0.287
0.058
DM Mathematics
0.0
0.0
0.019
0.0
0.0
0.101
0.026
0.037
Github
0.121
0.127
0.042
0.101
0.073
0.1
0.04
0.171
PhilPapers
0.006
0.0
0.0
0.0
0.0
0.019
0.0
0.0
Stack Exchange
0.024
0.204
0.146
0.001
0.02
0.054
0.022
0.015
Enron Emails
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
Gutenberg (PG-19)
0.001
0.147
0.01
0.265
0.017
0.0
0.0
0.045
Pile-CC
0.088
0.138
0.302
0.214
0.383
0.12
0.134
0.182
Ubuntu IRC
0.001
0.002
0.0
0.026
0.01
0.134
0.0
0.0
EuroParl
0.0
0.0
0.008
0.0
0.0
0.037
0.0
0.0
HackerNews
0.004
0.0
0.0
0.0
0.0
0,0
0.0
0.0
PubMed Abstracts
0.132
0.001
0.01
0.002
0.007
0.053
0.022
0.016
USPTO Backgrounds
0.189
0.001
0.255
0.0
0.007
0.002
0.134
0.008
Downstream
Performance (%)
Social IQA
33.51
33.4
33.59
33.52
33.53
33.49
33.16
33.56
HcllaSwag
36.75
36.97
40.81
38.25
40.28
35.71
37.37
37.39
PiQA
64.09
64.74
67.97
66.15
66.88
63.84
64.47
65.05
OpenBookQA
29.47
28.7
29.57
29.77
29.5
29.13
29.47
28.0
Lambada
26.69
33.0
31.6
33.08
31.49
27.69
26.99
29.54
SciQ
80.03
79.17
80.12
80.22
81.92
78.23
77.42
80.87
CОРА
67.67
65.5
69.0
65.67
68.33
63.33
64.67
67.17
RACE
30.05
30.19
30.96
30.37
30.08
29.62
30.13
29.92
ARC Easy
47.5
46.9
50.26
48.57
50.55
46.96
48.77
48.79
LogiQA
27.24
25.55
25.86
24.37
25.32
25.12
26.4
24.3
QQP
49.68
55.43
50.94
50.91
51.99
53.53
49.53
51.36
WinoGrande
51.68
52.12
51.93
51.5
52.32
51.67
52.13
52.63
MultiRC
51.24
51.91
50.33
52.42
52.52
54.04
52.05
53.04
Avg
45.82
46.43
47.15
46.52
47.29
45.57
45.58
46.28
29

---

Published as a conference paper at ICLR 2025
Model Index
36
37
38
39
40
Pre-training Domain Weights
ArXiv
0.105
0.295
0.142
0.279
0.052
0.251
0.239
0.157
FreeLaw
0.007
0.029
0.122
0.01
0.07
0.007
0.087
0,062
NIH ExPorter
0.0
0.0
0.001
0.0
0.253
0.007
0.0
0.0
PubMed Central
0.407
0.061
0.065
0.184
0.4
0.331
0.223
0.039
Wikipedia (en)
0.045
0.124
0.0
0.0
0.003
0.107
0.029
0.096
DM Mathematics
0.054
0.0
0.001
0.0
0.0
0.0
0.0
0.007
Github
0.017
0.006
0.006
0.108
0.033
0.13
0.049
0.057
PhilPapers
0.0
0.0
0.003
0.0
0.0
0.0
0.0
0.0
Stack Exchange
0.126
0.006
0.001
0.097
0.019
0.021
0.202
0.174
Enron Emails
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
Gutenberg (PG-19)
0.009
0.047
0.014
0.039
0.0
0.001
0.0
0.015
Pile-CC
0.167
0.364
0.618
0.198
0.031
0.006
0.156
0.181
Ubuntu IRC
0.0
0.0
0.001
0.0
0.0
0.12
0.0
0.0
EuroParl
0.007
0.026
0.0
0.0
0.0
0.0
0.0
0.089
HackerNews
0.0
0.004
0.0
0.0
0.018
0.0
0.0
0.012
PubMed Abstracts
0.047
0.0
0.0
0.083
0.002
0.005
0.012
0.016
USPTO Backgrounds
0.008
0.037
0.025
0.002
0.119
0.014
0.001
0.095
Downstream
Performance (%)
Social IQA
33.48
33.28
33.35
33.29
33.63
33.61
33.21
33.61
HcllaSwag
38.0
40.18
43.37
37.69
32.96
32.98
37.31
37.79
PiQA
65.3
66.68
69.04
66.46
62.25
60.17
65.24
65.32
OpenBookQA
29.43
30.37
30.43
27.63
26.43
26.83
27.97
28.7
Lambada
26.59
31.46
31.71
30.21
18.92
20.29
28.1
28.58
SciQ
79.82
80.58
82.13
80.83
76.73
77.9
79.12
79.6
CОРА
64.33
69.33
67.0
67.83
61.5
62.67
64.67
66.0
RACE
30.03
30.16
32.47
30.49
29.27
28.12
30.11
30.21
ARC Easy
48.86
49.88
52.22
48.32
44.86
45.54
48.15
48.86
LogiQA
25.91
24.3
23.35
24.96
26.19
27.68
25.47
25.37
QQP
56.06
56.56
52.57
56.7
52.54
48.04
49.81
57.12
WinoGrande
50.92
50.97
52.39
52.7
52.3
51.68
51.42
52.8
MultiRC
53.09
49.97
52.18
49.05
53.78
52.27
51.45
55.68
Avg
46.29
47.21
47.86
46.63
43.95
43.67
45.54
46.90
30

---

Published as a conference paper at ICLR 2025
Model Index
41
44
45
46
47
48
Pre-training Domain Weights
ArXiv
0.422
0.466
0.027
0.063
0.121
0.041
0.033
0.114
FreeLaw
0.213
0.075
0.041
0.089
0.008
0.025
0.048
0.116
NIH ExPorter
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
PubMed Central
0.08
0.07
0.116
0.219
0.093
0.111
0.22
0.081
Wikipedia (en)
0.019
0.006
0.021
0.001
0.008
0.092
0.027
0.038
DM Mathematics
0.001
0.0
0.001
0.05
0.016
0.062
0.002
0.031
Github
0.026
0.044
0.067
0.291
0.012
0.121
0.169
0.109
PhilPapers
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
Stack Exchange
0.003
0.078
0.137
0.002
0.408
0.124
0.082
0.001
Enron Emails
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
Gutenberg (PG-19)
0.01
0.0
0.001
0.0
0.006
0.0
0.057
0.021
Pile-CC
0.026
0.2
0.549
0.238
0.156
0.214
0.312
0.428
Ubuntu IRC
0.0
0.0
0.002
0.0
0.013
0.129
0.0
0.001
EuroParl
0.0
0.0
0.0
0.001
0,001
0.006
0.0
0.0
HackerNews
0.0
0.0
0.0
0,001
0.0
0,012
0.0
0.0
PubMed Abstracts
0.101
0.028
0.002
0.045
0.005
0.012
0.0
0.031
USPTO Backgrounds
0.099
0.031
0.037
0.0
0.153
0.052
0.05
0.029
Downstream
Performance (%)
Social IQA
33.49
33.43
33.07
33.28
33.44
33.08
33.78
33.17
HcllaSwag
34.51
37.59
42.69
37.37
38.31
38.3
39.67
41.07
PiQA
62.24
65.58
68.05
66.62
66.54
65.52
66.98
67.21
OpenBookQA
27.1
28.77
28.9
28.07
28.07
27.6
31.17
29.73
Lambada
22.78
26.99
31.34
29.51
27.87
29.47
30.34
32.71
SciQ
77.78
80.25
79.47
80.25
80.7
79.72
81.35
81.77
CОРА
64.0
66.33
67.0
67.0
67.33
68.33
67.17
67.67
RACE
28.33
28.82
30.78
30.8
30.08
30.24
30.24
30.67
ARC Easy
45.48
48.64
51.49
46.99
48.79
48.05
49.58
49.49
LogiQA
24.83
24.96
24.76
23.25
26.06
25.55
24.32
24.68
QQP
50.27
54.73
53.96
57.0
53.73
51.19
57.52
56.91
WinoGrande
51.79
51.63
51.32
50.76
53.18
52.45
50.72
52.24
MultiRC
54.03
53.96
48.91
50.74
53.01
50.89
47.63
53.84
Avg
44.35
46.28
47.06
46.28
46.7
46.18
46.96
47.78
31

---

Published as a conference paper at ICLR 2025
Model Index
52
53
54
55
56
Pre-training Domain Weights
ArXiv
0.082
0.091
0.194
0.011
0.039
0.294
0.012
0.25
FreeLaw
0.12
0.084
0.04
0.022
0.063
0.119
0.16
0.058
NIH ExPorter
0.0
0.0
0.022
0.0
0.0
0.0
0.0
0.0
PubMed Central
0.051
0.343
0.126
0.37
0.079
0.186
0.311
0.104
Wikipedia (en)
0.067
0.0
0.046
0.006
0.0
0.023
0.014
0.044
DM Mathematics
0.034
0.174
0.028
0.0
0.002
0.005
0.0
0.0
Github
0.205
0.144
0.048
0.14
0.482
0.023
0.117
0.028
PhilPapers
0.0
0.0
0.01
0.0
0.0
0.0
0.0
0.0
Stack Exchange
0.036
0.009
0.099
0.058
0.012
0.001
0.004
0.06
Enron Emails
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
Gutenberg (PG-19)
0.0
0.019
0.04
0.216
0.0
0.002
0.236
0.0
Pile-CC
0.371
0.122
0.229
0.101
0.269
0.213
0.037
0.363
Ubuntu IRC
0.0
0.001
0.0
0.033
0.0
0.023
0.007
0.0
EuroParl
0.0
0.003
0.002
0.0
0.0
0.0
0.0
0.0
HackerNews
0.0
0.001
0.0
0,002
0.0
0,0
0.0
0.0
PubMed Abstracts
0.029
0.006
0.089
0.026
0.002
0.024
0.007
0.086
USPTO Backgrounds
0.004
0.004
0.027
0.015
0.052
0.088
0.094
0.007
Downstream
Performance (%)
Social IQA
33.53
33.74
33.37
33.41
32.96
33.88
33.75
33.79
HcllaSwag
39.09
35.65
38.68
36.07
37.68
38.53
35.4
40.5
PiQA
66.81
64.58
65.68
63.99
65.85
65.76
64.51
66.89
OpenBookQA
29.13
27.57
28.27
29.1
29.43
28.73
28.3
29.87
Lambada
30.23
26.19
30.29
30.84
29.76
29.03
28.63
30.74
SciQ
79.9
80.83
78.4
80.03
81.38
80.92
77.75
82.07
CОРА
68.17
61.83
67.0
66.0
66.17
63.17
66.33
64.0
RACE
31.42
29.35
30.41
31.08
30.77
29.73
30.8
31.42
ARC Easy
49.54
47.71
49.02
47.64
48.38
49.36
46.96
51.22
LogiQA
24.99
24.58
25.32
24.91
25.17
26.22
24.63
24.91
QQP
54.06
56.48
50.96
56.62
56.45
53.86
53.85
53.26
WinoGrande
50.51
50.26
51.83
51.33
52.18
51.89
51.59
50.5
MultiRC
50.25
54.37
50.94
52.38
51.21
55.34
54.52
50.5
Avg
46.74
45.63
46.17
46.42
46.72
46.65
45,92
46.90
32

---

Published as a conference paper at ICLR 2025
Model Index
60
62
63
64
Pre-training Domain Weights
ArXiv
0.137
0.176
0.471
0.081
0.107
0.278
0.119
0.131
FreeLaw
0.085
0.007
0.038
0.153
0.016
0.141
0.085
0.006
NIH ExPorter
0.0
0.0
0.0
0.0
0.0
0.0
0.027
0.03
PubMed Central
0.085
0.05
0.218
0.17
0.218
0.257
0.294
0.075
Wikipedia (en)
0.059
0.122
0.005
0.017
0.003
0.099
0.02
0.0
DM Mathematics
0.0
0.001
0.0
0.033
0.0
0.009
0.073
0.093
Github
0.039
0.088
0.097
0.041
0.238
0.041
0.038
0.369
PhilPapers
0.0
0.069
0.0
0.048
0.0
0.0
0.0
0.0
Stack Exchange
0.017
0.05
0.016
0.077
0.113
0.027
0.046
0.06
Enron Emails
0.009
0.0
0.0
0.0
0.0
0.0
0.001
0.0
Gutenberg (PG-19)
0.007
0.0
0.018
0.001
0.0
0.0
0.026
0.002
Pile-CC
0.435
0.339
0.112
0.268
0.272
0.128
0.232
0.188
Ubuntu IRC
0.0
0.006
0.017
0.095
0.001
0.0
0.0
0.001
EuroParl
0.0
0.012
0.0
0.0
0.0
0.0
0.001
0.003
HackerNews
0.001
0.0
0.0
0.0
0.0
0.0
0.0
0.017
PubMed Abstracts
0.004
0.004
0.001
0.0
0.02
0.0
0.013
0.016
USPTO Backgrounds
0.122
0.077
0.006
0.016
0.013
0.02
0.025
0.009
Downstream
Performance (%)
Social IQA
33.24
33.3
33.56
33.54
33.42
33.84
33.32
33.55
HcllaSwag
41.74
39.63
35.36
38.83
38.53
36.46
38.8
36.43
PiQA
68.07
67.31
64.44
66.38
66.5
64.74
66.54
64.87
OpenBookQA
29.2
29.5
28.1
27.97
27.83
27.37
28.83
27.87
Lambada
31.79
31.11
27.32
30.17
28.75
26.22
30.38
26,25
SciQ
80.42
79.83
80.85
79.6
78.93
80.05
79.5
78.65
CОРА
66.17
69.0
64.0
64.83
67.0
64.0
66.0
66.83
RACE
31.39
29.82
29.67
30.08
29.98
29.46
30.37
29.19
ARC Easy
51.14
49.24
47.13
47.88
48.2
47.09
49.09
46,9
LogiQA
25.19
25.93
23.68
25.17
25.7
25.52
26.5
26.65
QQP
55.37
54.46
52.73
53.17
59.65
58.15
57.5
55.31
WinoGrande
53.21
51.46
50.83
52.16
52.37
51.41
51.63
51.85
MultiRC
53.58
52.31
52.22
53.03
50.41
52.17
52.27
51.5
Avg
47.73
47.15
45,38
46.37
46.71
45.88
46.98
45.84
33

---

Published as a conference paper at ICLR 2025
I RANK INVARIANCE HYPOTHESIS
The rank invariance hypothesis asserts that the relative rankings data mixtures should remain stable varying model sizes (e.g., from 1M 1B parameters) and token scales (e.g., from IB to 25B tokens). This hypothesis suggests that the comparative effectiveness of different data mixtures does not significantly change models grow in size
trained larger amounts of tokens.
To examine this hypothesis, we conducted a comprehensive set of experiments involving models of four distinet scales: 1M, 60M, 280M, and IB parameters. Each model was trained using 64 different data mixture configurations, where the token amounts varied systematically. For each setting, we measured the validation loss for all data mixtures and determined their rankings. To quantify the consistency of rankings across model sizes and token scales, we computed the Spearman rank correlation coefficient (p) for each comparison.
The results shown in Figure 15, reveal high rank correlation coefficients across all model and token scales. This finding provides strong empirical support for the rank invariance hypothesis, indicating that the relative utility of data mixtures is robust to changes in model size and training tokens.
1.00
1M (1B)
0.96
0.99
0.97
0.98
(2533) 280M (533 60M (1B)
0.96
0.96
0.94
0,96
0.99
0.96
0.98
0.94
0.92
0.97
0.94
0.98
0,90
(1B) (1B) (51) (25B)
Figure 15: Heatmap Spearman rank correlation coefficients (p) between validation loss rankings for different model scales (IM, 60M, 280M, and IB parameters) and token scales (IB, 10B, and 25B tokens). The consistently high correlation values across all settings support the rank invariance hypothesis, indicating that the relative ranking of data mixtures remains stable despite changes in model size and training scale.
Table 10: Performance comparison between Human and RegMix on 7B models trained on 100B tokens across 13 downstream benchmarks.
Beuchmark
Human
REGMMX
Social IQA
41.6
43.3
HellaSwag
55.0
63.3
PiQA
72.4
75.3
OpenBookQA QA
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
MultRCC
52.6
51.0
Average
54.5
56.5 (+2.0)
34

---

Published as a conference paper at ICLR 2025
J SCALING REGMIX TO 7B MODELS OVER 100B TOKENS
To further validate the effectiveness of our method on larger models, we conducted an experiment using Human baseline and our method data mixtures on 7B model trained on 100B tokens. The results, summarized in Table I, demonstrate that RegMix still outperforms the Human baseline, achieving an average performance boost of 2%.
To provide closer illustration, we benchmark the downstream performance of RegMix and Human on every dataset intervals of 25B tokens in Figure 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 RegMix and Human struggle to improve on certain benchmarks (e.g., MultiRC), even with increased token amounts. These findings suggest that RegMix can almost benefit 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 quite intriguing and warrants further investigation.
K USING LARGER PROXY MODEL
We conducted 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 main experiments). Both configurations used IB training tokens per proxy model. We limited our investigation to these configurations due to computational constraints that prevented us from exploring scenarios with more IB-parameter proxy models.
Table II: Performance comparison of optimized data mixtures derived from two proxy configurations (512x IM and 128x 1B) for 7B models trained on 100B tokens, evaluated across 13 downstream benchmarks.
Benchmark
REGMIX (IB proxy)
REGMIX (IM proxy)
Social IQA
43.4
43,3
HellaSwag
62.9
63.3
PiQA
75.1
75.3
OpenBook
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
To evaluate these configurations, we used their respective optimized data mixtures to train 7B models on 100B tokens and compared their performance. The results, summarized in Table K, 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 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 starting point to optimize data mixtures for language model pre-training.
35


Reviews:
[{"review_id": "zpE2b0tFPF", "reviewer": "Reviewer_YiAH", "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:\n- 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.\n- Data mixture has a substantial impact on downstream performance, with single-task performance differences reaching up to 14.6%.\n- General web corpora (such as CommonCrawl) outperform traditionally high-quality data like Wikipedia in driving downstream performance.\n- Domain interactions are complex and often counterintuitive, highlighting the value of automated approaches like REGMIX.\n- Data mixture effects go beyond scaling laws, with REGMIX capturing the complexity by jointly optimizing across all domains.", "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? \n     - Also, could you discuss the rationale behind this range and whether you conducted any sensitivity analyses to determine its impact on the results?\n\n2) 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.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "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.\n\n- 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.\n\n- 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.\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jbludrE3hB", "reviewer": "Reviewer_MoX3", "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.\n\nThe 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.\n\nThe 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.", "questions": "How do we decide the size for proxy model and training tokens.?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "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.\n\n2. The paper’s experimental setup is robust, with 512 small proxy models across diverse data mixtures, creating a solid regression model for data selection.\n\n3. 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.\n\nThe paper only experiments with 1M models with 1B tokens. \nIt is unclear how to decide the size of the proxy model parameter and training token.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0Tp44MacnB", "reviewer": "Reviewer_FGaq", "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.", "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?\n- 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?\n- 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.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "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.\n- 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.\n- 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.\n- 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.\n- 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. \n- 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?\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HZqm7E8ggY", "reviewer": "Reviewer_wxJ7", "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.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-written and easy to follow.\n2. The author identifies a helpful assumption, namely ranking invariance regarding training scales, which helps reduce the cost of tuning data mixtures in this paper.\n3. The proposed method is simple to implement and efficient, thus appealing to try in practice.\n4. 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.\n2. 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.\n3. 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.\n\n[1] Data Mixing Made Efficient: A Bivariate Scaling Law for Language Model Pretraining\n\n[2] Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance\n\n[3] Scaling Laws for Data Filtering— Data Curation cannot be Compute Agnostic", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XOVYMqyyjA", "reviewer": "Reviewer_uBLs", "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.", "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?\n2. 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?\n3. there are some related work on data selection that falls into the group-level data selection category: https://aclanthology.org/2020.acl-main.754/", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. the proposed method is novel and has practical impact to LLM training\n2. the authors evaluated the predicted mixture on 1B model and compared to a few prior methods based on the downstream performance of the model\n3. 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\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)