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

1 INTRODUCTION

• Manual selection can be suboptimal (Albalak et al., 2024), while existing proxy approaches become inefficient as pre-training data grows.Prior approaches (Xie et al., 2023a; Fan et al., 2023; Albalak et al., 2023) dynamically allocate data using proxies trained on substantial token budgets, such as 100B tokens.Source: §1 · p. 2
• RegMix fits a regression model to small-model performance across diverse mixtures, then predicts performance for untested mixtures.Its key assumption is rank invariance: mixtures retain their relative performance rankings across model sizes and training-token counts, enabling selection without exhaustively training every mixture.Source: §1 · p. 2
• RegMix outperforms human mixture selection and matches DoReMi (Xie et al., 2023a), requiring less total compute and allowing parallel training.Experiments with 512 models containing 1M parameters trained on 1B tokens predicted the optimal mixture among 64 candidate 1B-parameter models trained on 25B tokens.Source: §1 · p. 2
• 14.6% differences in single-task performance demonstrate the substantial downstream impact of data mixture selection.The reported difference is an upper bound observed across mixtures, not an overall benchmark improvement.Source: §1 · p. 2
• General web corpora, rather than Wikipedia, show the strongest positive correlation with improved performance across downstream tasks.CommonCrawl is one example of the general web corpora identified.Source: §1 · p. 2
• Domain interactions often contradict intuition, and data mixture effects transcend scaling laws.RegMix considers all domains together to capture this complexity; its regression model also provides insights into domain interactions that can inform data curation.Source: §1 · p. 2

2 RELATED WORK

• RegMix replaces DoReMi’s (Xie et al., 2023a) long single-model training with short, parallel small-model runs, improving scalability and yielding better-performing final models.RegMix learns fixed mixture weights before training the final model.Source: §2 · p. 3
• Data selection operates at token, sample, or group level; learned group mixtures use fixed offline weights or dynamically updated online weights.Token-level selection filters tokens, sample-level selection chooses individual examples, and group-level selection mixes data pools.Source: §2 · pp. 2–3
• Pre-training sample selection mostly relies on heuristics, while learned approaches use optimization algorithms, model perplexity, or LLM guidance.Sample-level selection is also commonly used to choose fine-tuning data.Source: §2 · pp. 2–3
• Unlike analytical scaling approaches (Hoffmann et al., 2022), RegMix directly optimizes the target metric through regression for from-scratch pre-training.The authors also find strong correlations between loss and downstream performance, especially for loss measured on web corpora.Source: §2 · p. 3
• Data-scaling research examines quantity, quality, and mixing proportions, including data constraints, dataset-specific scaling, continued pre-training, and downstream-performance prediction.Ye et al. (2024) and Ge et al. (2024) propose functional relationships predicting mixture effects on language-modeling loss; most surveyed work focuses on validation loss.Source: §2 · p. 3

3.1 TRAIN SMALL-SCALE PROXY MODELS

• Proxy-training mixtures are sampled from token-availability-based Dirichlet distributions to cover diverse domain weights without disproportionately emphasizing scarce domains.Sampling aims to cover domain weights from 0% to 100%. Multiplying the token distribution by values from 0.1 to 5.0 produces Dirichlet parameters spanning sparse and near-uniform distributions.Source: §3.1 · p. 4
• 1M-parameter proxy models are trained on 1B tokens each in the main experiment, then evaluated to provide regression targets.Targets can come from domain or benchmark evaluations, such as loss on a chosen domain; mixture weights become features and measured targets become labels.Source: §3.1 · p. 4

3.2 FIT A REGRESSION MODEL

• RegMix predicts continuous target values from domain weights using ridge regression or LightGBM, evaluating unseen mixtures without further model training.Ridge regression applies L2 regularization to prevent overfitting; LightGBM learns an ensemble of decision trees through gradient-based optimization.Source: §3.2 · pp. 4–5

3.3 SIMULATE AND PREDICT

• 1,000,000 candidate mixtures can receive regression predictions in less than 10 CPU seconds, enabling inexpensive search for the best predicted target value.The search evaluates potential mixtures using the fitted regression model rather than training additional language models.Source: §3.2 · p. 4, §3.3 · p. 5

3.4 LARGE-SCALE MODEL TRAINING

• 100 top-ranked mixtures are averaged for large-scale training with more tokens to improve the robustness of regression prediction.The selected mixture is used directly to train the larger model.Source: §3 · p. 3, §3.4 · p. 5

4.1 EXPERIMENTAL SETUP

