Data mixing--the strategic reweighting of training domains--is a critical component in training robust machine learning models. This problem is naturally formulated as a bilevel optimization task, where the outer loop optimizes domain weights to minimize validation loss, and the inner loop optimizes model parameters to minimize the weighted training loss. Classical bilevel optimization relies on hypergradients, which theoretically require the inner optimization to reach convergence. However, due to computational constraints, state-of-the-art methods use a finite, often small, number of inner update steps before updating the weights. The theoretical implications of this approximation are not well understood. In this work, we rigorously analyze the convergence behavior of data mixing with a finite number of inner steps T. We prove that the "greedy" practical approach of using T=1 can fail even in a simple quadratic example. Under a fixed parameter update budget N and assuming the per-domain losses are strongly convex, we show that the optimal T scales as Θ(logN) (resp., Θ((NlogN)1/2)) for the data mixing problem with access to full (resp., stochastic) gradients. We complement our theoretical results with proof-of-concept experiments.
Figures & tables
Setting
Optimal Horizon T
Deterministic (Sec. 6.2 )
Θ(logN)
Stochastic (Sec. 6.3 )
Θ(NlogN)
Table 1: Optimal horizon T (leading to best convergence bound) as a function of the total parameter update budget N , when the per-domain losses are strongly convex. In the deterministic (resp., stochastic) case, we have access to full (resp., stochastic) gradients.
Figure 1: Validation loss and the weight of the second domain (most aligned with validation data) w2 as a function of the horizon T , for N={1000,5000} . Note that the validation loss is lowest and w2 is highest when T is larger than 1 and sublinear in N .
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
Category
Value
Inner optimizer
SGD step size η=0.05
Outer (weights) update
Mirror descent step size α=0.5
Hessian approximation parameter
γ=0.01 (recall Hk=γI )
Batch size
batch_size=256
Train subset size
train_subset=12000 examples per domain
Validation subset size
val_subset=3000 examples
Appendix
Table 2: Hyperparameters for Algorithm 2 .
Figure 2: Validation loss, validation accuracy, and the weight of the second domain (most aligned with validation data) as a function of the horizon T , for N=1000 and 5000 .
Data mixture optimization studies how to allocate training examples across domains and tasks to improve model performance under data and compute constraints.
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.
Yuanjian Xu, Tianze Sun, Changwei Xu +7
Hong Kong University of Science and Technology (Guangzhou) · Harbin Institute of Technology University · OpenCSG +1
The capabilities of large language models (LLMs) significantly depend on training data drawn from various domains. Optimizing domain-specific mixture ratios can be modeled as a bi-level optimization problem, which we simplify into a single-level penalized form and solve with twin networks: a proxy model trained on primary data and a dynamically updated reference model trained with additional data. Our proposed method, Twin Networks for bi-level DatA mixturE optiMization (TANDEM), measures the data efficacy through the difference between the twin models and up-weights domains that benefit more from the additional data. TANDEM provides theoretical guarantees and wider applicability, compared to prior approaches. Furthermore, our bi-level perspective suggests new settings to study domain reweighting such as data-restricted scenarios and supervised fine-tuning, where optimized mixture ratios significantly improve the performance. Extensive experiments validate TANDEM's effectiveness in all scenarios.
Jiaxing Wang, Deping Xiang, Jin Xu +9
University of Oxford · Renmin University of China · University of Chinese Academy of Sciences
Recent research has established empirical scaling laws to predict model performance on multi-domain data mixtures. However, a theoretical understanding of these model loss behaviors remains absent. In this work, we propose a unified framework to explain the underlying mechanics of data mixing. Our approach extends theoretical perspectives originally developed for standard neural scaling laws (e.g., Kaplan and Chinchilla) to the multi-domain setting. Based on the distributional assumption that domains overlap on fundamental skills while diverging on specialized skills, we identify two key factors that govern the domain losses of models trained on different data mixtures: Capacity Competition, where the allocation of finite model capacity couples domain losses globally, and Noise Reduction, where optimal weights shift toward harder-to-learn domains to minimize overall noise. Empirical evaluations show that our framework outperforms existing baselines by fitting the loss landscape with a lower Mean Relative Error and identifying higher-performing training mixtures. Most importantly, our model successfully extrapolates across scales, predicting highly effective mixtures for large, unseen scales using parameters fitted on smaller ones. In addition, our model achieves these results using significantly fewer parameters compared to previous empirical laws. Our code is available at https://github.com/meiqwq/Explaining-Data-Mixing-Scaling-Laws.
Rui Dai, Shuran Zheng
Beijing Institute of Technology · IIIS, Tsinghua University