Rethinking Data Mixing from the Perspective of Large Language Models

Yuanjian Xu; Tianze Sun; Changwei Xu; XinLong Zhao; Jianing Hao; Ran Chen; Yang Liu; Ruijie Xu

cs.CLApr 9, 2026

Rethinking Data Mixing from the Perspective of Large Language Models

Authors: Yuanjian Xu, Tianze Sun, Changwei Xu, XinLong Zhao, Jianing Hao, Ran Chen, Yang Liu, Ruijie Xu,

Organizations: Hong Kong University of Science and Technology (Guangzhou) · Harbin Institute of Technology University · OpenCSG · China Mining Group

Abstract

Data mixing strategy is essential for large language model (LLM) training. Empirical evidence shows that inappropriate strategies can significantly reduce generalization. Although recent methods have improved empirical performance, several fundamental questions remain open: what constitutes a domain, whether human and model perceptions of domains are aligned, and how domain weighting influences generalization. We address these questions by establishing formal connections between gradient dynamics and domain distributions, offering a theoretical framework that clarifies the role of domains in training dynamics. Building on this analysis, we introduce DoGraph, a reweighting framework that formulates data scheduling as a graph-constrained optimization problem. Extensive experiments on GPT-2 models of varying scales demonstrate that DoGraph consistently achieves competitive performance.

Figures & tables

Figure 1: PCA projections of gradient directions at different training epochs. Colors denote data domains (C4, Wikipedia, ArXiv, Book, etc.). Initially, gradients form distinct clusters, showing strong domain bias. Over time, they overlap, indicating that the model homogenizes its domain perception. Experiments use 20% of SlimPajama trained on GPT2-Mini.

	Commonsense / Reasoning						RC		LM	Avg
Method	HellaSwag	PiQA	OBQA	COPA	LogiQA	WinoG	SciQ	ARC-E	Lambada
SlimPajama
Uniform	26.1	55.5	11.7	58.0	25.7	49.9	49.0	31.4	11.6	35.4
Dynamic Loss-Based	26.6	56.8	13.8	59.0	29.8	50.1	53.3	31.7	13.4	37.2
DoReMi	26.4	55.7	12.2	59.0	27.2	49.9	53.3	32.3	12.7	36.5
DOGE	26.2	55.8	11.5	62.0	27.2	50.4	52.8	31.3	11.6	36.5

Table 1: Downstream benchmark results ( accuracy % ) on SlimPajama (GPT-2 Medium). Tasks grouped into Commonsense/Reasoning , Reading Comprehension , and Language Modeling . Best results highlighted in bold.

Method	SlimPajama (Val PPL ↓)
Uniform	4.13
DYNAMIC LOSS-BASED	3.10
DoReMi	3.30
DOGE	3.31
RegMix	4.51
Data Mixing Law	4.50

Table 2: Pre-training results on SlimPajama. Validation perplexity (PPL) comparison across domain-mixing strategies. Lower values indicate better generalization.

Figure 2: Perplexity across GPT-2 model sizes.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

	Commonsense / Reasoning						RC		LM	Avg
Method	HellaSwag	PiQA	OBQA	COPA	LogiQA	WinoG	SciQ	ARC-E	Lambada
The Pile
Uniform	29.5	58.8	27.3	65.8	23.9	50.5	60.3	40.0	11.7	40.9
Dynamic Loss-Based	29.0	57.7	26.4	64.3	22.8	49.3	60.0	38.9	10.2	39.9
DoReMi	29.4	58.3	27.3	67.5	26.4	52.2	61.6	40.6	12.1	41.7
DOGE	29.2	58.5	27.1	64.5	23.2	49.8	60.1	40.0	11.7	40.5

Appendix

Table 3: Downstream benchmark results ( accuracy % ) on The Pile (LLaMA-1.1B). Tasks are grouped into Commonsense/Reasoning , Reading Comprehension , and Language Modeling . Our method, DoGraph, achieves consistently better and more balanced results across domains, demonstrating its competitiveness and generalization ability. The best results are highlighted in bold.

	Commonsense / Reasoning						RC		LM	Avg
Method	HellaSwag	PiQA	OBQA	COPA	LogiQA	WinoG	SciQ	ARC-E	Lambada
The Pile
Uniform	29.6	58.8	29.4	66.0	25.9	51.1	61.0	39.1	12.6	41.5
Dynamic Loss-Based	29.3	58.1	29.6	66.1	25.0	52.5	62.7	39.9	12.2	41.7
DoReMi	29.6	58.4	29.8	66.0	24.9	51.4	61.1	40.5	12.8	41.6
DOGE	29.7	56.9	29.2	64.0	25.5	50.6	61.7	40.6	11.9	41.1

Appendix

Table 4: Downstream benchmark results ( accuracy % ) on The Pile (LLaMA-3.2-3B). Tasks are grouped into Commonsense/Reasoning , Reading Comprehension , and Language Modeling . Our method, DoGraph, achieves consistently better and more balanced results across domains, demonstrating its competitiveness and generalization ability. The best results are highlighted in bold.

	Commonsense / Reasoning						RC		LM	Avg