• 17 copyright-compatible subsets of the Pile dataset (Gao et al., 2021), available on HuggingFace, underpin linear and LightGBM regressions targeting Pile-CC validation loss.The evaluation first tests small-model loss prediction, then tests ranking predictions across model sizes and training-token counts.Source: §4.1 · pp. 5–6, §4 · p. 5

4.2 EXPERIMENTAL RESULTS

• 97.12% Spearman rank correlation on unseen 1B-parameter models trained for 25B tokens supports rank invariance, while LightGBM outperforms linear regression across the reported metrics.The regression transfers predictions learned from 1M-parameter models trained for 1B tokens to substantially larger, longer-trained models.Source: Table 2 · p. 6, §4.2 · p. 6

4.1 EXPERIMENTAL SETUP

• 512 1M-parameter models trained for 1B tokens provide regression training artifacts, with evaluation spanning unseen mixtures at 1M, 60M, and 1B parameters.Evaluation uses 256 unseen mixtures each for 1M- and 60M-parameter models trained for 1B tokens, and 64 for 1B-parameter models trained for 25B tokens. Spearman correlation measures ranking agreement and MSE measures loss-prediction error; larger-model MSE values are excluded from analysis because of scale differences.Source: §4.1 · p. 6, Table 2 · p. 6

4.2 EXPERIMENTAL RESULTS

• 512 proxy models trained on 0.2B tokens outperform 128 trained on 0.8B tokens, favoring more models over longer training at the same FLOPs budget.Benefits from increasing proxy training tokens saturate after approximately 0.25B tokens. Increasing proxy-model count consistently improves performance, particularly for LightGBM.Source: §4.2 · p. 6

5.1 DATA MIXTURE SIGNIFICANTLY IMPACTS DOWNSTREAM PERFORMANCE

• 14.6 is the largest downstream performance gap, observed on Lambada among 64 models with 1B parameters trained using different data mixtures.Each model trains on 25B tokens, with task scores averaged across 0-shot to 5-shot evaluations. High-variance benchmarks are excluded; evaluation uses normalized accuracy when available, otherwise regular accuracy.Source: §5 · p. 6, Table 3 · p. 7, §5.1 · p. 7

5.2 WEB CORPORA BENEFITS DOWNSTREAM PERFORMANCE THE MOST

• Nearly 1.0 correlation links HellaSwag performance to negative Pile-CC validation loss, which correlates most strongly with most downstream tasks.This contradicts the expectation that Wikipedia validation loss would be most informative. More than 85% of C4100Domain domains (Magnusson et al., 2023) correlate very strongly with Pile-CC, suggesting diverse coverage may explain its downstream relevance.Source: §5.2 · pp. 7–8, Figure 5 · p. 8

5.3 DATA MIXTURE BY REGMIX IMPROVES DOWNSTREAM PERFORMANCE

• 6.8 points over Human selection on HellaSwag accompany REGMIX’s highest average score among evaluated methods.REGMIX beats the other three methods in 7 of 14 cases and surpasses the best randomly sampled mixture; Pile-CC Only also performs strongly. ODM and DoReMi weights are renormalized across 17 available domains, potentially disadvantaging these baselines.Source: §5.3 · pp. 8–9, Table 4 · p. 9
• REGMIX still outperforms baselines on Pile-CC validation loss when Pile-CC is entirely excluded from pre-training.LightGBM selects mixtures to minimize Pile-CC validation loss rather than losses across all domains. This proxy may not generalize to every downstream task.Source: §5.3 · pp. 8–9, Figure 6 · p. 9

5.4 DOMAIN INTERACTIONS ARE CHALLENGING FOR HUMANS TO UNDERSTAND

• Linear-regression analyses reveal complex domain interactions in Pile and The Stack (Kocetkov et al., 2022), with PhilPapers appearing beneficial to every other Pile domain.These non-obvious relationships illustrate why selecting mixtures through human intuition is difficult.Source: §5.4 · pp. 9–10

5.5 DATA MIXTURE EFFECTS TRANSCEND SCALING LAWS

• Most domains lack a simple log-log relationship between mixture weight and validation loss, motivating regression over the entire mixture.DM Mathematics exhibits a near log-log linear relationship, whereas domains such as Pile-CC show more complex patterns. REGMIX accounts for all domain weights simultaneously rather than predicting loss from one domain’s weight alone.Source: §5.5 · p. 10

5.6 EXTEND REGMIX TO 100 DOMAINS

• 100-domain preliminary experiments support REGMIX’s scalability, fitting regression models from 1,000 training runs of 1M-parameter models.Domains are defined by base URLs in FineWeb (Penedo et al., 2024a). Linear and LightGBM regressions are evaluated on 64 unseen mixtures at 1M and 60M parameters using rank correlation (ρ) and mean squared error (MSE).Source: §5.6 · p. 10, Table 5 · p. 11

