DRAW: Domain Weight Randomization with Bayesian Updating for LLM Pre-Training
Authors: Ruonan Wang
Yongqi Qiao
Zhonglin Xie
Kun Yuan
Venue: Accepted by TMLR ()

Abstract: Optimal pre-training data mixture is pivotal for large language model (LLM) performance, but searching for the best domain weights is computationally expensive. We present Domain Weight Randomization with Bayesian Updating (DRAW), a principled framework treating domain weights as Dirichlet-distributed random variables whose parameters scale with model width. Informative priors are first estimated using proxy models; the main model then refines these using Bayesian inference and parameter scaling, dynamically sampling domain weights during training. Theoretically, DRAW reduces generalization error at a rate $\mathcal{O}(1/\sqrt{n})$ as model width increases, ensuring stable convergence. Empirical results on open-domain corpora and diverse benchmarks show DRAW reliably outperforms fixed and adaptive baselines in both language modeling and downstream tasks, achieving better average and worst-case performance alongside strong robustness. DRAW not only highlights valuable data domains while suppressing noisy ones, but also introduces a scalable and effective mechanism for adaptive data mixing in LLM pre-training, facilitating efficient knowledge transfer from proxy to large models.

Published in Transactions on Machine Learning Research
DRAW: Domain Weight Randomization with Bayesian Up- dating for LLM Pre-Training
Ruonan Wang
wangyn//78nen..et.co
School of Mathematics and Statistics Northeast Normal Universily
Yongqi Qiao
miraclegiao223300qmmaa com
Academy for Advonced Interdisciplinary Studies Peking University
Zhonglin Xie
zloie edu.co
Beijing International Center for Mathematical Research Peking University
Kun Yuan
kunguan @pku. edu.cn
Center for Machine Learning Research
Peking University
Reviewed on OpenReview: https: // openreview. net/ form?? id= tc8TyD72yD
Abstract
Optimal pre-training data mixture is pivotal for large language model (LLM) performance, but searching for the best domain weights computationally expensive. We present Do- main Weight Randomization with Bayesian Updating (DRAW), principled framework treating domain weights Diricheetdiistributed random variables whose parameters scale with model width. Informative priors are first estimated using proxy models; the main model then refines these using Bayesian inference and paramcter scaling, dynamically sam- pling domain weights during training. Theoretically, DRAW reduces generalization error at a rate ((I/v)) as model width increases, ensuring stable convergence. Empirical results on open-domain corpora and diverse benchmarks show DRAW reliably outperforms fixed and adaptive baselines in both language modeling and downstream tasks, achieving better average and worst-case performance alongside strong robustness. DRAW not only high- lights valuable data domains while suppressing noisy ones, but also introduces scalable and effective mechanism for adaptive data mixing in LLM pro-training, facilitating efficient knowledge transfer from proxy to large models.
1 Introduction
The proportion of data sourced from each domain is increasingly recognized as a crucial factor shaping the quality and generalization capability of large language models (LLMs) during pre-training (Brown et al. 2020
Grattafiori et al.] 2024 Zhou et al.] 2025). Recent studies demonstrate that careful adjustment of these mixture ratios can have direct and significant impact downstream task performance and model robustness, particularly as the scale and diversity of training data expand Gu et al. 2025 McKinzie et al. 2024
Graitafiori el. al. 2024). While both heuristic (Gao et al. 2020 Shen et al. 2023) and optimization- based Xie et al. 2023 Fan et al. 2024 Liu et al. 2025) approaches have been proposed to determine optimal mixture ratios, substantial challenges remain. As the number of data domains and model size grow, existing methods are challenged by inefficiency and limited scalability, underscoring the need for more principled frameworks for data mixture adaptation.

---

Published in Transactions on Machine Learning Research (Apr/200) 2026)
Sampling Prior Domain Weights.
2
Trainning Proxy Model and Bayesian Update.
Plaaa BB)) = [1] Iafful
Priori distribution.
Dir(Br). B1
at".
Likelihood function.
Posteniori distribution
porameter
(n of eade ln ofeler (awwiv))
Domain Weight
proxy mode! to obtain
Demain
Optimal Weight
Arxiv 9:
0.4
Somple training data.
Arxiv
d'
0.6
GitHuub 6,
0.3
domain weights,
optimized domain weights.
GitHub
d"
FreelLow 6,
FreeLow
d"
samples for praxy
bucsviuuiie Parepater Flw. Trani
Sample domain weights
Scaled distribution.
Distributions converge toward optimal weights, points become derser.
Posterior distribution,
Domain
Weight
GitHub
Arxiv 4 0.62
4o 0.27
iteration.
Pa
Dynamically adjust training data mixing ratio.
FreeLow
6
Madel exponsion, distribution adaptation!
Trainning Main Model.
Scaling Distribution for Main Model.
Overview of the proposed workflow, illustrated with Arxiv, GitHub, and FreeLaw as example domains, First, based on the width of the proxy model and the number data domains, we determine the parameters of the Dirichlet prior distribution and sample a set of domain weights as the initial weights for the proxy model (Step 1). Next, we train the proxy model to obtain optimized domain weights, and perform Bayesian update to estimate the posterior distribution parameters (Step 2). The resulting posterior distribution is then scaled by factor determined by the ratio of the widths of the main and proxy models, enabling migration of the learned distribution to the main model (Step 3). Finally, during the training of the main model, we sample a set of domain weights from the final distribution at each iteration, use these weights to sample data, and thus dynamically train the main model (Step 4).
Figure 1:
A common approach to transferring data mixture policies across models assumes rank invariance- that is, the optimal mixture ratios determined small proxy model can be directly applied to larger models {Liu et al. 2025). However, this assumption is often violated in practice, as large language models tend to exhibit greater sensitivity to the composition of training data: even slight variations in domain proportions may cause substantial performance fluctuations, particularly for rare domains (Grattafiori et al. 2024]. However, there are also studies that have treated domain weights as adaptive parameters, learning their optimal values either through model scaling or jointly in specific data euvironment (Ye et al.] 2025 Shukor al.) 2025 Ge et al. 2025). Such methods adjust domain mixtures based on scaling laws model-specifie factors, thus alleviating some of the limitations of naive weight transfer and resulting in more robust and generalizable LLMs as both model and data scales increase. Nevertheless, it is important to note that most existing approaches maintain fixed data mixture ratios throughout the entire training process, and do not allow the mixture to evolve dynamically as training progresses or the model grows.
To address this, we propose the Domain Weight Randomization Model (DRAW). Instead of treating do- main weights as fixed constants, DRAW models them random variables drawn from distributions whose parameters are tied to model size. This reframes proxy model training from finding set of optimal weights to dynamically updating parameters of a weight distribution via Bayesian framework (van de Schoot al., 2021). We further introduce a scaling factor to adapt the paramcterized distribution to the main model. At each pre-training step, domain weights are sampled from this distribution to improve adaptability and gener- alization. DRAW consists of four primary steps (Figure Do: (1) sampling the initial weight distribution based on the proxy model width; (2) training the proxy to obtain optimized weights; (3) utilizing Bayesian updat- ing and scaling adjust the distribution for the main model; (4) during main model training, dynamically sampling weights to guide data mixing, enabling robust adaptation.
2

---

Published in Transactions on Machine Learning Research
Our experiments demonstrate that DRAW substantially improves the adaptability and generalization of large-scale model pre-training, outperforming mainstream alternatives in both average and worst-case val- idation loss and surpassing them on several downstream tasks (e.g., MMLU (Hendrycks et al. 2021). GSM8K (Cobbe et al.) 2021)). The method achieves effective weight stratification and sparsification locusing on key domains while suppressing noise and irrelevant data and both theoretical analysis and empirical results confirm reduced parameter estimation errors and stable training dynamics. Furthermore, DRAW exhibits remarkable robustness the choice of proxy model size and random initialization, con- sistently delivering
results with minimal performance variance across seeds. These findings validate the scalability and reliability of the proposed framework for adaptive domain weighting in large language model pre-training.
Our main contributions are as follows:
We propose Dirichlet-based domain weight randomization method, achieving more principled and effective data mixture strategies for large language model pre-training.
We develop a Bayesian updating and parameter scaling framework, which improves the stability and accuracy of domain weight transfer across model scales.
We provide theoretical analysis showing that randomized domain weights and Bayesian adaptation guarantee stable convergence and reduce estimation error at scale.
We experimentally show that DRAW consistently outperforms other mainstream methods, demon- strating its superior generality and robustness.
2
Related work
Our work builds on prior studies that explore how to design and adjust data mixing strategies for multi- domain pre-training, with particular focus on transferring effective mixing policies learned from small-scale proxy models to large-scale foundation models. We identify three main research directions relevant to this goal: data mixing, scaling laws and randomization of domain weights.
The scaling laws for large models, proposed by OpenAI in 2020 (Kaplan et al. 2020 It describes power-law relationship between training loss and factors such computational resources, model parameterization, and training data size (Hoffmann et al. 2022). In 2023, Greg Yang al. introduced Maximum Update Parameterization (MUP), which enables the optimal parameters determined on small proxy models to predict the loss of larger models through a modified power law (Yang & Hu 2021).
Scaling Laws.
The existing literature on data mixture scaling laws gencrally takes one two perspectives. One approach assumes that the optimal domain weights learned by small-scale models can be directly transferred to large- scale models without modification (Xie al.] 2023 Fan al.] [2024] Liu et al.) (022).. However, this static allocation overlooks the differences introduced by model scaling. Another approach fits domain weights parameters according to scaling laws, yet this method struggles to handle high-dimensional settings with many data domains (Gu et al. 2024 Goyal et al. 2024 Béthune et al.) 2025). To address these issues, we introduce a scaling factor that dynamically adjusts the domain weight distribution parameters when transferring from small models to large models. This enables effective adaptation of domain weights across different model sizes, thus enhancing the flexibility and generalization capability of large-scale language model pre-training.
Weight Randomization and Data Augmentation. Random parameter initialization is common technique in deep learning used to enhance exploration and prevent problems such as vanishing or exploding gradients during training (Goodfellow et al. 2016). In the computer vision domain, data augmentation strategies such as mixup (Zhang et al. 2018). CutMix (Yun et al. 2019),, and SnapMix (Huang et al. 2021) have shown remarkable effectiveness by mixing training data samples according to mixing coefficients sampled from a Beta distribution. This process of stochastic data mixing not only increases the diversity of training examples, thus improving generalization, but also acts as form of regularization to reduce overfitting. We extend these ideas beyond simple two-way mixes: for multi-domain data in language model
3

