PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Qian Liu, Xiaosen Zheng, Niklas Muennighoff, Guangtao Zeng, Longxu Dou, Tianyu Pang, Jing Jiang, Min Lin

arXiv:2407.01492v2cs.CLcs.AI

TL;DR

Large-language-model pre-training needs effective data mixtures, but manual selection does not scale as data grows. RegMix treats mixture selection as regression over small proxy-model experiments, then transfers the best predicted mixture to large-scale training. It predicts strong mixtures across scales, outperforming human selection and matching or exceeding DoReMi with less computation.

  • Problem

    As pre-training data grows in scale and diversity, determining an effective mixture becomes increasingly difficult, while manual selection may be suboptimal.

  • Method

    RegMix trains small models on diverse mixtures, fits a regression from domain weights to target performance, and uses it to select mixtures for large-scale training.

  • Results

    RegMix predicted the best mixture among 64 1B-parameter models trained on 25B tokens, achieved 97.12% rank correlation across scales, and matched or exceeded DoReMi with less computation.

  • Takeaways & Limitations

    Data mixtures substantially affect downstream performance, with complex domain interactions and stronger positive correlations for general web corpora than Wikipedia.

  • Takeaways & Limitations

    Rank-invariance validation currently covers models from 1M to 1B parameters; testing 3B models would exceed the available computational resources.

Abstract

from arXiv · show

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. 1000x larger and 25x 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.

1 INTRODUCTION

RegMix addresses the growing difficulty and limited scalability of manually selecting effective pre-training data mixtures by treating mixture selection as regression. It uses small proxy models to predict high-performing mixtures for larger models while exposing domain interactions and strong performance effects.

  • Manual data selection becomes increasingly difficult and may be suboptimal as the scale and diversity of pre-training data grow.
  • RegMix assumes data-mixture rankings remain consistent across model sizes and training-token counts.
  • RegMix trains diverse small proxy models, fits a regression model, and predicts promising unseen mixtures instead of exhaustively testing every mixture.
  • 512 1M-parameter models trained on 1B tokens predicted the best mixture among 64 1B-parameter models trained on 25B tokens.
  • Data mixtures changed single-task performance by up to 14.6%, while web corpora correlated more positively with downstream performance than Wikipedia.

2 RELATED WORK

Prior work studies data selection at token, sample, and group levels, using manual heuristics, learned methods, or proxy models. RegMix is positioned as an offline group-level approach within this broader landscape of mixture and scaling-law research.

  • Data selection methods operate at token, sample, or group level, with group-level methods mixing pools of data domains.
  • Pre-training data selection commonly uses heuristics, while learned approaches use optimization, perplexity, or language models to inform selection.
  • Group-level methods include offline fixed-weight selection and online dynamic reweighting during final-model training.
  • RegMix is an offline group-level selection method, differing from DoReMi’s category of learned mixture approaches.
  • Data-scaling-law research examines interactions among data quantity, quality, and mixing proportions as language models scale.

3 REGMIX: DATA MIXTURE AS REGRESSION

RegMix samples diverse domain mixtures, trains small proxy models, fits a regression model, simulates candidate mixtures, and transfers the best predicted mixture to large-scale training. Its regression inputs are domain weights and its target is an optimization value such as validation loss.

  • RegMix generates diverse mixtures, trains proxy models, fits regression, searches the simulated mixture space, and trains a large model with the selected mixture.
  • Dirichlet sampling based on token distributions exposes regression models to mixtures with extreme domain weights from 0% to 100%.
  • Proxy-model outputs provide target values such as StackExchange loss, while domain-mixture weights serve as regression features.
  • The regression model predicts target values for arbitrary mixtures without requiring additional model training.
  • 3.2 FIT A REGRESSION MODEL: Ridge regression uses mixture weights to predict targets while penalizing coefficient magnitude to reduce overfitting.
  • 3.2 FIT A REGRESSION MODEL: LightGBM provides an alternative ensemble-of-trees regressor designed for efficient and scalable prediction.
  • Predicting 1,000,000 candidate mixtures takes less than 10 CPU seconds, after which the top 100 mixtures are averaged for large-scale training.

4 EVALUATING ON REGRESSION PREDICTION

