Data Mixing as Mixture Experiment: Response Surface Methodology and Optimal Design for Large Language Model Pretraining

Yicheng Mao; Hongru Du

cs.AIAug 24, 2026

Data Mixing as Mixture Experiment: Response Surface Methodology and Optimal Design for Large Language Model Pretraining

Authors: Yicheng Mao, Hongru Du

Organizations: Department of Mathematics and Statistics, University of Calgary, University Drive NW, Calgary, T2N 1N4, Canada · Department of Systems and Information Engineering, University of Virginia, Charlottesville, VA, USA

Abstract

Data mixing is a central design problem in large language model pretraining: given a fixed token budget, practitioners must decide how much data to allocate to each domain. Recent proxy-based methods address this problem by training small models on candidate mixtures, fitting a response model, and using the response to select mixtures for larger-scale training. We show that this workflow has the structure of a classical mixture experiment. Under this view, data domains are mixture components, token shares are component proportions, proxy-training runs are experimental design points, and validation loss defines a response surface over the probability simplex. We develop this formulation using sparse second-order Scheffé response-surface models and construct model-robust $\mathcal{I}$ -optimal designs for proxy data-mixing experiments. Using RegMix as an empirical case study, we demonstrate how the framework can both interpret observed mixture responses and design more efficient proxy experiments. The Scheffé analysis shows that domain value is strongly relational: several domains that are weak under additive effects become favourable through pairwise interactions, especially through combinations with web-derived text. The sparse Scheffé model preserves mixture rankings across model scales and remains competitive with a flexible machine-learning predictor while providing an explicit decomposition of additive and interaction effects. In a simulation study calibrated to observed proxy-training responses, model-robust $\mathcal{I}$ -optimal designs recover the relevant mixture ordering after removing about 25% of the original proxy runs. These results suggest that LLM data mixing should be treated not only as a prediction problem, but also as an experimental-design problem in which the proxy mixtures themselves can be chosen to improve statistical efficiency.

Figures & tables

Label used in this paper	Pile component	Brief description
arxiv	ArXiv	Scientific papers and preprints
freelaw	FreeLaw	Legal opinions and court documents
nih_exporter	NIH ExPorter	NIH grant and project records
pubmed_central	PubMed Central	Full-text biomedical articles
wikipedia_en	Wikipedia (en)	English Wikipedia articles
dm_mathematics	DM Mathematics	Mathematical problem and proof text

Table 1: Pile domain labels used in the RegMix main experiment.

Proxy-based data-mixing workflow	Mixture experiment interpretation
Data domains	Mixture components
Domain token shares $x_{i}$	Component proportions
Constraint $\sum_{i}x_{i}=1$	Simplex constraint
Sampled proxy-training mixtures	Experimental design points
Proxy-model validation loss or downstream score	Continuous response
Predictive model fitted to mixture–performance pairs	Response-surface model

Table 2: Correspondence between proxy-based data mixing and classical mixture experiment methodology.

Figure 1: First-order domain coefficients from the sparse Scheffé response-surface model. Because the response is validation loss, smaller coefficients correspond to lower fitted loss under the additive component.

Figure 2: Selected pairwise interaction coefficients from the sparse second-order Scheffé model. Blue cells denote negative coefficients, corresponding to lower fitted validation loss relative to the additive mixture baseline. Red cells denote positive coefficients, corresponding to higher fitted validation loss relative to the additive baseline. White or near-white cells indicate interactions that are weak or not selected by the sparse model.

Model	Scale	$n$	$\rho$	PRA
First-order Scheffé	1M	256	0.902 [0.869, 0.925]	0.867 [0.863, 0.870]
First-order Scheffé	60M	256	0.892 [0.859, 0.917]	0.859 [0.855, 0.863]
First-order Scheffé	1B	64	0.879 [0.767, 0.942]	0.864 [0.849, 0.878]
Sparse Second-order Scheffé	1M	256	0.937 [0.914, 0.952]	0.894 [0.891, 0.898]
Sparse Second-order Scheffé	60M	256	0.939 [0.916, 0.954]	0.896 [0.893, 0.900]
Sparse Second-order Scheffé	1B	64	0.975 [0.949, 0.985]	0.937 [0.926, 0.948]

Table 3: Cross-scale rank-preservation performance of response-surface models fitted to the RegMix proxy-training data. Values in brackets are 95% percentile bootstrap confidence intervals. The intervals are defined by the 2.5th and 97.5th percentiles of the resulting bootstrap distributions.

Figure 3: Geometry of the model-robust $\mathcal{I}$ -optimal design at $n=512$ compared with the original 512-run RegMix proxy design. Panel A shows the realised mixture proportions for each domain across all design points, with vertical jitter added for visibility. Panel B shows the empirical cumulative distribution function of all realised domain proportions pooled across domains and design points. The RegMix design spreads proportions more continuously over the simplex, whereas the model-robust $\mathcal{I}$ -optimal design concentrates most proportions at the boundary values $0$ , $0.5$ , and $1$ .

Figure 4: Relative

\mathcal{I}

-efficiency of

1{,}000

Dirichlet designs at

n=512

, evaluated under the sparse second-order Scheffé basis with the model-robust

\mathcal{I}

-optimal design as reference. Values below one indicate higher average prediction variance than the model-robust