---

Published in Transactions on Machine Learning Research
pre-training, domain weights can be considered random vectors on the simplex, naturally modeled by the Dirichlet distribution, a multidimensional generalization of the Beta distribution (Bae et. al. 2024). This allows adaptive sampling of domain proportions when constructing each batch, enabling the model to exploit both the diversity and hetcrogencous quality of different data sources during large-scale pro-training.
Dynamic Refinement Methods. DRAW is also related to recent methods (Jiang et al. 2024 that refine data or domain weights online during training based on model feedback. The key distinction is that DRAW separates estimation from sampling: proxy model is first used to infer an informative domain prior, and the main model is then trained by stochastic sampling around this prior. In contrast, online refinement methods update weights directly from short-term training dynamics. This proxy-first design aims to retain the benefits of adaptive sampling while reducing sensitivity to transient minibatch fluctuations.
3 Method
3.1 DRAW
Setup. Suppose we have a dataset Derain A (Di,..., D&} consisting a mixture of distinet domains, where each domain D; has an associated dataset. Each domain D; is assigned a domain weight a;, and together these weights define the data sampling distribution Pa Σ ag (n,) where unif(D) [D] ExeD 6 denotes the uniform distribution over dataset D, and (x') 1 if x' = æ and otherwise. The domain weight vector a _ fai...., ax] E A is treated as a random variable drawn from a Dirichlet distribution.
We consider Transformer-baset language models that share the same architectural depth (i.e., number of layers) but differ in width: smaller proxy model with width n1, and a larger main model with width 72. Here, width refers to the number of hidden units per layer, as is standard in prior work (Yang & Hu 2021] For mainstream LLMs such as Qwen3 (Qwen Team 2025), increasing model width expands the hidden state, the intermediate feedforward network (FFN) dimension, and the input/outpin embedding sizes, while keeping the number of layers fixed. The proxy model is first trained to optimize domain weight distributions suitable for its scale, and these learned distributions are then used initialize the main model. The main model is subsequently trained minimize the training loss on the pre-defined objectives, dynamically adapting data mixture via the domain weights informed by the proxy.
Step 1: Sampling Prior Domain Weights.
We begin by establishing a Dirichlet prior over the domain weights, with concentration parameters set as a function of both the proxy model width n1 and the number of domains k. Specifically,
V VTI (VT k
B1
(1)
k
Domain weights are then sampled as a ~ Dir(Bi). Note that this parametcrization does not change the mean of the Dirichlet distribution, which remains uniform over domains, i.e., E[az] = 1/kk for all i. Instead, increasing nn increases the total concentration of the prior, causing samples to cluster more tightly around the mean and reducing sampling variance. In other words, a larger proxy model induces a sharper and more confident prior over domain mixtures, rather than a more uniform one. This design reflects the intuition that stronger proxy models can provide more stable estimates of desirable data mixtures, while smaller proxies may benefit from broader stochastic exploration. The sampled prior then serves as principled starting point for subsequent Bayesian refinement of the domain weights.
Step 2: Training Proxy Model and Bayesian Updating.
To estimate the parameters governing the domain weight distribution, we combine a DRO-style proxy opti- mization step with Bayesian updating. At the proxy stage, we solve
min max C(0, a) a
Σ a Di
Σ [(z; zED;
(2)
4

---

Published in Transactions on Machine Learning Research (Apr/2000
Samples
Samples
Meen (Center)
Mean (Center)
0.01
0.00
-011
0.10
0.05
010
0.05
0.55
0.00
PC2
0.00
PC2
-0.05
0.30
-0005
0.10
-0.10
Figure 2: 3D visualization of domain weight samples over the 17 sub-domains in The Pile. Left: untrained random weights from a Dirichlet distribution with parameter /n2 model after DRAW training.
k 1. Right: learned weights of the main
where 0); (x))) (1r the excess loss relative to reference model (Xie et al. 2023). This formulation is inspired by Group DRO (Sagawa et al. 2020 robust reweighting, but in DRAW it is used only to infer proxy-level domain prior. The resulting optimized weights a are then treated pscudo- counts for Bayesian updating, after which the main model is trained via stochastic Dirichlet sampling around the proxy-informed prior rather than by directly optimizing a fixed DRO mixture.
To refine the distribution, we perform a Bayesian update using a pscudo-obserrations,, yielding a posterior Dirichlet distribution:
P(a a², P(o") a) P(a (1) x
II
a
Bato -1
which is Dirichlet with updated concentration parameters ".T B1 a.. This posterior integrates both prior knowledge and observed domain importance, establishing statistically principled basis for the mixing distribution used in the subsequent main model training.
Step 3: Scaling Distribution for Main Model.
Instead of directly applying the domain weight distribution optimized for the proxy model (width n1), we scale the posterior Dirichlet parameters to account for increases in model capacity. Given the updated concentration vector grew, we define the main model's parameters as:
n2
B2
gTew
(4)
where 32 specifies the Dirichlet distribution for the main model width n2. This scaling results a sharper domain weighting distribution, i.e., lover-variance sampling around the proxy-informed mean mix- ture. Specifically, each Dirichlet parameter be written
A = 122
12 k
a
+
1,
(5)
vn
where is k-dimensional vector of oncs. This procodure transfers the proxy-ostimated mixture the main model in a scale-aware stochastic form.
Remark Aithough in Eq. (5) the uniform baseline term L may dominate in magnitude when is small ny n2, this does not mean that the influence of a² is diminished. In. Dirichlet distribution, each
5

---

Published in Transactions on Machine Learning Research
concentration parameter affects the mean via
o
E
[ay
71 +
Thus, a" shifts the expected proportions away from the uniform distribution, even when its scale is numeri- cally much smaller than the baseline term.
This design is intentional: a large baseline concentration ensures low-variance and slable domain weights, preventing overfitting to the fine-grained biases learned by large prozy model, while o still provides di- rectional bias toward the prozy-learned structural preferences. As nı grows, the relative size of the a" term decreases, reflecling a robuslness mechanism rolher Ihan loss of useful information. Increusing the concen- tration of the Dirichlet distribution does not make the minture more uniform; instead, reduces sampling variance around the promy-informed mean. As observed in Suriya et al [Gunasekar et ol.) 2023), this appears beneficial for larger main models, which benefit from higher signal-to-noise ratio in domain sampling while still relaining controlled stochastic exploration.
Step 4:
Training Main Model.
At each iteralion of main model training, domain weights are dynamically sampled from the scaled Dirichlet prior Dir(B2) (see Figure 2), and data batches are drawn accordingly. As illustrated in Figure 2 the learned domain weight distribution after scaling is both more concentrated and noticeably non-uniform, with prominent biases towards certain domains. Such concentration enables the model to focus on more informative domains, thereby further enhancing the robustness and efficiency of the DRAW framework under heterogeneous data distributions.
3.2 Theoretical Guarantees
Preliminaries and Notation. Let F(O) constitute the objective function parameterized by 0 E IRd, and let ||- denote the Euclidean (f2) norm. The following definitions are used:
Simplex (As): The probability simplex of dimension k- 1 is denoted as A. = far E R E a; = 1,0; 2 0}.
Strong Convexity: is p-strongly convex if for any 0, E Rdd and À E [0, 1],
F(A0 + (1 - A)0') s + - P = X(1 A) - -
I-smoothness: is L-smooth is differentiable and |VF VF(0')| s I.|0 -0'l for all 0,0' ER.
Convergence in Probability: sequence { converges in probability to 0*, denoted as Or P 0", if limy p 0* c) for any e >>> 0.
Asymptotic Notation (0): For a quantity æ depending on scaling parameter (e.g., width), == 0(1) implies that there exist constants C1,C2 > 0 such that C₁ |æ| C2 for sufficiently large n.
The loss function L(0, a) is My-Lipschitz with respect to a if L(O, T((),o')) & Mi a a'l| for any a, a' E Ag-
Mi-Lipschitz Continuity:
We provide a rigorous theoretical analysis of DRAW under adaptive, randomized domain weights (sce Ap- pendix C for proofs). Specifically, we derive formal guarantees governing model output stability, optimization trajectory convergence, and generalivation error scaling.

---

Published in Transactions on Machine Learning Research
Lemma 1. Based on the Tensor Progrom and Marimal Update Parametrization, the size of the model output J after t ilerations is (11).
Lemma demonstrates the inherent stability of the DRAW approach. By theoretically bounding the vari- ance of the stochastic gradients induced by domain weight sampling, it ensures that the training process remains controlled, effectively mitigating risks of optimization divergence. This theoretical property is cor- roborated by the stable losses observed Figure 5 Moreover, this result justifies our choice of Dirichlet hyperparameters, as they directly regulate the trade-off between sampling diversity and training stability.
Thcorem Suppose F(Ø) = Ea [a[ ai) is p-strongly convez and L-smooth. If the parameters 0 are updated using slochastic gradients compuled with domain weights ar Dir(B) using slep size sequence {ne) satisfying Σ and Σ n < 00, then the optimization guarantees convergence in mean square to the unique minimizer 0* of F()):
lim E |100- 0*112) 0. T-000
0r 4 0".
(6)
Theorem [ shows that, under standard strong-convexin and smoothness assumptions, stochastic optimiza- tion with Dirichlet-sampled domain weights remains mean-square convergent to the minimizer of the expected objective F(0). The result should be interpreted as stability guarantee for randomized domain weight- ing: the additional stochasticity introduced by sampling ag does not, by itself, prevent convergence under a standard diminishing-steyy
regime. Importantly, this theorem does not claim that practical large-scale pretraining is globally convex, nor that stochastic weighting alone guarantoes better final model quality; rather, formalizes that DRAW's randomized weighting mechanism is compatible with stable optimization in idealized setting.
Theorem 2. Let a)) be the loss function of a model with width n2 parameterized by 0, under domain misture weights Assume the loss is M1 -Lipschcit regulor with respect to Define the expected risk as:
F(0) := (C(
where = + a² (implying concentration parameter n1). Let 0* denote the population minimizer and Ô the parameter obtained after optimization.
The generulization error, governed by the interplay belween domain variance (scale nı) and model width (scale n2), is bounded by:
B(BB 3i) 33) (2 + +m1)
F(O) - F(0)) S M1
Σ
where = + is the sum of Dirichlet parameters. In the regime of large data concentration (n1 >> 1) and wide models (nz > 1), the scaling law simplifies to:
F(0) =0 - (0(( M
(
Remark The theoretical bound specifically isolates the generulization error induced by the stochasticity of Dirichlet-sampled data miring, independent of optimization noise (e.g., SGD randomness). Assuming the model reaches a parameter configuration close to the population minémizer, this upper bound primarily reflects the variance introduced by data mirture variability.
Theorem 2 bounds the excess risk induced by stochastic variation in Dirichlct-sampled domain mixtures, and shows that this contribution decreases as either the proxy concentration scale ii or the main model width n2 grows. In particular, the 0((nnn))-1/3 rate should be interpreted as a variance-control result: larger n, yields a more concentrated sampling distribution, and larger n2 reduces the impact of this mixture stochasticity under the theorem's 's assumptions. This statement does not contradict that larger models can be more sensitive to data composilion. The theorem concerns robustness to random fluctuations around

---

Published in Transactions on Machine Learning Research (Apr/2026
DM Mathematiccs 15 49GIB
EuroFan
9.17GIB
PubMee Abstracts 38.53GIB
BoakCorpus?
9.45GiB
Wikipedia (en) 19. 13GIB
95 16GiB
Ubuntu
PubMed
11 11.03GiB
USPTO Backgrounds
Gutenberg (PG-19 19)0penSubiitles 27.19Gi8
19.47Gi8
ArKiv 112 42GiB
FreeLaw 76.73GIB
Stack Exchange 64, 64.39GG
227.126i8
Books3 151.44GiB
OpenWebText2 125. 54GIB
Figure 3: Treemap of The Pile's 17 subsets by data volume.
given mixture policy. A larger model may therefore be more robust to stochastic sampling noise while still being more sensitive to whether the average domain allocation is well chosen.
This asymmetry confirmed empirically in Table 2 Increasing the proxy size from 150M to 510M results in only marginal improvements in validation loss for a 1B main model (a decrease of merely ~ 0.015). This indicates that the proxy model has entered regime of diminishing returns, where the loss landscape becomes insensitive to further reductions in weight variance. Conversely, increasing the width of the main model n2 yields substantially greater gains (e-g., drop of 0.53 from 150M to 1B), highlighting that main model capacity is the docisive factor for improving generalization performance.
Consequently, the proxy model only needs to be large enough to ensure accurate weight estimation; scaling it beyond this sufficiency threshold is computationally inefficient. To effectively minimize generalization error, computational resources are best allocated to scaling the main model.
4 Experiments
4.1
Experimental Setup
Datasets and Models. We utilize 17 copyright-free subdomains from The Pile (Gao al. 2020 dataset for training. with specific domains and their respective sizes detailed in Figure We employ decoder- only Transformer language models (Vaswani al. 2017 based on the GPT architecture, trained using the standard next-token prediction objective.
Training Setup. Our primary configuration consists a 70M-paramcter proxy model (hidden dimension 512) and a 1B-parameter main model (hidden dimension 2048). To conduct the ablation studies presented in Table 3 we vary both the proxy and main model sizes. Unless otherwise stated, all models (including both proxy and main models) sequence length of 1024 tokens, corresponding to approximately 13.1B training tokens in total. Additional hyperparameters are detailed in Appendix D.
trained for 50,000 optimization steps with a global batch size of 256 and
Evaluation. We assess model performance using the held-out validation set of The Pile to measure validation loss. For downstream evaluation, we employ a comprehensive suite of benchmarks, including MMLU (Hendrycks et al. (022). MMLU-Pro (Wang et al.] 2024h, Natural Questions (NQ) (Kwiatkowski et al. 2019), GSM8K (Cobbe et al.) 2021), BoolQ, (Clark al. 2019). ARC-C (Clark al.) 2018). and DROP. We adopt 5-shot setting for MMLU, MML.U-rro, NQ, and ARC-C, 10-shot setting for GSMSK,
8

---

Published in Transactions on Machine Learning Research
and 0-shot setting for DROP. We follow standard evaluation protocols and report accuracy for all bench- marks. We note that, for models at this scale, performance on some difficult downstream benchmarks can be close to chance level in absolute terms. Accordingly, we interpret downstream accuracy primarily as a complementary signal, while using held-out validation loss (or perplexity) on The Pile as the main metric for assessing the quality of domain reweighting.
Baselines. We benchmark DRAW against several domain weighting strategies, including the official Pile weights (Baseline) (Gao et al. 2020 DoReMi (Xie al. 2023), RegMix (Liu et al 2025), and random domain weights (Random). All methods evaluated under identical settings on both the Pile validation set and downstream benchmarks.
4.2
Experimental Result
9x
Baseline Doremi
Baseline
Average Validation Loss
100
8 10
Doremi
Draw
Worst Validation Loss
Draw
5 x
Regmix Random
x
Regmix
Random
6 10
4x
10%
x000
4 10
10k
20k
30k
40k
50k
10k
20k
30k
40k
50k
Training Steps
Training Steps
Figure 4: Average validation loss across domains in Pilo-CC (left) and worst-casc validation loss on the most challenging domain in Pile-CC (right).
Accuracy (%) various downstream benchmarks. For each benchmark, the best result is high- lighted in bold.
Table 1:
Benchmark
DRAW
Baseline
Random
DoReMi
RegMix
MMLU
31.58
24.38
25.44
23.28
25.33
MMLU-PRO
16.61
9.77
10.61
9.14
10.15
NQ
6.02
5.09
2.31
3.24
6.02
GSM8K
1.60
1.47
1.47
2.41
2.51
BoolQ
5.11
8.17
5.02
7.43
7.75
ARC-C
25.00
25.88
19.30
28.95
22.81
Average
14.32
12.46
10.69
12.41
12.43
Superior Language Modeling Performance. The DRAW-traint main models consistent outperform those using manually tuned fixed domain weights on standard language modeling metries, achieving lower average and worst-case validation losses The Pile (see Figure 4 Notably, our experiments reveal that models with purely random domain weights also surpass weight baselines, validating the underlying benefit of stochastic regularization in exploring the loss landscape. However, DRAW significantly improves upon blind randomization. By leveraging informed Bayesian updating and parameter scaling, DRAW effectively guides the optimization trajectory, resulting in faster convergence and lower final loss than the Random baseline. Furthermore, despite the dynamic nature the weighting process, results over multiple random seods demonstrate minimal variance in downstrcam performance (sce Figure 5). This empirically confirms the stability and reproducibility of our method, validating our theoretical guarantee that DRAW converges to consistent global minimizer.

---

Published in Transactions on Machine Learning Research
Robust Downstream Transfer. We evaluate the downstream task performance of DRAW using prehensive suite of benchmarks following standard accuracy-based protocols. As shown in Table 1 DRAW achieves the best results on most benchmarks particularly on knowledge-intenssive tasks such as MMLU, MMLU-PRO, and NQ-and significantly improves average accuracy over all baseline methods, whether fixed or random. These results demonstrate that DRAW enables robust transfer of domain knowledge from the proxy to the main model. By minimizing the distribution mismatch (parameter estimation error), DRAW ensures that the main model is trained on the most inforhation-rich data mixture, thereby fully utilizing its capacity to enhance pre-training performance.
Table 2: Domain weights learned by different. methods. Comparison implies a distinct. preference change: While DoReMi and RegMix heavily rely on the noisy Pile-CC domain, DRAW shifts focus high-quality knowledge sources like Wikipedia and StackErchange, effectively filtering out web noisc
Domain
Baseline
DoReMi
RegMix
DRAW
Wikipedia (en)
0.0204
0.0699
0.0159
0.4171
StackExchange
0.0685
0.0153
0.0003
0.2477
Euron Emails
0.0019
0.0070
0.0017
0.2770
PubMed Central
0.1922
0.0046
0.0031
0.0241
Github
0.1014
0.0179
0.0002
0.0213
DM Mathematics
0.0166
0.0018
0.0003
0.0086
ArXiv
0.1195
0.0036
0.0012
0.0010
FreeLaw
0.0817
0,0043
0.0015
0.0008
PubMed Abstracts
0.0410
0.0113
0.0242
0.0006
EuroParl
0.0097
0.0062
0.0000
0.0004
USPTO Backgrounds
0.0487
0.0036
0.0025
0.0003
Pile-CC
0.2418
0.6057
0.8701
0.0003
Gutenberg (PG-19)
0.0290
0.0072
0.0016
0.0002
NIH ExPorter
0.0040
0.0063
0.0012
0.0002
PhilPapers
0.0051
0.0274
0.0000
0.0002
Ubuntu IRC
0.0117
0.0093
0.0642
0.0002
HackerNews
0.0083
0.0134
0.0119
0.0001
Interpretability of Learned Weights: Quality over Quantity. To investigate the source of DRAW's performance gains, we compare the optimal domain weights learned by different methods (Table 2 The results reveal a striking divergence in data selection strategies. First, DRAW acts as an aggressive noise filter. While conventional methods like DoReMi and RegMix assign dominant weights (approximately 60% 87%) to Pile-CC a massive but noisy web-crawl dataset DRAW effectively prunes this domain, assigning negligible weight of 0.0003. Second, DRAW amplifies high-signal domains. It reallocates the probability mass to structured, information-denss sources, with over 94% of the total weight concentrated on just three domains: Wikipedia (en) (41.7%), Enron Emails (27.7%), and StackErchange (24.8%). This automatic stratification suggests that DRAW's proxy model successfully identifies the 'signal-to-noise ratio of each domain relative to the target distribution, avoiding resource wastage on low-quality web text in favor of knowledge-rich data.
Impact of Proxy and Main Model Scale. To validate the theoretical scaling properties of DRAW, conduct two complementary experiments analyzing the decoupling of proxy estimation and main model representation. The results are summarized in Table
Diminishing Returns of Proxy Size (nI). The left section of Table 3 presents the performance of fixed IB main model trained with weights derived from proxies increasing sizes (150M to 510M). We observe that increasing the proxy size yields only marginal gains (c.g., Average Validation Loss decreases from 2.9803 2.9650). This empirically confirms our theoretical bound that the estimation error decays slowly with nn, suggesting that moderate-sizel proxy is sufficient to estimate the optimal distribution direction.
10

---

Published in Transactions on Machine Learning Research
Mean: 2.96995
2.974
Min: 2.96844
Max: 2.97202
Average Validation Loss
2.97202
2.972
Average Validation Loss
2.97126
+0.00207
+0.00131
2.970
-0.00105
00000011
00000151
2.96914
.2.96890..
2.968
2.96844
2.966
17
42 Seed
115
123
Figure 5: The average validation loss on The Pile for main models trained with DRAW, where the proxy model's domain weights are initialized using different random seeds.
Dominance of Main Model Size (n2). In sharp contrast, the right section of Table shows the effect of scaling the main model (from 150M to IB) using weights from a fixed, small 70M proxy. Here, the performance gains are substantial (Loss drops from 3.5002 to 2.9689). This verifies the core insight of our Scaling Law (Theorem : the main model's capacity (n2) the dominant factor in generalization. Once the proxy provides reasonably accurate direction, allocating compute scale the main model yields the highest return investment.
Table 3: Impact of model scaling on validation loss. Left: Scaling the Proxy Model (n1) while keeping the Main Model fixed 1B shows diminishing returns. Right: Scaling the Main Model (n2) using weights from small 70M proxy yields significant performance gains, confirming that ny the dominant factor.
Effect of Proxy Scaling (Fixed Main Model: 1B)
Effect of Main Model Scaling (Fixed Proxy: 70M)
Configuration
150M
280M
510M
150M
280M
510M
1B
Avg. Val Loss
2.9803
2.9753
2.9650
3.5002
3.4885
3.2707
2.9689
Worst Val Loss
4.0216
4.0153
4.0031
4.4913
4.4844
4.2972
3.9905
Stability and Reproducibility. A potential concern with stochastic reweighting is performance variance. To address this, we evaluate the robustness of DRAW by initializing the proxy model's domain weights with different random seeds and transferring the learned weights to the main model. As illustrated in Figure 5 the average validation loss The Pile remains remarkably stable, fluctuating narrowly between 2.96844 and 2.97202. The negligible deviation (< 0.004) confirms that despite the stochastic nature of the Dirichlet sampling process, DRAW converges to a consistent global solution. This stability effectively validates the reliability of the DRAW framework for large-scale pre-training, ensuring robust performance independent of initialization noise.
Dynamic vs. Static Reweighting. We compare DRAW with a static variant that fixes proxy-derived domain weights during main-model training. In the 70M proxy - 150M setting, dynamic DRAW achieves
11

---

Published in Transactions on Machine Learning Research
average validation loss of 3.5002, compared with 3.6415 for the static variant. This suggests that DRAW benefits not only from good proxy-informed mixture, but also from stochastic exploration during training.
5 Discussion
5.1
Conclusion
We introduced DRAW, a framework for domain weight selection LLM pre-training that models domain mixtures as random variables in Bayesian framework. DRAW turns data mixing into a controlled stochastic reweighting problem and enables transfer from a proxy model to larger main model.Our main contributions are threefold:
Scale-aware domain weighting. We show theoretically and empirically that useful domain weights can be estimated with a proxy model and transferred across scales, decoupling weight estimation from main-model training.
Adaptive filtering of heterogeneous data. DRAW downweights lower-quality sources (e.g., Common Crawl) while emphasizing more informative domains.
Improved training behavior. Experiments show that DRAW yields better convergence and stability than fixed-weight and naive random baselines.
Overall, DRAW provides a practical and principled approach to scale-aware domain reweighting for large- scale pre-training.
5.2 Limitations
Closed-world assumption. DRAW assumes the target distribution lies within the span of the observed training domains, and is therefore not designed for fully out-of-distribution settings.
Modest gains in some settings. Although DRAW consistently improves validation perplexity and outper- forms weight baselines, its advantage over random stochastic weighting sometimes limited, suggesting that part of the gain comes from stochastic exploration itself.
Theory uses simplifying assumptions. Our analysis relies on standard assumptions for tractability and should be viewed as providing intuition and statistical support rather than complete account of modern LLM optimization.
Limited domain interaction modeling. The current Dirichlet formulation does not explicitly model richer correlations among domains beyond the simplex constraint.
Small-scale downstream evaluation. Some downstream results for smaller models are chance level difficult benchmarks, validation perplexity remains the main metric for assessing reweighting quality.
5.3 Future Work
Online reweighting. An online version of DRAW that updates domain weights during main-model training could improve adaptivity and connect more directly dynamic refinement. methods.
Richer priors. Future work could explore hierarchical correlated priors, such as Dogisticnnrrml distri- butions, to better capture structure across domains.
Multimodal extension. Extending DRAW to multimodal pre-training is a promising direction, especially given the different scaling behavior and noise profiles of text, image, and video data.
Task-aware reweighting. Incorporating downstream task feedback into the reweighting loop may enable fask-specific domain mixture optimization.
12

---

Published in Transactions on Machine Learning Research
References
HecSun Bac, Seungjae Shin, Byconghu Na, and II-haul Moon. Dirichlet-basod per-sample weighting by transition matrix for noisy label learning. In International Conference on Learning Represenlations, 2024.
Louis Béthune, David Grangier, Dan Busbridge, Elconora Gualdoni, Marco Cuturi, and Pierre Ablin. Scaling laws for forgetting during finetuning with pretraining data injection. In International Conference on Machine Learning, 2025.
Tom Brown, Benjamin Mann, Nick Ryder, Mclanic Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. In Advances in Neural Information Processing Systems, pp. 1877-1901, 2020.
Christopher Clark, Kenton Lee, Ming-Wa Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. BoolQ: Exploring the surprising difficulty of natural yes/no questions. In Proceedings of the 2019 Conference of the North American Chapler of the Associalion for Compulational Linguislios: Human Language Technologies, pp. 2924 2936. Association for Computational Linguistics, 2019.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? Try ARC, the A12 reasoning challenge. arXiv preprint arXiv:1803.0544 2018.
Karl Cobbe, Vineet. Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Malthias Plappert, Jerry Tworek, Jacob Hilton, Reüichrro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint 1xxx:21100144000 2021.
Simin Fan, Matteo Pagliardini, and Martin Jaggi. DOGE: Domain reweighting with generalization estima- tion. In Proceedings of the 41st International Conference on Machine Learning, pp. 12895-12915, 2024.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa. Nabeshima, Shawn Presser, and Connor Leahy. The Pile: An 800GB dataset of diverse text for language modeling. arXie preprinl arXiv:2101.00027, 2020.
Ce Ge, Zhijian Ma, Daoyuan Chen, Yaliang Li, and Bolin Ding. Bimix: A bivariate data mixing law for language model pretraining. In International Conference on Learning Representations, 2025.
Tan Goodfellow, Yoshua Bengio, Aaron Courville, and Yoshua Bengio. Deep learning, volume MIT Press, 2016.
Sachin Goyal, Pratyush Maini, Zachary C Lipton, Aditi Raghunathan, and J Zico Kolter, Scaling laws for data filtering: Data curation cannot be compute agnostic. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 02702-22711, 2024.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al- Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, Amy Yang, and et Angela Fan. The Llama Herd of Models. arxiv preprinl ar orXiv:2407.21756 2024.
Jaawe Gu, Zace Yang, Chuanghao Ding, Rui Zhao, and Fei Tan. CMR scaling law: Predicting critical mixture ratios for continual pre-training of language models. In Proceedings of the 2024 Conference Empirical Methods in Natural Language Processing, pp. 16143 16162, 2024.
Xinran Gu, Kaifeng Lyu, Jaazheng Li, and Jingzhao Zhang. Data Mixing Can Induce Phase Transitions in Knowledge Acquisition. In ICLR 2025 Workshop on I Can't Believe It's Not Beller, 2025.
Suriya Gunasekar, Yi Zhang, Jyoti Aneja, Caio César Teodoro Mendes, Allie Del Giorno, Sivakanth Gopi, Mojan Javaheripi, Picro Kauffmann, Gustavo de Rosa, Olli Saarikivi, Adil Salim, Shital Shah, Harki- rat Singh Behl, Xin Wang, Sébastien Bubeck, Ronen Eldan, Adam Tauman Kalai, Yin Tat Lee, and Yuanzhi Li. Textbooks are all you need. arXiv preprint or Xix:22300. 11644, 2023.
13

---

Published in Transactions on Machine Learning Research
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Stein- hardt. Measuring Massive Multitask Language Understanding. In International Conference on Learning Representations, 2021.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, et al. Training compute-optinal large language models. In Advances in Neural Information Processing Systems, pp. 30016 30030, 2022.
Shaoli Huang, Xinchao Wang, and Dacheng Tao. SnapMix: Semantically proportional mixing for augmenting fine-grained data. In Proceedings of the AAAI Conference Artificial Intelligence, pp. 1628 1636, 2021.
Yiding Jiang, Allan Zhou, Zhili Feng, Sadhika Malladi, and J. Zico Kolter. Adaptive data optimization: Dynamic sample selection with scaling laws. arXiv preprint 2rXiv:2/10.11822 2024.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jesfrey Wu, and Dario Amodei. Scaling laws for neural language models. arXie preprinl arXiv:2001.0880 2020.
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Iilia Polosukhin, Jacob Devlin, Kenton Lee, et al. Natural questions: A benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:452 466, 2019.
Qian Liu, Xiaosen Zheng, Niklas Muennighoff, Guangtao Zeng, Longxu Dou, Tianyu Pang, Jing Jiang, and Min Lin. RegMix: Data mixture regression for language model pre-training. In The Thirteenth International Conference on Learning Representations, 2025.
Brandon McKinzie, Zhe Gan, Jean-Philippe Fauconnier, Sam Dodge, Bowen Zhang, Philipp Dufter, Dhruti Shah, Xianzhi Du, et al. Mml: Methods, analysis and insights from multimodal Ilm pre-training. In European Conference on Computer Vision, pp. 304-323, 2024.
Qwen Team. Qwen3 technical report. arXiv preprint or arxiv:200. 09388, 2025.
Shiori Sagawa, Pang Wei Koh, Tatsunori B. Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. arXiv preprint
urXiv:1911.08791, 2020.
Zhiqiang Shen, Tianhua Tao, Liqun Ma, Willie Neiswanger, Zhengzhong Liu, Hongyi Wang, Bowen Tan, Joel Hestness, Natalia Vassilieva, Daria Soboleva, and Eric Xing. SlimPajama-De Understanding data combinations for LLM training. arXiv preprint ar 2rXi::2309.100 108188 2023.
Mustafa Shukor, Louis Bethune, Dan Busbridge, David Grangier, Enrico Fini, Alaaeldin El-Nouby, and Pierre Ablin. Scaling laws for optimal data mixtures. In Advances in Newral Information Processing Systems, 2025.
Rens van de Schoot, Sarah Depaoli, Rochelle King, Bianca Kramer, Kaspar Märtens, Mahlet G Tadosso, Marina Vanmucci, Andrew Gelman, Daan Veen, Jolien Willemsen, and Christopher Yau. Bayesian statistics and modelling. Nalure Reviews Melhods Primers, 1(1):1, 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, pp. 1898-6008, 2017.
Yubo Wang, Xueguang Ma, Ge Zhang, Yuansheng Ni, Abhranil Chandra, Shiguang Guo, Weiming Ren, Aaran Arulraj, Xuan He, Ziyan Jiang, et al. Mmlu-pro: A more robust and challenging multi-task language understanding benchmark, 2024.
Sang Michael Xic, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu, Percy Liang, Quoc V Le, Tengyu Ma, and Adams Wei Yu. DoReMi: Optimizing data mixtures speeds up language model pretraining. In Advances in Neural Information Processing Systems, volume 36, pp. 69798 69818, 2023.
14

---

Published in Transactions on Machine Learning Research
Greg Yang and Edward J. Hu. Tensor programs IV: Feature learning in infinite-widtt neural networks. In Proceedings of the 38th International Conference on Machine Learning, pp. 11727-11737. PMI.R, 2021.
Jiasheng Ye, Peiju Liu, Tianxiang Sun, Jun Zhan, Yunhua Zhou, and Xipeng Qiu. Data mixing laws: Optimizing data mixtures by predicting language modeling performance. In International Conference Learning Representations, 2025.
Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo. CutMix: Regularization strategy to train strong classifiers with localizable features. In Proceedings of the IEEE/CVE International Conference on Compuder Vision, pp. 6023- 6032, 2019.
Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. In International Conference on Learning Representations, 2018.
Xuanhe Zhou, Junxuan He, Wei Zhou, Haodong Chen, Zirui Tang, Haoyu Zhao, Xin Tong, Guoliang Li, Youmin Chen, Jun Zhou, Zhaojun Sun, Binyuan Hui, Shuo Wang, Conghui He, Zhiyuan Liu, Jingren Zhou, and Fan Wu. A survey of LLM x DATA. arXiv preprint arXiv:2505. 18458, 2025.
A Scaling Law and Moments of the Dirichlet Distribution
The Dirichlet distribution, often referred to the multivariate Beta distribution, is the canonical choice for generating random probability vectors the simplex, making it especially suitable for modeling domain weights in large language model (LLM) data mixture. In this section, we review the definition and statistical moments of the Dirichlet, then provide the concrete parameter scaling law used for transferring proxy- optimized weights to larger models in DRAW. All key formulas are explicitly numbered for clarity.
A.1 Definition and Properties of the Dirichlet Distribution
Let 8 -Bi, 38)) denote the concentration parameter vector of the Dirichlet distribution. The random variable a (a au,), representing the domain weights, is drawn from the Dirichlet as follows:
f(a ()
Il
8-1 0
B(B)
where a; 2 0, Σ = 1. The normalization constant B(B), called the multivariate Beta function, is given by
B(B) I. I(3.) r E
where r(-) denotes the gamma function. When all B3 are large, the distribution concentrated near the uniform distribution; conversely, small 3, leads to mass near the simplex corners.
The first moments of the sampled weights as are
3
Ela.]
Σ/ -5=1 3j
Bi
EE 35- -B,)
(7)
Varlos] =
E 33)
2
- B, 1,+1)
15

---

Published in Transactions on Machine Learning Research
A.2 Expectation and Variance After Scaling
Suppose the proxy model and main model have widths ny and n2, respectively, and the number of data domains is k. The Dirichlet parameter for the i-th domain in the main model is set
n2
n1 + aa))
(8)
k
where of is the optimized proxy domain weight for the i-th domain. To facilitate later expressions, define the normalization sum:
c
U22
S=
Σ 32,4 vn2
(9)
VT
By substituting Eqs. equation 7 with the above scaling, the expectation and variance for the sampled weight in the main model are
Vug
E Eaa:
B2
VnI
+
oz
(10)
S
TU2 +
u V
Varlos] = B2 B2,1)
( + 11
/u22 Vni
ymi k
+ ox;)
112 +
//n
Vmi
+ az
Vna
(VTee +
1 VRI
+1)
vmi
+ +a, b
Vnn +
vii - a
(11)
(yri + 1)2
(VT2l//0 + 1) + n2
B Computing the Posterior Parameters of the Dirichlet Distribution
We now derive the posterior distribution for domain weights under a Dirichles-mulliomiiaa Bayesian model, which forms the core of our Bayesian updating in domain weight learning.
Prior. Suppose the prior over domain weights a is Dirichlet distribution with concentration parameter B1 181.1...,,11.... where each 81 yai The prior PDF is given by:
f(a Bi)
(12)
B(B,) 1=1
where B(31) is the normalization constant.
Likelihood, Let a* (ai....,aa) af) denote the proxy-optimized domain weights, interpreted pscudo- counts. The multinomial likelihood is:
N!
Lia" a) = I at! t=1
(13)
where N Σ 0f.
16

---

Published in Transactions on Machine Learning Research
Posterior.
By Bayes' theorem, the unnormalized posterior is proportional to the product of prior and
likelihood:
f(a a², B1) x f(a B1) L(a" a).
Substituting 12 and 13
1(a a3, x II
81.
1+0
02
1=1
That is,
f(a a², (r)
II
0g
i.e., the posterior is Dirichlet with updated parameters:
giew = (31,1 + aj,..., BL + +a)).
The normalization constant for the posterior (the multivariate Beta function) is
II T(3,, + TE (81. aa))
B(
Therefore, the normalized posterior density is
f(a a', (1)
a B -i-1
(14)
B gnew
C Theoretical Analysis and Proofs
C.1 Proof of Lemma
Lemma 1. Based on the Tensor Program and Maximal Update Parumetrization, the size of the model oulpul f after t iterations is @(1), even if the domain weights of the training data obey the Dir( (n) D
Proof of Lemma Before proving Lemma D we establish the statistical properties of the data mixture under Dirichlet scaling in Lemma 2
Lemma 2. Let each domain E [l....,k} k} characterized by data distribution P. with mean μ4 and covari- ance Σ.. At each training iteration, domain weights are sampled from Dir(g) where $ (Vn/k..... vn/k). Inpuls x are drawn from the mixture Pa(x) = E.= oz P.(r))). Then, as n - co, the impuls possess stable statistics:
E E[x]
M. 9(1),
(15)
k
Cov(a:)
1
Σ
[ + H
E|-[E[]] OOO (1/2) O(1).
(16)
Proof of Lemma The moments the Dirichlet distribution Dir(Bo1) with Bo = vn/k satisfy:
A h-1 B|a.] K Var(a.) 1²(vn++1) 1(22//2).
1
(17)
Using the Law of Total Expectation, E[x] Σ Bla,),. Using the Law of Total Covariance, Cov(x) [[(ov():(a)] + (o((a))) The first term converges to the average within-domain second moment. The
17

---

Published in Transactions on Machine Learning Research
second term, Covl_ axfs),, depends on Cov(as, ay), which scales ( 9/2). Thus, the variation in imput statistics due to Dirichlet sampling vanishes as -> 00, leaving Cov(2) = ((1).
Completing the Proof of Lemma !
Consider a standard 1-hidden-layer perceptron (MLP) under Maximal Update Parameurization (xP). The network output is given by:
f(z) _ 6(wx),
(18)
Vn"
where W E Raxd and E R" To ensure maxinal feature learning, parameters are initialized as Wji (N0 1/d) and 0j ~N(0 o.) (typically of 1 in this scaling).
From Lemma the input x has (1)) mean and covariance. Consequently, the pre-activations = Ww satisly:
h Σ W Ji;; E[hy]= 0, Var(h,) O(1)
Since is L-Lipschitz, the activations 2j = (b₂)) also possess bounded moments independent of n. The output f(x) is sum of n independent (initially) random variables, scaled by 1/ v/:
[[f(x)] = 00
1 Var[f(x)]
(a(())z))== _.. Var(t))B[23] O(1).
n
By the Central Limit Theorem, as n + 0c, f(x converges in distribution to a Gaussian (N0, where of 0(1.
During training via SGD, the weights are updated by AW and Av. Under jP logie, the learning rales are scaled (typically 7w x 1, x 1/n) such that the change in the pre-activations Ah and output Af remains 0(1) throughout training time f. Explicitly, since the input distribution moments (driven by a) are stable within ((n-1/2) 1/2) (Lemma the gradient moments
E
V Ea, 1 V parems
remain stable. Specifically, the fluctuation introduced by the randomized is of order O(n (nn-//3,, which is negligible compared to the (1) signal for large
Therefore, at any finite step I, the model output (x) retains O(1) statistics, ensuring stability.
C.2 Proof of Theorem 1
Theorem 1. Suppose F(0) = Ea[C(O, x)) is p-strongly convex and L-smooth. If the paramelers O are updated using stochastic gradients computed with domain weights at ~ Dir(B) using a slep size sequence {nc} satisfying Σn oo and Σni n² < 0o, then the optimization guarantees convergence in mean square to the unique minimizer 0* of F(0):
lim E ||0 0-112- =0, T
= 0, -
(19)
Proof. Assume F()) is p-strongly convex and L-smooth. At iteration 6, the update rule is:
041 = 0, nett, where = VoL(O,,
(20)
Since a is sampled from the Dirichlet distribution, the stochastic gradient unbiased, i.e., E[g = VF(O.). Furthermore, we assume the variance of the stochastic gradient bounded by o² due the bounded support of the Dirichlet weights and the local Lipschitz continuity of the loss:
E lg
VF(0.) < 2
(21)
18

---

Published in Transactions on Machine Learning Research
Now, consider the squared distance to the minimizer 0*. Expanding the iteration:
0++1 1² ||0, 1: 11
= 110, 0 112 2n (gt, 0 07) + n 9el1
(22)
Taking the expectation conditioned on Oy:
E 0+11
0" = |00 1-1 - 2m (VF()) -0)
8-) Bll 1..)
(23)
Using the variance decomposition |
| ((.) B[l V
0, 1 - (V - --)
+ (1 +
(24)
By jestrong convexity, (V.( 0. 0-) 1 0, 0 + -
IVF (0) However, a simpler bound using
just smoothness suffices. Since F is L-smooth, |VF L 0 Also using strong convexity, -VF((O)), 0²) S -u 0, 0"l 1/1 Combining these into Eq. equation
24
E| t+1 0" 0 S (1 () 0- + 11 L 10,
0-11 + 02
= (1 2pm + n 1 1-1 + +n
(25)
Taking the total expectation over all randomness history:
E 0| 11 - (1 - 2pmp +
0- + n,?0"
(26)
For sufficiently large f, since 1) - 0, we have I.2 < μr, so the contraction factor becomes (1 - pm))). Thus:
E 0,+1 0- 12 < (1 Bi 1- 0/
(27)
Applying the standard lemma for Robbins-Monro sequences (e.g., Chung's Lemma) with Σ Σ mh < co, we conclude that:
and
lim 0 0- 0.
(28)
Convergence in mean square implies convergence in probability (0r 3 0*). This completes the proof.
C.3 Proof of Theorem 2
Theorem 2. Let C(O,00 a)) be the loss function of a model with width n2 parameterized by O, under domain misture weights Assume the loss is M1 -i-Lssschälz regulor wilh respect to a. Define the expecled risk as:
F(0) DEI
where B == a (implying concentration parameter n1). Let 0* denote the population minimizer and Ô the parameter oblained after oplimization.
The generalization error, governed by the interplay between domain variance (scale ni) and model width (scale n2), is bounded by:
k
B 3i) B +m1)
F(O) - F(0") S Mi
Σ
where B = + is the sum of Dirichlet parameters.
the regime of large data concentration (n) >> 1) and wide models (ri > 1), the scaling law simplifies to:
((0)-F(0') =0 Mi
\
19

---

Published in Transactions on Machine Learning Research
Proof. Let (ir(3) with B= = 1. The term F(Ø) F(0)) represents the excess risk induced by the stochasticity of the domain weights a relative to the ideal population distribution.
For a neural network of width n2, the output (and consequently the loss) can be viewed as an aggregation of feature mappings. According to the Law Large Numbers applied to the width the network, the variance of the loss fluctuation due to input distribution shifts scales inversely with the width n2- Specifically, if Va represents the variance intrinsie the Dirichlet sampling, the effective variance observed by the wide model is reduced:
1 Var(a:).
VaLmodel x
722
We proceed by bounding the deviation using the Lipschitz property. Since L is M1-Lipschitz w.r.t az:
[2(6, ox) L(0",Ea)) & M1 la Ea|22.
(29)
However, this deviation is damped by the model width. The expected squared error (variance contribution) is:
E - Σ El Ea,)
n2
More precisely, substituting the variance of the Dirichlet distribution Var(o;) 8,(B-8)) Bc). The term in the
specifically derived bound combines the Dirichlet variance denominator (B + 1) with the model width factor. Ideally, the variance scales as 1 (from data) and 1 (from model width). The coupled term in the W:
denominator (ng + m1) reflects this interaction:
2 n
(n2B + n1) ss 12n11
Total denominator ss n2n2..
The numerator to oi (as shown in Lemma 2). Thus, the expression inside the square rool behaves
nf
ni 12
10
Formalizing the bound:
B(BB 3()
F(O) F(O*) < M1
nn + ()2(n((1 + + + n1)
Taking the asymptotic limit ny > 1,72 > 1:
O(fi)
(MM
Mi
RHS Miy
0
(
nf O(112))
This confirms that the generalization error decreases with the geometric mean of the data concentration parameter and the model width n2.
D Pseudocode of DRAW
Algorithm I outlines the complete DRAW framework, illustrating how domain weights are transferred from proxy model a main model via Bayesian updating. The process of obtaining the optimal proxy weights ax (as specified in Algorithm 1 relies Distributionally Robust Optimization (DRO).
20

---

Published in Transactions on Machine Learning Research (Apr/2020)
Algorithm DRAW: Domain Weight Randomization with Bayesian Updating
Require: Dataset D (D1..... D with & domains; Proxy model width n1; Main model width ng: Total training steps T.
Ensure: Trained main model parameters 0*
// Phose 1: Prozy Model Optimization
Initialize prior concentration parameter: $1 <- vni
Eq. (1)
Sample initial domain weights: Oinit Dir(31)
Random Initialization
3: Train proxy model with width (imitialized with ainit) to obtain optimal weights a":
4: a" f Solve DRO objective Eq. (2)
// Phose 2: Boyesion Updating And Scaling
Compute scaling factor from scaling laws: 4
Update concentration parameters for the main model:
7: 82 + 8. 8.50""" a''a"'"-- + LRz.
D Eq. (5)
Phahase 3: Main Model Training
Initialize main model parameters with width n2
for t 1 to T do
10:
Sample dynamic mixture weights: a, Dir(32)
11: Sample batch B, from D according to distribution a
12:
Compute loss L(6, B,)
13:
Update model parameters: 0 + 0 - VVC((0 B.)
14: end for
15: return
E
Experimental Details
E.1 Data Preparation and Preprocessing
For all datasets, we tokenized the text using the SentencePiece tokenizer and separated the data into shards by domain. Each example was chunked to a maximum sequence length of 1024 tokens, and all samples were indexed by domain for hierarchical sampling. During training, we first sample domain weights from Dirichlet distribution, then sample domain according to these stochastic weights, and finally sample example from the selected domain. To improve training efficiency and reduce padding, also implement sequence packing, occasionally mixing samples from multiple domains into the same packed sequence.
E.2 Model Architecture and Hyperparameters
Table 4: Architecture hyperparameters for various model scales used in the paper.
Layers
Attention heads
Attention head dim
Model dim
Hidden dim
70M
4
64
256
1024
150M
8
512
2048
280M
12
12
64
768
3072
510M
12
16
64
1024
8192
760M
12
20
64
1280
8192
IB
16
32
64
2048
8192
8B
32
32
128
4096
24576
Table 4 provides an overview of the key architectural hyperparameters used across different model scales in our experiments. All models are based on the standard Transformer decoder architecture, with configurations that include varying numbers of layers, self-attention heads, attention head dimensions, model dimensions, and hidden layer sizes. This design allows for systematic scaling from small proxy models (70M parameters) to
21

---

Published in Transactions on Machine Learning Research (Apr/202)
large main models (up to 8B parameters), which facilitates the study of scaling laws and the transferability of domain mixing strategies. Larger models employ wider and deeper structures, resulting in increased expressive power and capacity for more complex language modeling tasks. All models utilize the same core architectural components to ensure comparability and isolate the impact of scale.
E.3 Training Setup
We standardizod the training budget across all modcls (proxy and main) to 50k steps, using global batch size of 256 and sequence length of 1024 (approx. 13.1 billion lokens). Training was implemented using the AdamW optimizer (31 0.9, 32 = 0.99) and schedule with linear warmup followed by exponential docay (1 x 10-3 to 1 x 10-1). To facilitate hicrarchical domain sampling, domain woights were drawn from a Dirichlet distribution and updated every 10 steps for main runs. Evaluations were performed every 100 or 200 steps. All experiments utilized four NVIDIA A100-PCIe-40GB GPUs supporting CUDA 12.4.
E.4 Reference Model
The reference model adopts the same architecture and training procedure as the proxy model, but differs in that it uses fixed domain weights derived from the established baseline configuration. Specifically, the baseline domain weights for The Pile were computed as follows: after chunking each domain's data into 1024- token examples, we counted the number of such examples per domain and multiplied by the number epochs assigned to that domain in the original work. The resulting counts were then normalized to obtain the final domain weight proportions, which remain fixed throughout training. This enables a controlled comparison between learned randomized weighting schemes and traditional static mixing commonly practiced in prior literature.
Table Validation loss across domains for Baseline, DoReMi, RegMix, and DRAW methods.
Domain
DoRcMi
Bascline
Random
RegMix
DRAW
ArXiv
2.7343
2.3825
2.5041
3.2113
2.4381
DM Mathematics
1.7261
1.6215
2.5917
2.3177
2.8594
Enron Emails
3.4327
3.2207
3.0088
3.1261
3.5042
EuroParl
4.2285
3.6318
1.6137
5.6617
1.5440
FreeLaw
3.0444
2.8175
2.6689
3.3294
2.5293
Github
1.6345
1.7982
3.1087
2.7782
3.0553
Gutenberg (PG-19)
4.0417
3.7704
3.0145
3.8131
2.9464
HackerNews
3.8879
3.6171
1.9959
3.4490
1.9152
NIH ExPorter
3.3692
3.3904
3.9611
3.6346
3.8892
PhilPapers
4.4872
4.2033
3.4648
4.8690
3.3987
Pile-CC
3.9710
3.8146
3.2215
3.5740
3.1611
PubMed Abstracts
2,9910
3.1552
2.9816
3.4259
3.0042
PubMed Central
2,8132
2.6569
3.2368
3.1551
2.7489
StackExchange
2.7135
2.6690
3.8000
3,4060
3.7361
USPTO Backgrounds
3.1326
2.9615
4.0785
3.2677
3.9905
Ubuntu IRC
3.6097
2.9608
3.2862
2.6253
3.2227
Wikipedia (en)
3.1959
3.5380
2,8384
3.3974
2.7691
Average Validation Loss
3.2361
3.0711
3.0428
3.4730
2.9689
Worst Validation Loss
4.4872
4.2033
4.0785
5.6617
3.9905
22

---

Published in Transactions on Machine Learning Research
F
Additional Experiment Results
To further validate the advantages of the proposed DRAW method in multi-domain mixture, adaptive weighting, and model generalization, we present and analyze supplementary experimental results from mul- tiple perspectives, inchuding weight evolution, downstream task performance, and proxy model robustness.
Firstly, to provide an intuitive understanding of how domain weights evolve under DRAW, Figure visualizes the distribution of Dirichlet-sampled weights before and after training. At initialization, the samples exhibit an approximately isotropic distribution across all data domains, indicating no prior domain preference. As training progresses, the distribution becomes more concentrated around certain regions, reflecting that the model has Iearned emphasize domains that contribute more to downstream objectives. This adaptive behavior of DRAW underpins its improved downstream performance.
To further assess the impact of this adaptive weighting, compared several weighting strategies (Baseline, DoReMi, RegMix, and DRAW) validation losses across 17 data domains, as summarized in Table DRAW consistently achieves lower validation losses in the majority of domains, with particularly pronounced advantages in the worst-case loss metric. This demonstrates that adaptive domain mixing not only enhances average task performance but also improves robustness to domain imbalance and tail distributions.
Table 6: Accuracy for the main model (1B) major downstream benchmarks during DRAW pretraining.
Step
MMLU
MMLU-PRO
NQ
GSM8K
BoolQ
ARC-C
Average
2,000
65.00
41.83
0.46
0.00
6.82
68.58
30.45
4,000
40.41
24.38
1.39
0.00
0.98
42.04
18.20
6,000
38.74
23.19
2.31
0.13
3.24
34.29
16.99
8,000
29.94
19.03
4.63
1.07
4.34
27.65
14.44
10,000
25.12
13.13
3.24
1.20
4.68
25.66
12.17
12,000
25.85
11.75
6.02
0.94
3.85
25.44
12.31
14,000
25.70
11.27
4.17
1.60
9.27
25.00
12.83
16,000
25.71
13.85
6.02
0.67
10.00
24.78
13.50
18,000
26.78
15.25
3.24
1.47
7.83
27.21
13.63
20,000
26.59
13.04
5.09
0.94
5.84
26.11
12.93
22,000
26.91
12.95
6.02
1.34
7.80
27.88
13.82
24,000
28.38
15.87
6.48
1.20
6.02
29.20
14.53
26,000
27.21
13.83
4.63
0.94
9.85
27.21
13.94
28,000
31.21
20.55
4.17
1.60
7.95
28.76
15.71
30,000
28.64
16.87
6.48
1.34
6.88
25.66
14.31
32,000
29.94
20.93
3.70
0.94
6.39
29.87
15.29
34,000
30.23
16.20
6.02
1.74
3.46
26.11
13.96
36,000
26.99
15.04
6.02
1.74
2.63
28.32
13.46
38,000
28.33
15.97
3.70
1.60
4.25
28.98
13.81
40,000
30.46
17.33
5.09
1.74
4.25
30.53
14.90
42,000
30.46
17.89
4.17
1.74
3.15
28.76
14.36
44,000
28.01
17.60
6.48
1.47
3.70
24.78
13.67
46,000
30.87
22.45
5.56
1.07
3.79
25.00
14.79
48,000
26.86
16.10
5.09
0.80
3.43
23.67
12.66
50,000
31.58
16.61
6.02
1.60
5.11
25.00
14.32
Beyond final performance, the dynamics model improvement during training are also crucial. Table 7 details the accuracy of the 1B main model on major downstream tasks during DRAW pretraining, with the learning curves further illustrated in Figure 6 We observe that DRAW consistently accelerates accuracy improvements across tasks, often reaching higher accuracies with fewer training steps compared to
23

---

Published in Transactions on Machine Learning Research
the baseline. Detailed trend analysis also reveals that DRAW stabilizes carly-stage training and mitigates performance volatility, highlighting its efficacy in guiding more efficient. and stable learning.
MMLU
MMLU-PRO
45
Daseline
Daseline
60
35-
50-
Accuracy
20-
30
10-
Training Steps
Tainiin Steps
GSMBK
2.5
MW
Daseline
Develine
m
Avlly m
Accuracy
Accuracy
1.5
2.0
0.0
10k
20k
30k
42k
50k
10k
20k
40k
S0k
Training Steps
Training Steps
BoolQ
ARC-C
14
Our
Our Daseline
Daselne
60
10
Accuracy
Accuracy
50
30-
0k
10k
20k
30k
40k
50k
0k
10k
20k
30k
40k
50k
Training Steps
Training Steps
Figure 6: Accuracy(%) curves of six benchmarks (MMLU, MMLU-PRO, NQ, GSMSK, BoolQ, ARC-C) comparing DRAW and Bascline.
The stability of learned domain weights under stochastic initialization is analyzed by comparing the optimized weights across different random seeds for 17 domains, as shown in Table g While some local variances are observed, the overall weighting trends remain consistent, evidencing the robustness of the DRAW framework. This outcome suggests that even for rare or long-tail domaīns, the adaptive mechanism reliably allocates meaningful weights, which is crucial for real-world deployment.
systematically explores the effect of proxy model capacity the learned domain weights. We compare results using proxy models of various sizes (70M, 150M, 510M, 280M). It's found that although a larger proxy can provide finer adjustments, the overall weighting trend remains essentially invariant across model scales. This indicates that. DRAW's core mechanism does not depend on excessively large proxy models, supporting practical scalability and engineering feasibility.
Lastly, Table 7
In summary, this appendix systematically demonstrates the advantages of DRAW in terms of weight evo- lution, domain adaptability, and the robustness of the learned ceeturre-strongh supporting the main con-
24

---

Published in Transactions on Machine Learning Research
Table Domain weights learned by different proxy model sizes (70M, 150M, 510M, 280M) for DRAW.
Domain
70M
150M
510M
280M
ArXiv
0.0010
0.9709
0.7306
0.4735
DM Mathematics
0.0086
0.0027
0.0005
0.0004
Enron Emails
0.2933
0.0002
0.0002
0.0002
EuroParl
0.0004
0.0001
0.0002
0.0002
FreeLaw
0.0008
0.0003
0,0005
0.0006
Github
0.0225
0.0219
0.0719
0.0766
Gutenberg (PG-19)
0.0002
0.0002
0.0003
0.0003
HackerNews
0.0001
0.0002
0.0003
0.0003
NIH ExPorter
0.0002
0.0002
0.0003
0.0004
PhilPapers
0.0002
0.0002
0.0002
0.0003
Pile-CC
0.0003
0.0003
0.0007
0.0010
PubMed Abstracts
0.0007
0.0008
0.1904
0.4398
PubMed Central
0.0244
0.0004
0.0006
0.0007
StackExchange
0.2395
0.0004
0.0006
0.0007
USPTO Backgrounds
0.0003
0.0003
0.0004
0.0005
Ubuntu IRC
0.0002
0.0002
0.0002
0.0002
Wikipedia (en)
0.4075
0.0006
0.0021
0.0043
clusion that DRAW enables more effective and adaptable language model pre-training on large-scale, multi- domain datasets.
Table 8: Domain weights across 17 data domains under varying random seeds for DRAW.
Domain
Seed 1
Seed 115
Seed 123
Seed 17
Seed 42
ArXiv
0.0010150
0.0010100
0.0009531
0.0009537
0.00096446
DM Mathematics
0.0085540
0.0085540
0.00855400
0.0085550
0.0085530
Enron Emails
0.2933000
0.3295000
0.3036000
0.2641000
0.2770000
EuroParl
0.0003762
0.0003897
0.0003868
0.0003912
0.0003880
FreeLaw
0.0008116
0.0008223
0.0008085
0.0008287
0.0008117 17
Github
0.0225100
0.0215800
0.0219400
0.0216300
0.0213500
Gutenberg (PG-19)
0.0001532
0.0001557
0.0001554
0.0001563
0.0001562
HackerNews
0.0001374
0.0001394
0.0001391
0.0001400
0.0001396
NTH ExPorter
0.0002121
0.00021 23
0.0002127
0.0002132
0.0002123
PhilPapers
0.0001590
0.0001594
0.0001592
0.0001594
0.0001588
Pilc-CC
0.0002578
0.0002607
0.0002607
0.0002605
0.0002619
PubMed Abstracts
0.0006736
0.0006458
0.0006520
0.0006363
0.0006345
PubMed Central
0.0244400
0.0247600
0.0250200
0.0237100
0.0241000
Stacknxchange
0.2395000
0.2262000
0.2261000
0.2410000
0.2477000
USPTO Backgrounds
0.0002996
0.0002919
0.0002947
0.0002926
0.0002920
Ubuntu IRC
0.0001835
0.0001799
0.0001801
0.0001822
0.0001787
Wikipedia (cn)
0.4075000
0.3852000
0.4106000
0.4368000
0.4171000
25


Reviews:
[{"review_id": "odQVEjf2lw", "reviewer": "Reviewer_pg11", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper revisits domain weighting for LLM pretraining by modeling domain weights as random variables and introducing a Bayesian framework (DRAW) to transfer mixture policies from a proxy model to a larger model. The proxy model comes up with an initial policy using the DRO frameworks which have been extensively explored in the community. The method uses the DRO based proxy model distribution to initialize a Dirichlet distribution which in turn guides the domain weights for the main model. Empirically, the paper shows that stochastic domain weighting—especially when informed by a proxy—can outperform fixed-weight baselines on language modeling and downstream tasks.", "claims_and_evidence_yn": "Yes", "claims_explanation": "- The claims are well backed by a detailed theoretical analysis in pretty nice setups, which is a rare thing in LLM pretraining works and the strongest contribution of this work\n- Empirical evidence is also decently good.\n- Some very intriguing findings that even random sampling surpasses fixed weight training \n- The dirichlet based random sampling of domain weights results are well backed by theoretical insights, which explain both the above observations that why randomness can help find a better local minima", "audience_yn": "Yes", "audience_explanation": "Domain weighing during LLM pretraining is one of the biggest problems in both academia and industry. Overall, this paper would be of interest to lot of people.", "requested_changes": "Weakness\n\n- The method is largely a reformulation of DRO-style domain weighting into a Bayesian framework. The proxy model still optimizes domain weights via a min–max objective over α, minimizing excess loss under worst-case domain weights. This implicitly assumes adversarial domain mixtures, which does not reflect practical scenarios where only a small subset of domain weightings are realistic. As a result, many limitations of DRO-based LLM reweighting carry over unchanged.\n\n- All the limitations of DRO based LLM weighing would continue to trickle in because the dirichlet distribution continues to be guided by the same. These include the overly adversarial optimization, harder convergence of proxy model training, \n \n- Step 3: making weight distribution for bigger model more sharper. Why would a bigger model benefit from a more sharper distribution? Is there any intuition behind this?\n\n- The paper repeatedly samples domain weights from the same Dirichlet distribution throughout training. THis is not a “dynamic refinement” which has been more of focus in recent works like that of Jiang et al (https://arxiv.org/abs/2410.11820). Do authors have any thoughts on this?\n\n- The biggest challenge in domain weight sampling continues to be how to model cross domain interaction. Although I acknowledge that the authors mention this in limitations of the work as well.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "OBE2HiDDeX", "reviewer": "Reviewer_tPWY", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces DRAW (Domain Weight Randomization with Bayesian Updating) as a data mixture optimization algorithm for Large Language Model (LLM) pre-training. Instead of a common assumption that the optimal data mixture should be transferrable between different model scales, DRAW searches for a scale-dependent optimal data mixture by modelling the sensitivity upon domain weights as the concentration factor in a Dirichlet-distributed prior.\nThe main optimization problem is identical to DoReMi [1], where the Group DRO is applied to prioritize the worst-case excess loss among training domains. Various $\\beta$ parameters in a Dirichlet distribution $P(\\alpha; \\alpha^{\\star}, \\beta)$ are applied to different model scales, where $\\alpha^{\\star}$ is from a small-scale proxy using DoReMi algorithm.\n\n## Strength:\n\n(1) Theoretical soundness of incorporating scale-variant factor with Dirichlet priors: The intuition of *the greater sensitivity of larger models to domain imbalances* is reasonable and can be well reflected by the strength of $\\beta$ parameter in Dirichlet distribution. \n\n(2) The paper is well-written with the main theorems and results clearly presented with proper notations.\n\n(3) The empirical results presented on downstream evaluation accuracies (table 1) and validation perplexities (Figure 5) are able to support the main claim that DRAW does outperform the baselines. \n\n## Weakness and Questions:\n\n(1) How related is Theorem 1 to the main domain reweighting algorithm? It appears to me that if the $\\alpha$ distribution has a computable constant expectation, the loss will converge under the given assumptions.\n\n(2) Questioning the logic and purpose of Theorem 2: What is supposed to be the main conclusion from theorem 2 to the data mixture optimization algorithm? According to Theorem2, the generalization gap decreases both according to data concentration factor n1 and model width n2, which indicates that the model's performance would be more stable with a wider model and be less sensitive to the data distribution. However, it conflicts with the main hypothesis that \"large language models tend to exhibit greater sensitivity to the composition of training data\". How shall we approach this theoretical-empirical gap?\n\n(3) Based on the formulation of Equation (1), does it mean a larger proxy model would lead to a more uniform domain weight distribution? Does it also conflict with the main hypothesis of domain weight sensitivity?\n\n(4) The method does not seem to be generalizable to out-of-domain generalization case (Similar as DoGE [2]), where the target is not included in the training domains.\n\n(5) A contradict of Remark 1 and empirical results in Table2: how could the domain weights from DRAW being extremely different from the $\\alpha^{\\star}$ (domain weights from DoReMi), by only incorporating a scale-dependent factor? Also, it seems the impact of $\\alpha^{\\star}$ does diminish when applying DRAW.\n\n[1] DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining.\n[2] DoGE: Domain Reweighting with Generalization Estimation.", "claims_and_evidence_yn": "Yes", "claims_explanation": "The empirical results presented on downstream evaluation accuracies (table 1) and validation perplexities (Figure 5) are able to support the main claim that DRAW does outperform the baselines.", "audience_yn": "Yes", "audience_explanation": "The targeted problem, data mixture optimization for langauge model pretrianing is a critical and challenging problem in today's foundation model training. Especially, the scalability of data attribution mechanisms is one of the most important chanllenges when applying those gradient-based methods on large-scale model and datasets. \n\nDespite several confusions and questions, I still believe this paper conduct an impressive and interesting trial on the scaling invariant version of data mixture optimization but without fitting the traditional scaling law, which often requires hundreds of reruns.  Therefore, i think this paper can be interesting to a lot of machine learning researchers in both data-centric AI and foundation model training fields, after it improves on certain aspects.", "requested_changes": "(1) **Clarify the algorithm and implementation steps more**: now how the reweighting iteration is conducted is still not very clear. Include a pseudo-code can help.\n\n(2) **Resolve the theoretical-empirical contradict, and the conflict between hypothesis and the theoretical/statistical intuitions**: resolve questions 1,2,3,5 in the Weakness and Quesitons part above.", "broader_impact_concerns": "None. This paper uses established, copyright-free subsets of The Pile for the experiments, which does not introduce ethical issues or risks.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "wshPmZ0elP", "reviewer": "Reviewer_9DuQ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes DRAW, a framework for adaptive data mixture selection in large language model pre-training. \nThe key idea is to model domain weights as Dirichlet-distributed random variables, learn an informative posterior using a smaller proxy model via Bayesian updating, and transfer this distribution to a larger model using width-based scaling. \nDuring main-model training, domain weights are dynamically sampled at each step. \nThe paper provides theoretical analysis suggesting convergence and favorable scaling behavior, and presents empirical results on language modeling and downstream benchmarks.", "claims_and_evidence_yn": "No", "claims_explanation": "The total training token budget is not specified, the notion of an “epoch” under dynamic reweighting is ambiguous, and many downstream benchmark results are at or near chance level. \nAs a result, it is difficult to assess whether the observed gains reflect meaningful improvements or artifacts of undertraining and evaluation noise. \nThe evidence supports relative comparisons within the experimental setup, but does not convincingly substantiate broader claims about robustness and downstream transfer. \n\nAlso, the theoretical guarantees rely on assumptions that are far removed from real LLM training: strong convexity and smoothness of the objective; Lipschitz continuity w.r.t. domain weights; convergence analysis resembling classical SGD on convex objectives\nWhile this is common in theory sections, here the theory is used to justify specific scaling laws and practical design choices (e.g., proxy sufficiency, variance decay), yet no empirical validation directly tests the theory’s quantitative predictions. The analysis does not account for non-convexity, representation collapse, or memorization effects typical in LLMs. The theory is largely illustrative, but the paper sometimes presents it as explanatory, which risks overclaiming.", "audience_yn": "Yes", "audience_explanation": "The conceptual framing of domain weights as distributions rather than fixed parameters is novel and aligns well with current interest in principled data mixing and robustness. \nEven if the empirical evaluation is currently limited, the methodological ideas and theoretical perspective may inspire follow-up work or alternative implementations, making the paper relevant to part of the TMLR community.", "requested_changes": "1. Clearly report the total number of training tokens and optimization steps used for both proxy and main models.\n2. Explicitly define how training progress and “epochs” are measured under dynamic domain reweighting.\n3. Re-evaluate or better justify downstream benchmarks, especially given that many results are at or near random guessing.\n4. Add ablations to isolate the effects of dynamic sampling vs. static weighting, and of Dirichlet scaling choices.\n5. Improve reproducibility by clarifying evaluation protocols, hyperparameters, and baseline tuning.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)