Method	HellaSwag	PiQA	OBQA	COPA	LogiQA	WinoG	SciQ	ARC-E	Lambada
SlimPajama
Uniform	26.0	55.4	13.8	57.2	22.8	49.3	32.6	30.6	12.0	33.3
Dynamic Loss-Based	26.2	56.1	13.2	55.3	26.0	49.2	53.8	31.8	12.6	36.2
DoReMi	26.1	55.7	12.3	53.5	26.8	48.8	52.4	30.9	12.4	35.4
DOGE	26.2	55.0	14.4	60.5	23.5	49.0	31.1	30.8	11.4	33.5

Appendix

Table 5: Downstream benchmark results ( accuracy % ) on SlimPajama (GPT-2 Small). Tasks are grouped into Commonsense/Reasoning , Reading Comprehension , and Language Modeling . Our method, DoGraph, achieves consistently better and more balanced results across domains, demonstrating its competitiveness and generalization ability. The best results are highlighted in bold.

Data Source	Byte Proportion
Common Crawl	54.1%
C4	28.7%
GitHub	4.2%
Books	3.7%
ArXiv	3.4%
Wikipedia	3.1%

Appendix

Table 6: Byte proportion of data sources in the SlimPajama-6B dataset.

Component	Effective Size
Pile-CC	227.12 GiB
PubMed Central	180.55 GiB
Books3	151.44 GiB
OpenWebText2	125.54 GiB
ArXiv	112.42 GiB
Github	95.16 GiB

Appendix

Table 7: Overview of the Pile dataset with datasets that are no longer available due to copyright issues marked in gray. Merged into a single column list.

	GPT-2 Small	GPT-2 Medium	LLaMA-1.1B	LLaMA-3.2-3B
Parameters	210M	300M	1.1B	3B
Layers	24	36	22	28
Attention Heads	16	24	32	24
Embedding Dim.	768	768	2048	8192
Hidden Dim.	3072	3072	2048	3072
Max Seq. Length	512	512	2048	131072

Appendix

Table 8: Model architectures used in our experiments.

	GPT-2 Small	GPT-2 Medium	TinyLLaMA-1.1B	LLaMA-3.2-3B
Minibatch Size	48	48	64	64
Learning Rate $(\times 10^{-3})$	0.50	0.50	0.50	0.50
Learning Rate End $(\times 10^{-4})$	1.0	1.0	1.0	1.0
Warmup Steps	500	500	500	500
$r$	0.4	0.4	0.4	0.4
Training Steps	20,000	20,000	25,000	25,000

Appendix

Table 9: Training hyperparameters for GPT-2 and LLaMA models in our benchmark evaluations.

Figure 3: Evolution of per-sample gradients and the emergence of model-centric structures. Points are colored by their original source datasets. Initially, gradients are separated by domain bias. Over time, the model homogenizes its perception of these sources, leading to significant overlap. Despite this mixing, DoGraph identifies m=11 distinct model-centric domains within the gradient space. Experiments use 20% of SlimPajama trained on GPT2-Mini.

Method	Complexity	PPL
DoGraph (variance)	$O(m^{2})$	3.24
DoGraph (robust softmax)	$O(m)$	3.31
DoGraph (alignment)	$O(m^{2})$	3.15
DoGraph (uncertainty)	$O(m^{2})$	3.09

Appendix

Table 10: Comparison of optimization objectives in DoGraph.

Figure 4: Impact of cluster granularity $m$ on validation perplexity. The U-shaped curve demonstrates that $m=11$ provides the optimal balance; insufficient granularity fails to resolve gradient structures, while excessive partitioning leads to signal inconsistency.

Figure 5: Pre-training GPT-2 Mini on SlimPajama under a 100B-Token Computational Budget. We report the total training time (GPU hours) using $2\times$ NVIDIA H200 GPUs. While our dograph method introduces a sophisticated data-driven decision process, the resulting overhead is minimal (only 4.51% over regmix ), while establishing a new SOTA performance baseline. The marginal increase in budget is well-justified by the superior convergence quality and data selection efficiency.

LLM Training

Large language model training studies the data, objectives, optimizers, and systems used to build and adapt language models.

See all 340 papers

Data Mixing for Large Language Models Pretraining: A Survey and Outlook
Mar 25, 2026Zhuo Chen, Yuxuan Miao, Supryadi +1Data Mixture Optimization Language Model Pretraining
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
DomainPilot: Domain-Level Loss-Guided Two-Stage Data Mixture Optimization for Efficient Language Model Fine-Tuning
Jul 24, 2026He ZhangSupervised Fine-Tuning Large Language Model-Guided Optimization
Tsinghua University
Beyond Data Schedule: Fast LLM Pretraining on Data Mixtures by Loss-Reweighting
Oct 23, 2026Jiazheng Li, Huaqing Zhang, Jing Xu +2Language Model Pretraining Adaptive Loss Weighting
Tsinghua University · ByteDance Seed

Mar 25, 2026J/K move · Enter open · S save

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

Zhuo Chen, Yuxuan Miao, Supryadi, Deyi Xiong

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

Open paper AskPDF arXiv

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.

Abstract

LLM Training

Explore similar work

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

DomainPilot: Domain-Level Loss-Guided Two-Stage Data Mixture Optimization for Efficient Language Model Fine-Tuning

Beyond Data Schedule: Fast LLM Pretraining on Data Mixtures by Loss-Reweighting