The evaluation tests whether regression predictions from small proxy models transfer to unseen mixtures across model scales and token budgets. Results support the rank-invariance hypothesis and show that increasing proxy-model count can be more effective than extending each proxy model’s training.

  • 4.1 EXPERIMENTAL SETUP: The study evaluates predictions on unseen mixtures for 1M, 60M, and 1B-parameter models using Pile domains and validation loss as the target.
  • 4.1 EXPERIMENTAL SETUP: The experiments use 17 available Pile domains and compare linear regression with LightGBM regression.
  • 4.1 EXPERIMENTAL SETUP: Spearman rank correlation measures agreement between predicted and actual mixture rankings, while MSE measures squared prediction error.
  • 4.2 EXPERIMENTAL RESULTS: 97.12% Spearman correlation was achieved when 1M models trained on 1B tokens predicted unseen mixtures for 1B models trained on 25B tokens.
  • 4.2 EXPERIMENTAL RESULTS: LightGBM outperformed linear regression across all three metrics, with larger advantages on larger models trained with more tokens.
  • 4.2 EXPERIMENTAL RESULTS: 512 proxy models trained on 0.2B tokens outperformed 128 proxy models trained on 0.8B tokens, showing that proxy count can matter more than token count beyond a threshold.

5 EVALUATING ON DOWNSTREAM TASKS

Data mixtures substantially affect downstream performance, and REGMIX identifies strong mixtures by modeling validation-loss signals and searching mixture space. The experiments find web-corpus losses correlate most strongly with downstream results, while domain interactions remain complex and REGMIX remains effective under out-of-distribution evaluation.

  • 5.1 DATA MIXTURE SIGNIFICANTLY IMPACTS DOWNSTREAM PERFORMANCE: 14.6 was the largest performance ∆ across downstream tasks among 64 1B-parameter models trained with different data mixtures.Performance was averaged across 0-shot to 5-shot evaluations.
  • 5.2 WEB CORPORA BENEFITS DOWNSTREAM PERFORMANCE THE MOST: Pile-CC validation loss showed the strongest correlation with most downstream tasks, contrary to the expectation that Wikipedia loss would be most predictive.The Pile-CC correlation with HellaSwag was remarkably close to 1.0.
  • 5.2 WEB CORPORA BENEFITS DOWNSTREAM PERFORMANCE THE MOST: More than 85% of C4100Domain domains showed very strong correlation with Pile-CC, suggesting diverse topical coverage may explain its downstream relevance.The www.ign.com domain closely mirrored Pile-CC’s overall correlation graph.
  • 5.3 DATA MIXTURE BY REGMIX IMPROVES DOWNSTREAM PERFORMANCE: REGMIX improved HellaSwag by 6.8 over Human selection, beat the other three methods in 7 of 14 cases, and achieved the highest average score.The comparison included Pile-CC Only, perplexity filtering, Online Data Mixing, and DoReMi.
  • 5.3 DATA MIXTURE BY REGMIX IMPROVES DOWNSTREAM PERFORMANCE: REGMIX surpassed the best model found in the 64-model comparison, while Pile-CC Only also performed surprisingly strongly against the baselines.These results support automatic mixture selection over random search and reinforce the downstream value of web corpora.
  • 5.3 DATA MIXTURE BY REGMIX IMPROVES DOWNSTREAM PERFORMANCE: REGMIX still outperformed baseline approaches when Pile-CC was excluded from pre-training, supporting robustness when the target domain is out of distribution.The target validation domain was excluded while the remaining domains were used to find the mixture.

6 CONCLUSION

The paper presents REGMIX as an automatic regression-based approach for selecting high-performing data mixtures. It uses small proxy models to predict promising mixtures for large-scale training and studies mixture effects, loss correlations, and the difficulty of human selection.

  • 6 CONCLUSION: REGMIX formulates data-mixture selection as regression, using small models to predict mixture impacts and identify high-performing combinations for large language model pre-training.The method was demonstrated by predicting the best mixture among 64 1B-parameter models.

ETHICS STATEMENT

REGMIX raises ethical concerns because optimizing for performance can skew domain representation and still requires substantial computation. The paper also frames computational efficiency and open availability as contributions toward broader access to data-mixture research.

  • ETHICS STATEMENT: Optimized mixtures may favor some domains while underrepresenting or misrepresenting others, potentially producing biased results.The paper identifies domain imbalance as an ethical risk of performance-oriented optimization.
  • ETHICS STATEMENT: Searching for an optimal mixture still consumes computational resources, increasing energy use and environmental impact.The authors identify reducing this computation cost as an open direction.
  • ETHICS STATEMENT: REGMIX reduces mixture-search overhead to less than 2% of final training costs and releases datasets and trained models publicly.These design choices are presented as supporting more accessible research within academic budgets.
  • ETHICS STATEMENT: The study reports substantial performance variation across 64 mixtures and 12 benchmarks, with domain interactions that can exceed human intuition.PhilPapers is highlighted as a case study of complex domain interactions.

B LIMITATIONS