• 64 1B-parameter models were compared, with REGMIX successfully predicting the best data mixture among them.Source: §6 · p. 11
• REGMIX trains small models to predict how different data mixtures affect performance, enabling efficient automatic mixture selection through regression.Source: §6 · p. 11
• The large-scale study provides insights into data-mixture effects, the relationship between loss and downstream performance, and human experts’ difficulties in selecting optimal mixtures.Source: §6 · p. 11

ETHICS STATEMENT

• Optimized mixtures may underrepresent or misrepresent certain domains, causing poor performance or biased outputs for those domains despite improving overall performance.The ethical concern arises because optimization can favor particular domains.Source: ETHICS STATEMENT · p. 12
• Searching for optimal mixtures still consumes computational resources and energy, creating environmental impacts despite the method’s efficiency.The authors identify further reductions in computation cost as worthwhile future work.Source: ETHICS STATEMENT · p. 12

A BROADER IMPACT

• Less than 2% of final training costs is the computational overhead of RegMix’s mixture optimization using 1M-parameter proxy models.The authors identify this use of ultra-small proxy models as previously unexplored and emphasize its accessibility within academic budgets.Source: A BROADER IMPACT · p. 18
• Publicly available datasets and trained models complement computational efficiency in supporting RegMix’s goal of democratizing language-model pre-training research.The authors frame this availability as a commitment to open science.Source: A BROADER IMPACT · p. 18
• Automatic data-mixture optimization outperforms human intuition-based approaches, with PhilPapers illustrating complex domain interactions that transcend human intuition.PhilPapers serves as an in-depth case study of these interactions.Source: A BROADER IMPACT · p. 18
• 64 distinct mixtures, studied using 1B-parameter models evaluated across 12 benchmarks, demonstrate significant performance variation.The authors describe this as the first comprehensive demonstration of such variation across data mixtures.Source: A BROADER IMPACT · p. 18

B LIMITATIONS

• 1B parameters is the largest scale verified for optimal-mixture prediction, while rank invariance is investigated from 1M to 1B parameters.Statistically meaningful correlations at 3B would require 64 models trained on 50B tokens each, equivalent to one 3B model trained on 3.2T tokens, exceeding available compute.Source: B LIMITATIONS · p. 18
• Nearly random MMLU performance (Hendrycks et al., 2021) and negligible GSM8K performance (Cobbe et al., 2021) prevent assessing their correlations with validation loss.The authors attribute these results to scarce relevant data in the Pile corpus and the relatively small 1B-parameter model.Source: B LIMITATIONS · p. 18
• Limited domain data remain difficult to incorporate systematically, despite the no-Pile-CC experiments addressing this issue.Most mixing methods assume unlimited data per domain. Combining RegMix with the data-reuse decay coefficient proposed in Muennighoff et al. (2023) is suggested as future work.Source: B LIMITATIONS · p. 18
• RegMix assumes each example’s domain is known, making application difficult when domain boundaries are unclear and assigning examples to domains is challenging.This assumption is shared with existing data-mixture methods.Source: B LIMITATIONS · p. 18
• Proxy-derived domain weights assume identical tokenizers and vocabulary sizes in proxy and large models; transferring weights across tokenizers remains challenging.The authors identify this assumption across existing data-mixture methods.Source: B LIMITATIONS · p. 18

C.1 THE REGRESSION PREDICTION VISUALIZATION

• Near 100% Spearman Rank Correlation ρ is achieved by LightGBM on 1M-parameter models, outperforming the linear model.The visualization compares predicted and true losses for both regression models.Source: Figure 9 · p. 19, §C.1 · p. 19

C.2 LOSS AND RANK PREDICTION ON SMALL MODELS FOR OUT-OF-DISTRIBUTION SETTING

• Nearly 100% Spearman Rank Correlation ρ is achieved by LightGBM, outperforming the linear model when Pile-CC is excluded from pre-training.The regression model uses artifacts from 512 models with 1M parameters trained for 1B tokens each, and is evaluated on unseen mixtures. The objective is to minimize Pile-CC validation loss using the remaining domains.Source: Table 6 · p. 19, §C.2 · p. 19

C.3 THE DERIVED DATA MIXTURES

• 0.87 is the weight REGMIX assigns to Pile-CC, which the authors describe as aligning with human intuition.Table 7 presents the derived domain weights for different methods.Source: §C.3 · p. 19, Table 7 · p. 20
• 280M-parameter DoReMi weights come from its best-performing configuration; Human, DoReMi, and Online weights are renormalized to available domains from Xie et al. (2023a).Online uses the method’s final domain weights.Source: Table 7 · p. 20