\mathcal{I}

-optimal design.

Figure 5: Monte Carlo average ranking performance of the model-robust

\mathcal{I}

-optimal designs as a function of design size

n

(solid line with markers), compared with the original RegMix design at

n=512

(dashed line). Left: average Spearman rank correlation

\bar{\rho}_{n}

between the predicted and reference 1B orderings. Right: average pairwise ranking accuracy

\overline{\mathrm{PRA}}_{n}

. Each point averages

1{,}000

simulated datasets generated from the sparse Scheffé surface.

Data Mixture Optimization

Data mixture optimization studies how to allocate training examples across domains and tasks to improve model performance under data and compute constraints.

See all 39 papers

Mar 25, 2026

Data Mixing for Large Language Models Pretraining: A Survey and Outlook

Large language models (LLMs) rely on pretraining on massive and heterogeneous corpora, where training data composition has a decisive impact on training efficiency and downstream generalization under realistic compute and data budget constraints. Unlike sample-level data selection, data mixing optimizes domain-level sampling weights to allocate limited budgets more effectively. In recent years, a growing body of work has proposed principled data mixing methods for LLM pretraining; however, the literature remains fragmented and lacks a dedicated, systematic survey. This paper provides a comprehensive review of data mixing for LLM pretraining. We first formalize data mixture optimization as a bilevel problem on the probability simplex and clarify the role of data mixing in the pretraining pipeline, and briefly explain how existing methods make this formulation tractable in practice. We then introduce a fine-grained taxonomy that organizes existing methods along two dimensions: static versus dynamic mixing. Static mixing is further categorized into rule-based and learning-based methods, while dynamic mixing is grouped into adaptive and externally guided families. For each class, we summarize representative approaches and analyze their strengths and limitations from a performance-cost trade-off perspective. Building on this analysis, we highlight challenges that cut across methods, including limited transferability across data domains, optimization objectives, models, and validation sets, as well as unstandardized evaluation protocols and benchmarks, and the inherent tension between performance gains and cost control in learning-based methods. Finally, we outline several exploratory directions, including finer-grained domain partitioning and inverse data mixing, as well as pipeline-aware designs, aiming to provide conceptual and methodological insights for future research.

Zhuo Chen, Yuxuan Miao, Supryadi +1

The International Joint Institute of Tianjin University, Tianjin University, Tianjin 300072, China · TJUNLP Lab, School of Computer Science and Technology, Tianjin University, Tianjin 300350, China

Data Mixture Optimization Language Model Pretraining

Ask PDF

Jun 17, 2026

RegMix-D: Dynamic Data Mixing via Proxy Training Trajectories

Data mixture selection is critical for Large Language Model pretraining. Existing methods such as RegMix select a single static mixture by fitting a regression model on small-scale proxy runs. We propose RegMix-D, a simple extension of RegMix to dynamic mixing. Our key observation is that proxy runs produce not only endpoint losses, but also full loss trajectories, which can be used to further improve data mixture. By training regression model on these trajectories, we can predict optimal mixtures at multiple training stages. RegMix-D supports two deployment modes: an offline variant that generates a complete mixture schedule before target training, and an online variant that adapts the mixture during training using observed loss. Experiments on 25B tokens of the Pile dataset with a 1B parameter target model show that RegMix-D consistently improves over RegMix and DoReMi across 13 downstream tasks while remaining proxy-efficient: it surpasses RegMix even with only 128 proxy models (25% of RegMix's proxy compute budget).

Kaiyan Zhao, Zhongtao Miao, Akiko Aizawa +1

The University of Tokyo · National Institute of Informatics

Language Model Pretraining Training Data Selection LLM Training

Ask PDF

Mar 9, 2026

Capacity-Aware Mixture Law Enables Efficient LLM Data Optimization

A data mixture refers to how different data sources are combined to train large language models, and selecting an effective mixture is crucial for optimal downstream performance. Existing methods either conduct costly searches directly on the target model or rely on mixture scaling laws that fail to extrapolate well to large model sizes. We address these limitations by introducing a compute-efficient pipeline for data mixture scaling. First, we propose CAMEL, a capacity-aware mixture law that models validation loss with the nonlinear interplay between model size and mixture. We also introduce a loss-to-benchmark prediction law that estimates benchmark accuracy from validation loss, enabling end-to-end performance prediction for the target model. Next, we study how to allocate a fixed compute budget across model scales to fit the law and reduce prediction error. Finally, we apply our method to Mixture-of-Experts models with up to 7B-A150M parameters to fit the law, and verify the optimal mixture derived from the law by extrapolating to a 55B-A1.2B target model. Compared to prior methods, we reduce mixture optimization costs by 50% and improves downstream benchmark performance by up to 3%.

Jingwei Li, Xinran Gu, Jingzhao Zhang

Institute for Interdisciplinary Information Sciences, Tsinghua University · Shanghai Qizhi Institute

Data Mixture Optimization LLM Training Language Model Scaling Laws

Ask PDF

Abstract

Data Mixture Optimization

Explore similar work

Data Mixing for Large Language Models Pretraining: A Survey and Outlook

RegMix-D: Dynamic Data Mixing via Proxy Training Trajectories

Capacity-Aware Mixture Law Enables Efficient LLM Data Optimization