The paper’s evidence and applicability remain bounded by scale, benchmark coverage, data availability, domain identification, and tokenizer compatibility. Its regression visualization shows strong predictive performance, including nearly 100% rank correlation in reported settings.

  • Rank-invariance scope: The rank-invariance investigation covers 1M-to-1B parameter scales, while testing 3B models would require 64 models trained on 50B tokens each.That experiment would equal training one 3B model on 3.2T tokens and exceeds the authors’ resources.
  • Maximum model parameters: REGMIX has been verified for models up to 1B parameters, leaving 7B and 70B applications for future work because of compute constraints.The authors do not establish that small-model predictions transfer to these larger scales.
  • Benchmark coverage: Benchmark coverage is limited because 1B-scale performance is nearly random on MMLU and negligible on GSM8K, preventing correlation analysis there.The limitation follows from the available Pile data and the relatively small model size.
  • Data availability: Systematically modeling limited data availability remains challenging because most mixture methods assume unlimited data for each domain.The paper suggests combining REGMIX with data-reuse decay as future work.
  • Domain assumption: The method assumes known domain assignments, but assigning examples to domains is difficult when domain boundaries are unclear.This assumption can make application challenging in corpora without clear domain labels.
  • Tokenizer assumption: Proxy and large models are assumed to share a tokenizer and vocabulary size, making cross-tokenizer weight transfer difficult.The paper identifies tokenizer generalization as a fundamental challenge for existing mixture methods.
  • Regression prediction: LightGBM outperforms linear regression and achieves near 100% Spearman rank correlation on 1M-model predictions.The visualization compares predicted and true loss pairs for both regressors.
  • Out-of-distribution evaluation: LightGBM also achieves nearly 100% Spearman rank correlation in the out-of-distribution setting that excludes Pile-CC.This setting evaluates prediction using the remaining domains against Pile-CC validation loss.

C.3 THE DERIVED DATA MIXTURES

REGMIX derives domain weights that can differ from human and baseline methods, while its evaluation reports stronger performance than the Human baseline on six benchmarks. Additional figures examine domain correlations and mixture stability across proxy-model sizes.

  • Derived data mixtures: REGMIX assigns the Pile-CC dataset a weight of 0.87, aligning with human intuition in this experiment.Table 7 compares the derived domain weights across methods.
  • Evaluation results: REGMIX outperforms the Human baseline on 6 benchmarks and has higher average performance than DoReMi and Pile-CC Only.The comparison uses the LightEval evaluation setup described for Table 8.
  • URL-domain correlations: Figures 10–13 visualize correlations between URL domains within C4 subsets and downstream performance.The four figures present the visualization in four parts.
  • Mixture stability: REGMIX produces similar data-mixture distributions when using 1M and 60M proxy models, supporting method stability across these proxy sizes.The figure uses a log-scaled y-axis for visualization.

E IMPLEMENTATION DETAILS

The implementation details specify default hyperparameter usage and provide a table describing model configurations across different model sizes.

  • Implementation settings: Hyperparameters are set to their default values.
  • Model configurations: Table 9 lists the detailed model configuration for different model sizes.
  • Model configurations: The implementation section distinguishes model sizes through their corresponding configuration details.The supplied table caption indicates that configurations vary across model sizes.

F THE STABILITY OF OUR METHOD

RegMix is presented as more stable than prior proxy-model approaches, whose predicted mixtures can vary substantially with proxy-model size.

  • Prior proxy-model methods can suffer from instability, with DoReMi reporting substantially different predicted mixtures for different proxy-model sizes.
  • RegMix’s consistency across proxy-training scales suggests improved stability compared with previous approaches.

G DETAILED EXPERIMENTAL RESULTS

The experiments evaluate RegMix’s regression procedure, rank-invariance assumption, scalability, and performance across model and training-data scales. Results support stable mixture rankings and improvements over human selection, while proxy-size comparisons show that many small proxies remain competitive.

  • RegMix shares the data mixtures and downstream performances of 64 trained 1B-parameter models to support future research.
  • The regression model is fitted from proxy-model evaluations on randomly sampled data mixtures and predicts performance for unseen mixtures.
  • RegMix assumes that relative data-mixture rankings remain stable across model sizes and training-token scales.
  • High Spearman rank correlations across 1M, 60M, 280M, and 1B parameters and 1B, 10B, and 25B tokens support rank invariance.
  • 2% average performance boost: RegMix outperforms Human selection on 7B-parameter models trained on 100B tokens.
  • 50% acceleration on most benchmarks and up to 75% on PiQA: RegMix’s performance benefit persists as training tokens increase, but not on MultiRC.The reported gains concern tasks whose performance increases with training data; MultiRC showed little improvement for either method.
  • Similar average performance across downstream tasks: 128 1B-parameter proxies and 512 1M-parameter proxies produce competitive mixtures for 7B models trained on 100B tokens.The comparison used 1B training tokens per proxy and was limited by computational constraints.
Loading 2407.01492v2…