C.4 THE EVALUATION RESULTS USING LIGHTEVAL

• 6 benchmarks show REGMIX outperforming the Human baseline, while its average performance exceeds DoReMi and Pile-CC Only.Human uses The Pile weights (Gao et al., 2021), while DoReMi uses weights from Xie et al. (2023a).Source: Table 8 · p. 20, §C.4 · p. 21
• Five runs underpin each task’s reported average zero-shot performance and standard deviation in the LightEval comparison.The comparison also estimates mixture-selection compute in FLOPs and uses Cohen’s d to identify scores significantly outperforming Human.Source: Table 8 · p. 20
• LightEval follows FineWeb (Penedo et al., 2024b), selecting benchmarks with low sample variance, monotonic training improvements, and above-random scores for 1B-parameter models.These criteria prioritize benchmark stability and suitability for the evaluated model scale.Source: §C.4 · p. 20

D URL DOMAIN CORRELATION GRAPH

• Figures 10–13 visualize correlations between URL domains within C4 subsets and downstream performance.Figure 14 additionally shows similar REGMIX mixture distributions using 1M- and 60M-parameter proxy models, supporting the method’s stability.Source: Figure 10 · p. 22, Figure 11 · p. 23, Figure 12 · p. 24

E IMPLEMENTATION DETAILS

• Hyperparameters are set to their default values.The supplied fragment does not identify which hyperparameters.Source: E IMPLEMENTATION DETAILS · p. 26
• Table 9 provides detailed model configurations for different model sizes.Only the caption is supplied, so specific configuration values cannot be reported.Source: Table 9 · p. 26

F THE STABILITY OF OUR METHOD

• 1M and 60M proxy models produced remarkably similar, but not identical, prediction-distribution patterns for RegMix.This consistency suggests improved stability compared with previous approaches when varying proxy-model scale.Source: F THE STABILITY OF OUR METHOD · pp. 26–27
• 0.67 and below 0.20 were DoReMi’s (Xie et al., 2023a) Pile-CC weights using 280M and 1B proxy models, respectively, illustrating sensitivity to proxy scale.Source: F THE STABILITY OF OUR METHOD · p. 26
• 1024 top predictions were simulated after fitting regression models to training logs from 1M and 60M proxy models.Source: F THE STABILITY OF OUR METHOD · p. 26

G DETAILED EXPERIMENTAL RESULTS

• 64 trained 1B-parameter models have their data mixtures and corresponding downstream performances shared to facilitate future research.Source: G DETAILED EXPERIMENTAL RESULTS · p. 27

H PESUDOCODE OF REGMIX

• RegMix fits regression to proxy-model performance on sampled mixtures, predicts candidate performance, and trains a larger model using the mixture minimizing the target metric.Proxy mixtures are sampled from Dir(α = λ · x0), with λ ∈ [0.1, 5.0], and each proxy trains for a fixed token budget. Top-performing mixtures may optionally be averaged for robustness.Source: H PESUDOCODE OF REGMIX · pp. 27–28

I RANK INVARIANCE HYPOTHESIS

• High Spearman correlations support rank invariance: data-mixture validation-loss rankings remain stable across model sizes and training-token scales.Experiments evaluate 64 mixtures at each setting across 1M, 60M, 280M, and 1B parameters. Figure 15 compares rankings across 1B, 10B, and 25B training tokens.Source: Figure 15 · p. 35, I RANK INVARIANCE HYPOTHESIS · p. 35

J SCALING REGMIX TO 7B MODELS OVER 100B TOKENS

• 2% average performance boost over the Human baseline is achieved by RegMix on a 7B model trained on 100B tokens.The comparison covers 13 downstream benchmarks.Source: Table 10 · p. 35, J SCALING REGMIX TO 7B MODELS OVER… · p. 36
• 50% pre-training acceleration is reported on most benchmarks, including HellaSwag, reaching 75% on some benchmarks, including PiQA.The performance advantage does not diminish as training tokens increase. Both methods struggle to improve MultiRC, suggesting mixture optimization may not help tasks whose performance does not improve with additional training data.Source: J SCALING REGMIX TO 7B MODELS OVER… · p. 36

K USING LARGER PROXY MODEL

• 512 1M-parameter proxies and 128 1B-parameter proxies produce mixtures with similar average downstream performance for 7B models trained on 100B tokens.Each proxy trains on 1B tokens; computational constraints prevented testing more 1B-parameter proxies. Because larger proxies add substantial overhead without significantly outperforming smaller ones, the authors recommend starting with ultra-small proxies and prioritizing their number.Source: Table 11 · p. 36, K USING LARGER PROXY MODEL · p. 36