Domain2Vec: Vectorizing Datasets to Find the Optimal Data Mixture without Training
Authors: Mozhi Zhang
Howe Tissue
Lu Wang
Xipeng Qiu
Venue: ICML 2025 poster (2025)
Keywords: Language models
Pretraining
Data mixutre
Primary area: deep_learning->large_language_models

Abstract: We introduce *Domain2Vec*, a novel approach that decomposes any dataset into a linear combination of several *meta-domains*, a new concept designed to capture the key underlying features of datasets.
*Domain2Vec* maintains a vocabulary of meta-domains and uses a classifier to decompose any given dataset into a domain vector that corresponds to a distribution over this vocabulary.
These domain vectors enable the identification of optimal data mixture for language model (LM) pretraining in a training-free manner under the ***D**istribution **A**lignment **A**ssumption* (DA$^{2}$), which suggests that when the data distribution of the training set and the validation set is more aligned, a lower validation loss is achieved.
Moreover, *Domain2Vec* can be seamlessly integrated into previous works to model the relationship between domain vectors and LM performance, greatly enhancing the efficiency and scalability of previous methods.
Extensive experiments demonstrate that *Domain2Vec* helps find the data mixture that enhances downstream task performance with minimal computational overhead.
Specifically, *Domain2Vec* achieves the same validation loss on Pile-CC using only $51.5$\% of the compute required when training on the original mixture of The Pile Dataset.
Under equivalent compute budget, *Domain2Vec* improves downstream performance by an average of $2.83$\%.

Domain2Vess
Vectorizing Datasets to Find the Optimal Data Mixture without Training
Mozhi Zhang'
Howe Tissue
X
Lu Wang?
Xipeng Qiu
Abstract
Qwen Team, 2024; DeepSeek-AI, 2024). The training datasets for LLMs are typically composed of multiple do- mains, each derived from different sources. Recent research has shown that the mixture proportions of these domains (named as the data mixture) can significantly influence the effectiveness of LMs (Hoffmann al., 2022b; Xie et al., 2023b), with data from domain potentially affecting the performance in others (Guo et al., 2022). Typically, the data mixtures used for training LLMs are often determined heuristically or based on downstream performance metrics. However, these methods are not scalable and always result in a suboptimal data mixture. Thus, identifying the optimal data mixture in scalable and efficient manner remains a critical and challenging research question.
We introduce DOMAIN2VEC; a novel approach that decomposes any dataset into a linear combi- nation of several meta veea-ddmainns, a new concept designed to capture the key underlying features of datasets. DOMAIN2 VFC EEC maintains a vocabulary of meta-domains and uses classifier to decom- pose any given dataset into domain vector that corresponds to a distribution over this vocabulary. These domain vectors enable the identification of optimal data mixture for language model (LM) pretraining in a training-free manner under the Distribution Alignment Assumption (DA?), which suggests that when the data distribution of the training set and the validation set is more aligned, lower validation loss is achieved. Moreover, DOMAIN2VEC can be seamlessly integrated into previous works to model the relationship between domain vectors and LM performance, greatly en- hancing the efficiency and scalability of previ- ous methods. Extensive experiments demonstrate that DOMAIN2 VEC helps find the data mixture that enhances downstream task performance with minimal computational overhead. Specifically, DOMAIN2 achieves the same validation loss on Pile-CC using only 51.5% the compute re- quired when training on the original mixture of The Pile Dataset. Under equivalent compute bud- get, DOMAIN2VER improves downstream perfor- mance by an average of 2. 2.83%.
Recently, researchers have proposed various methods to pre- dict the optimal data mixture. The first line of prior works implicitly adjusts the data mixture by selecting high-quality data from different domains or datasets (Lin et al., 2024; Ankner al., 2024; Thakkar et al., 2023). The second line of work focuses on modeling the relationship between the data mixture and the performance of LLMs, and explic- itly adjusts the data mixture across different datasets (Rae et al., 2022; Xie et al., 2023a; Sagawa* et al., 2020; Fan et al., 2023; Ye et al., 2024; Ge et al., 2024; Gu et al., 2024a; Que et al., 2024). While prior work has shown promising results, there are some key issues: 1) Computa- tional Efficiency: For example, although the proxy model in DoReMi (Xie al., 2023a) has only 280M parameters, its estimated FLOPs are high to 3.7 x 10 for calculating only 22 datasets. Moreover, The computational complexity of these methods will grow non- non-linearly as the number of datasets increases. 2) Lack of Scalability: After establish- ing the functional relationship between data mixtures and model performance (Ye et al., 2024; Liu et al., 2024), if the dataset composition changes (e.g.. by adding new datasets or filtering low-quality data, etc), previously fitted functions cannot be directly applied. This requires resampling new data mixtures, retraining proxy models, and refitting the functions, severely limiting the scalability of these methods.
1. Introduction
Through training on large-scale text corpora, Large Lan- guage Models (LLMs) have demonstrated remarkable gen- eralization capabilities (Touvron et al., 2023; OpenAI, 2024;
"ccoool of Computer Science, Fudan University, Shanghai, China "Ritzz-A AI. Correspondence to: Howe Tissue (project lead) <h- ch-sun2088 28t t wwt.tttinnnnaa.eee org. en>.
To address these issues, we introduce DOMAIN2VES a novel framework designed to vectorize datasets. This en- ables us to perform all operations for computing optimal
Proceedings of the ndnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Domain2Vec: Vectorizing Datasets to Find the Optimal Data Mixture without Training
mixing ratios in domain vector space, which has broad appli- cability when datasets change. Specifically, DOMAIN2 VEC maintains vocabulary of meta-domains, and we hypothe- size that any dataset can be approximated as a linear combi- nation of several meta-domains with a specific distribution. This distribution could serve as the vector representation (domain vector) of a given dataset.
Key Assumption. DOMAIN2VEC maintains a vocabulary, set of meta-domains. Assume we have n. meta-domains Dj (1 و S n),, where Dj is represented ej. a one-hot vector where only the j-th element is 1. We hypothesize that, for any given dataset DD it could be represented as a domain vector v, by linear combination of these meta- domains. Specifically,
To efficiently identify the meta-domain composition of any given dataset, we propose to use meta-domain classifier to generate the corresponding domain vector. Building upon DOMAIN2VEC, we introduce the Distribution Alignment Assumption (DA?) to find optimal data mixtures for LM pretraining. DA² states that lower validation loss can be achieved when the domain vector of the training dataset bet- ter aligns with the domain vector of the validation dataset. Based on 2 we easily find the optimal data mixture without training.
Dj Cjz
/=1
where each element Vj of v represents the projection (weight) of the dataset D on Dj. Thus, v [m) V2, V2,,.. ..... Un] can be representation (distribution) of the dataset D over the meta-domains. However, an ideal approach for con- structing these meta-domains remains to be established. Next, we will introduce how we construct meta-domains from large-scale unlabeled text corpora.
Moreover, DOMAIN2VEC can be seamlessly integrated into prior works like RegMix (Liu et al., 2024). Unlike previous methods that model the relationship between data mixtures and language model performance (Liu et al., 2024; Ye et al., 2024), we model the relationship between domain vectors provided by DOMAIN2 and model performance, further enhancing efficiency and scalability of previous works.
Constructing the Vocabulary of DOMAIN2VEC. With the above key assumption, define meta-domains collection of actual datasets (or set of domains) that serve as a basis in the domain vector space, allowing for linear combinations of these concrete datasets to represent any unknown domain in this space. These constructed meta- domains, which could represent datasets from any source, should satisfy the following three properties, similar the properties of basis in linear algebra:
In summary, we highlight our contributions as follows:
1. We introduce DOMAIN2 to vectorize datasets and propose viewing datasets as combinations of meta- domains. We present an efficient pipeline for vectoriz- ing datasets using a meta-domain classifier.
Spanning Set. The domains that compose meta- domains should be diverse and comprehensive as possible.
2. We propose the Distribution Alignment Assumption (DA?), training- free method for identifying the op- timal data mixture. We further demonstrate how Do- MAIN2 can be seamlessly integrated into prior work to improve efficiency and scalability.
2. Linear Independence. There should be distinct differ- ences between these constructed meta-domains.
Computational Efficiency (Optional). The method for constructing meta-domains should be computation- ally efficient.
3.
We validate the effectiveness of DOMAIN2VEC+DA and +RegMix in text generation and downstream tasks. Experimental results show that our method can accu- rately predict the performance of various data mix- tures without training proxy models. Moreover, we can identify data mixtures that achieve downstream per- formance comparable to DoReMi (Xie et al., 2023a), using only 0.26% of its computational cost.
For diverse and comprehensive meta-domains, we collect data from more than 100 coarse sources across English, Chinese', and Code. After deduplication, we obtain around 5.2 TB text data including more than 1 billion documents. The large corpora have a similar source composition as the standard large-scale L.MM pretraining, including common crawl (CC), Wikipedia, social media platform, arXiv, code, news, books, etc. One could assume that the corpora already include as diverse and comprehensive contents as possible. corresponding to the requirement "spanning set"
2. Domain2 Vec
In this section, we introduce DOMAIN2VEC, an algorithm that decomposes dataset into linear combination of var- ious meta-domains and allows us to represent underlying features of datasets through a normalized vector. We also outline a pipeline for constructing the vocabulary of Do- MAIN2VER VFC and training a meta-domain classifier.
'In this paper, we primarily aim at languages English and Chinese.
DDue to deduplication pre-processing and the native difference among the corpora, the requirement "linear independence" is also naturally satisfied.
2

---

Vee:in2Vec: Vectorizing Datasets to Find the Optimal Data Mixture without Training
Ple-00
Pub
thu
Reclaw
Stocked
sackgrounds
PubMie
Gutenborg IP6 19)
Wicpedls len)
one Wathemal
buroher
Hacket fown
PhiPspore
NH ExPorter
Figure 1. The domain vector of each sub-dataset of The Pile (Gao et al., 2021), where each row corresponds to sub-dataset and each column corresponds to meta-domain. The higher the proportion of data belonging to a particular meta-domain, the closer the color of the corresponding cell to blue. We display distribution some English meta veta-domains for clarity. The full picture is shown in Figure 7.
Inertia vs Number of Meta Domaanss
hensive representation captures the overall domain charac- teristics of the entire dataset. Formally, assume that there is meta-domain classifier, for any given document doe DD
Chinese Dataset
1.50
English Dataset
p Pi P2. .... p.. Classificr(doc).
(2)
where p: represents the probability that doe belongs to the i- th meta-domain such that ||||| 1. For the unseen dataset D, we could sample N documents then take the average of domain vector of these samples. Formally, the domain vector u of dataset D is,
240
N Σ
V =
P
Veta-Domains
Figure 2. The number of meta-domains vs. Inertia.
Then, we could use the vector to approximately represent the feature of any unseen dataset 'D. Meanwhile, during the pretraining phase of LLMs, we typically have training datasets from many sources Dtrain = {D1, D...... Dm We can convert each of these datasets into domain vec- tors following Equation and 3. Therefore, Dtrain can be approximately represented Virain [V1, U2, Um. where Virain E paxm and is the number of meta- domains.
After getting the corpora, aim extract the meta- domains in the corpora, that is, to divide the corpora into some (semantically) different clusters, to serve the meta- domains. We employ k-means (Macqueen, 1967; Arthur & Vassilvitskii, 2006) clustering algorithm to implement the separation and utilize bge-sma 1- 5 and bgg-small-22m 2h (Xiao al., 2023) to compute embeddings for the English and Chinese documents, respec- tively. See Figure 2 for the relationship between the number of meta-domains and Inertia (measuring the distance be- tween each data point and its centroid). Besides, we divide the code data directly according to the programming lan- guage. Ultimately, we construct 260 (120 Chinese 120 English + 20 Code) unique meta-domains. Each document in the corpora is labeled which meta-domain it originates from.
Specifically, we train a 260-class classifier to determine which meta-domain a given document originates from. We finetune Qwen2-1.5b-base to balance accuracy and efficiency. After training, the meta-domain classifier achieves a classification accuracy of 74.73% in our constructed test set. For further evalu- ating the performance of the meta-domain classifier, we also sample 1, 000 examples from each sub-dataset of The Pile (Gao et al., 2021). Following Equation 3, we obtain domain vectors predicted by the meta-domain classifier for each sub-dataset, as shown in Figure The distributions of sub-datasets of The Pile over meta-domains exhibit dis- tinctive patterns. This phenomenon indicates not only that
5b-base (Qwen Team, 2024)
Meta-Domain Classifier. We now present approach for representing unseen dataset using the previously estab- lished meta-domains. The methodology is straightforward yet effective: we assign each document in the unseen dataset to its corresponding meta-domains and then calculate the aggregate distribution across all documents. This compre-
IIn this paper, we N 1000, which is enough for an accurate and stable domain vector.
3

---

Domain2Vec: Vectorizing Datasets to Find the Optimal Data Mixture without Training
the various meta-domains have significant semantic differ- ences, but also that our classifier can accurately distinguish semantic features from different unseen datasets.
of The Pile (Gao et al., 2021) and RedPajama (Weber et al., 2024), as shown in Figure 3. The results of more validation sets can be seen in Figures 8 and 9. There are two findings:
3. Methodology
Given a validation set, there exists an optimal mix- ture ratio. For different validation sets, the ranking of mixture ratios varies significuntly.
In this section, we first introduce the task formulation of optimal data mixture discovery. We then present methodolo- gies for identifying the optimal data mixture using Do- MAIN2VEC without requiring additional training. We introduce two approaches: the first is grounded in the Distribution A Alignment Assumption (DA²). Moreover, we demonstrate how our DOMAIN2VEC can be integrated with previous works that model the relationship between mixture ratios and final performance, significantly enhancing the scalability of these existing approaches.
For the same validation set, the data mixture ratio rank- ing does not (nearly) change across model sizes. We calculate the correlation coefficients of data mixture rankings between the 83M model and the 1.6B model across diverse validation sets. The analysis yields a Spearman coefficient of 0.9743 and Pearson coeffi cient of 0.9947, providing robust statistical evidence for this consistency. These exceptionally high correla- tion values strongly support our finding that optimal mixture ratios are largely invariant to model size when evaluated on the same validation benchmark.
3.1. Task Formulation
During the pretraining phase of L.I,s, we typically col- lect training datasets Dtrain = {D1, D2,... ..., } from m sources (e.g., arXiv, Wikipedia, etc.). We also pre-define validation set D Duatid, which is of high quality and cor- responding to the final performance. Note that Deatid is often independently and identically distributed with Dtrain- For example, Liu et al. (2024) adopts Pile-CC (Gao et al., 2021) as the validation set and Gu et al. (2024b) adopts LIMA (Zhou et al., 2023) as the validation set. Accord- ingly, the data mixture = [ri. P2,.... F'm] TT, ,0 I ry S 1. >''r' T2 1 specifies the mixture ratio of the datasets. Let the trained LM be denoted as 8, and the validation loss of the LM be denoted as Lg. The objective of finding the optimal data mixture r" is usually to minimize the vali- dation loss, as shown formally in Equation 4. We denote C (r) as the validation loss of a LM pretrained on the data mixture r. r- = arg min min Lg
These finding aligns with prior work by Liu et al. (2024), which indicates that it is possible to find the optimal data mixture without training (Section 3.3) or simply training small models (Section 3.4).
3.3. Distribution Alignment Assumption (DA2)
We introduce how we directly apply our proposed Do- MAIN2 on finding optimal data mixture. We notice an intuitive law, that a lower validation loss [ is achieved when the data distribution of the training set is better aligned with the given validation set'. One essential question is that How do we model the data distribution of various datasets? Fortunately, according to Section 2, for the training dataset Devain. we obtain the vector represen- tation Virain E Rnxm, which models semantic features of DD Dirain- Correspondingly, for the validation set D Dralid- we also have its vector representation Uralid- After mixing Drrain with data mixture r, the final distribution over meta-domains of Dirain is Utrain Virain Therefore, based on the distribution alignment assumption, Equation 4 can be equivalently written as:
D.aated (r)) - arg min Peati (r)
(4)
3.2. Pilot Study: Mixture Ratio Ranking Holds across Model Sizes
r" = arg min Dist(V (Vrain TT, Dpatid)
(5)
We first conduct a pilot study for a critical research question: Could the optimal data mixture generalize across different model sizes? If the answer is Yes, it opens up the promising possibility that we could determine the optimal mixture ra- tio by simply training a small proxy model-og or even more efficiently, without training any model at all. To answer the questions, we mix C4 (Raffel et al., 2020) and Knowl- edge Pile (Fei et al., 2024) with different data mixtures (0, 0.2, ... 1.0) in Table 1. We pretrain two LMs with 83M and 1.6B parameters from scratch using the standard LM loss. During pretraining, we evaluate the validation loss of models trained with different mixture ratios on 20 subsets
where Dist(..) :) is a distance function used to measure the similarity between two vectors. Theoretically, numerous dis- tance function options are available, including Wasserstein (optimal transport) distance, Fuclidean distance, etc. In this paper, we use Huber Loss (Huber, 1964; Hastic et al., 2009) between two vectors to measure the distance. We also dis- cuss the choice of different distance functions in Appendix E. We present the pseudo code of DOMAIN2VEC+DA 2 in Ap- pendix B.
*We provide the detailed description in the Appendix A.
4

---

Vee: Vectorizing Datasets to Find the Optimal Data Mixture without Training
Model Pretrained
1.65
Hodel
Hodel
3.4
5.5
04:0.4
CL. 0.4
0.5
02:0.2
KP 0.5
D4.0.2 Kn
0.5
04:0.0
+ 1.0
06:0.0
1:0
C4:10
C4:100
KP:00
KP:0.
0.2
08
0.2
KP:08
0.0
C4: DC
K2
10
CR: 0.0
K2:10
1.00
196
8222
8502
8150
196
1150
1509 8150 30030
1908 855 10:0
9750
5003
1150
10:00
1756
15:00
1030
1964
1000
8003
8000
12:03
180
1360
5:00
10:00
19.0
5000
12:00
Figure 3. The validation loss on the EuroParl (The Pile) and Stackexchange (RedPajama) of models trained using data mixture in Table 1. The validation loss on other validation sets are shown in Appendix D.
3.4. Applying Domain2Ves to Prior Work
Equation 6 enables the prediction of validation loss on any meta-domain given a data mixture, which is also the function that we aim to fit. For unseen validation dataset, recall that any dataset including Pealid can also be viewed as linear addition of meta-domains and the domain vector of Dealid is denoted as Unatid have
There is one typical line of research focused on determining optimal mixture ratios, which aims to model the relationship between these ratios and the final validation loss using vari- ous functional approaches. That is, these approaches model [.Paaia (r) f(r) where f()) can take various reasonable forms as proposed in previous works. For example,
[g), -2,-" .4m] Therefore, we
D (r) =
Σ qi L (r)()) Σ qs fi(trriin)
Data Mixing Law (Ye et al., 2024) adopts f(r) = C + k - eex [Σ}= tig ri) to predict the validation loss on training domain i, where ci, ki, li are all unde- termined parameters to fit.
i=1
Σ qi f((Virain r).
RegMix (Liu et al., 2024) initially adopts Linear Regression approach, modeling the validation loss fr) w where needs fitting. Furthermore, it advances this concept by employing LightGBM (Ke et al., 2017) to more effectively fit the function f(-).
Now, we connect validation loss to the mixture ratio in the the domain vector space via our proposed DOMAIN2VEC. It is feasible to search the optimal mixture ratio r* by mini- mizing CPoatee (r). Note that this connection is built only the top of the meta-domains (i.c., fi for 1 I and can adapt with no cost to (1) any unseen validation set; (2) any unseen training set; (3) any number of training sets. Thanks to this property, we realize the efficiency and scalability by DOMAIN2V for prior approaches.
n),
We can directly integrate DOMAIN2VEC with these ap- proaches without modifying their core function, but in- stead perform the computations in the domain vector space. Thereby, we address two inherent limitations of these ap- proaches: (1) Efficiency: for modeling f(.) with m vari- ables 71..m 5 is expected to run experiments ((m)) times for different r to collect fitting points; (2) Scalability: When a new training source is introduced, one must re-collect fit- ting points and re-fit f()), which lacks of scalability.
Following RegMix (Liu et al., 2024), we use LightGBM (Ke et al., 2017) f()) to fit Equation 6 for each meta-domain (named as DOMAIN2 D Reg+Reg The pseudo code of DOMAIN2 + RegMix are shown in Appendix B. We sample 10, 500 diverse mixture ratios from a Dirichlet dis- tribution and we get the validation losses on each meta- domains by training 10, 500 small LMs. We also reserve some mixture ratios as testset and run experiments for eval- uating whether fitted function f() can accurately predict the validation loss for unseen mixture ratios. For various mixture ratios in the testset, we use the Spearman coefficient to measure the correlation between the predicted ranking and the actual ranking of performance under unseen mixture ratios. Note that we adopt correlation coefficient because it is more general metric than mean loss error with the goal
Specifically, novelly build the relationship fi(-) between the validation loss on the i-th meta-domain Dj (notated as CP:) and the domain vector Uraai after mixing training datasets by ratio r, that is, Vtrain r. Formally, for each meta-domain, we have
_i (r) = f. ((((aai)))= = ((Vrrai r),15t n.
(6)
'm can scale 10a in modern LLM training. For example, Fineweb (Penedo et al., 2024) consists of over 30k data dumps.
5

---

Domain2Vec: Vectorizing Datasets to Find the Optimal Data Mixture without Training
include C4 (Raffel et al., 2020) and Knowledge Pile (Fei et al., 2024). C4 is colossal and cleaned version of Com mon Crawl corpus. Knowledge Pile is a high-quality dataset that significantly improves the performance of LLMs in knowledge-related and mathematical reasoning tasks. We conduct our experiments on various validation datasets to perform comprehensive evaluation. We select 20 valida- tion datasets from The Pile (Gao et al.. 2021) and RedPa- jama (Weher et al., 2024). Since the optimal mixture ratio varies among the validation datasets, we instead predict the performance ranking across different preset mixture ratios. Specifically, we mix C4 and Knowledge Pile with different data mixtures as the training set as shown in Table 1.
50
Sptammn
2000
4000
5000
0000
10000
Num of Data Mature
Table 1. The preset data mixture ratios.
Figure 4. Relationship between the number of trained data mix- tures and the Spearman correlation.
Dataset
Data Mixture
C4
0
0.2
0.4
0.6
0.8
Knowledge Pile 1.0
0.8 0.6 0.4 0.2
0.0
to find the better mixture ratio than others. Moreover, the pilot study suggests that the predicted ranking holds across model sizes while the predicted loss becomes meaningless for inconsistent model sizes. As shown in Figure 4, the Spearman coefficient increases with the number of mixture ratios that we use for training and collecting fitting points. And finally we get an over 90% Spearman coefficient, which is quite accurate for predicting good mixture ratio for vari- ous domanns
Training & Evaluation Setup. We pretrain LLaMA- like (Grattafiori et al., 2024) models with 83M and 1.6B parameters from scratch using standard language modeling loss. Both models have a batch size of 1.5M tokens and a maximum sequence length of 4, 096. We use the AdamW optimizer (Loshchilov & Hutter, 2017) with gradient clip- ping at 1.0. The learning rate linearly warms up to 2e-4 over the first 100 steps, then decays to 2e-5 using cosine sched- uler over 10, 000 steps. More parameters are detailed in Ta- ble 7. Then, we evaluate DOMAIN2VEC using the Spearman and Pearson correlation coefficient between the predicted ranking and the actual ranking. We compare DOMAIN2 VFC with randomly ranking and an embedding denoted as kNN. Specifically, we use bge- 11v 5 to obtain embeddings and apply mean pooling to generate unique embeddings for each dataset and meta-domain. We then apply ANN based on Euclidean distance to compute the probability distributions of training and test datasets origi- nating from each meta-domain, treating these distributions as new domain vectors.
4. Experiments
In this section, we elaborate the implementation and exper- imental results to show how our proposed DOMAIN2VFE F.C helps find the optimal data mixture with less computation. The goal of optimizing the data mixture is to enhance the performance of LMs. The performance of LMs can be evaluated from two perspectives: 1) Text generation, often measured by LM loss on a validation set. We aim to min- imize the validation loss through finding the optimal data mixture; 2) Downstream task performance. The objective is to optimize task performance. As an overview for exprimen- tal results, by applying DOMAIN2VEC, we can accurately predict the ranking of data mixtures under various settings (e.g., training and validation sets). We also achieve compa- rable validation loss with the original data mixture from The Pile but only spend 51.5% training computational resources. Moreover, we use only 0.26% of the computational costs re- quired by DoReMi to find a data mixture with performance comparable to strong basclines like DoReMi.
baseline.
Experimental Results. We present the validation loss curves for various data mixtures in Figure 3 and Appendix D. It can be observed that, on most validation sets, incorpo- rating a certain amount of Knowledge Pile significantly re- duces the validation loss, even on the C4 validation set from RedPajama. We apply two DOMAIN2VEC-bare methods described in Section 3 to rank the data mixture from Table 1.
As demonstrated in Table 2, the ranking predicted by Do- MAIN2VEC exhibits a strong positive correlation with the actual ranking, significantly outperforming random guess- ing and ANN. The effectiveness of the INN method partially validates the rationale behind our meta-domain vocabulary construction. also important to note that our method
4.1. Validation Loss Minimization
Dataset & Data Mixture. We design some training and validation datasets to evaluate the performance to minimize the validation loss of our methods. Our training datasets
6

---

Demain2Vee: Vectorizing Datasets to Find the Optimal Data Mixture without Training
103M Model Pretrained on 25 Tokens
290M Model Pretrained on 05 Tokans
505H Model inetrained on 125 Tokens
10 Model iestrained on 208 Tokens
Domeen2iec= + DA-2
Domeen2Wer: DA-2
2.5
Dema nolte DA"2 Dama Reghlia
Demancleec DA12
DamairzWec Regghts
Dominazeec Rights
Demainzzee: HogMt HUMAN
3.5
3.2
2.20
33
6--
1.2
45%
-48%
310
55%
34-
-48%
-4.37%
-5.315
-4.93%
20-
-4 260
8000
500
104p
5009
Figure 5. The validation loss on the Pile-CC subset. DOMAIN2VEC achieves the comparable validation loss of Human (The model using original data mixture from The Pile), which only uses almost 51.5% training computational costs of Human. Using the same training cost, DOMAIN2VEC can reduce the validation loss by approximately 4. 72% compared to Human.
is a training-free approach, unlike prior works that rely on training small proxy models to rank data mixtures. Despite the more challenging setup, our method accurately predicts the rankings of different data mixtures.
ct al., 2024). The Pile-CC Only baseline (which trains the model solely on the Pile-CC subset) is included to verify the strong correlation between Pile-CC validation loss and downstream performance. The data mixtures for each base- line are shown in Table 5.
Table 2. The results of deploying the DOMAIN2VEC to predict the ranking of different Validation sets.
Training & Evaluation Setup. We pretrain LLaMA- like (Grattafiori et al., 2024) models from scratch using standard language modeling loss with model sizes ranging from 106M to 1B parameters. Following Hoffmann et al. (2022a), the token count for each model is 20 times corre- sponding parameter size. All models adopt batch size of 1M tokens and a maximum sequence length of 4, 096. We apply the AdamW (Loshchilov & Hutter, 2017) optimizer with gradient clipping at 1.0. The learning rate linearly warms up to 6e-4 over 1, 000 steps, then decays to 0 using a cosine scheduler at the end of training. More parameters are detailed in Table 7. For evaluation, we track the performance on Pile-CC CCC validation loss across different model sizes. Be- sides, we evaluate the performance of different data mixture using following benchmarks: Social IQA (Sap et al., 2019), HellaSwag (Zellers et al., 2019), PiQA (Bisk et al., 2019), OpenBookQA (Mihaylov et al., 2018), Lambada (Paperno et al., 2016), SciQ (Welbl et al., 2017), ARC Easy (Clark et al., 2018), COPA (Gordon et al., 2012), RACE (Lai et al., 2017), LogiQA (Liu et al., 2021), WinoGrande (Sakaguchi et al., 2021), and MultiRC (Khashabi et al., 2018). We uti- lize LM Evaluation Harness (Gao et al., 2024) to evaluate these models and report the average score across O-shot to 5-shot settings in Table 3.
Metries
Random
ANN
DOMAIN2VEC4DA?
DoMAIN2VEc+RepMll
Peanson
0.0300
0.4014
0.583
0,3881
Speurman
0.0497
0.3543
0.6657
0.4629
4.2. Downstream Task Performance Maximization
In this section, we demonstrate how DOMAIN2 can be used to identify the optimal data mixture for maximizing downstream task performance. One challenge is modeling the relationship between data mixture and downstream per- formance. Fortunately, Liu et al. (2024) finds that validation loss on Pile-CC correlates most strongly with downstream performance across their evaluations. To align with prior work, we follow and use the same validation datasets Liu et al. (2024). Thus, our goal is to identify data mixture that minimizes validation loss on Pile-CC. Experimental results show that DOMAIN2VER EC predicts a data mixture with per- formance comparable to DoReMi (Xie et al., 2023a), while using only 0.26% computational cost.
Datasets & Baselines. We follow RegMix (Liu et al., 2024) and use The Pile (Gao et al., 2021) as our training datasets. The Pile is an 825 GB English text corpus used for LLM pretraining. In line with RegMix, we use only the 17 components of The Pile that do not have copyright issues. Our goal is to identify the data mixture that min- imizes validation loss on the Pile-CC subset to improve downstream task performance. We compare our approach with several baselines, including Human (the original data mixture), DoReMi (Xie et al., 2023a), and RegMix (Liu
Implementation Details. We predict the optimal data mixture by applying Equation 5 (DoMain2 and Equation 7 (DOMAIN2VEC+RRgNN 100, 000 data mixtures from a Dirichlet distribution based on the token distribution of these components. Using these mixtures, we predict the optimal data mixture by our pro- posed two methods. We select top-100 predicted data mix-
We generate
7

---

S (Di), Domain2Vec: Vectorizing Datasets to Find the Optimal Data Mixture without Training
Table Average downstream task performance different models pretrained on different data mixtures. Similar to Liu et al. (2024). Human refers the original data mixture from The Pile. Pile-CC is a golden training set which can 100% correspond to validation set to validate our propose DAA², data mixtures are shown in Table and Table calculated data mixture are shown in Table
Benchmark
Human
DoReMi
Pile-CC Only
RegMix
DOMAIN2VEC + DA
DOMAIN2VEC
RegMix
Social IQA
0.367
0,380
0.381
0.382
0.372
0.375
HellaSwag
0,319
0.346
0.351
0.351
0.335
0,338
PIQA
0.615
0.639
0.644
0.647
0.635
0,639
OpenBookQA
0.264
0.275
0.276
0.276
0.275
0.272
Lambada
0.199
0.240
0.247
0.241
0.219
0.232
SciQ
0,710
0.695
0.688
0.708
0.701
0,701
ARC Easy
0,411
0,428
0.436
0.438
0,427
0,426
COPA
0.621
0.651
0.660
0.653
0.638
0.641
RACE
0.274
0.291
0,288
0.288
0.279
0,282
LogiQA
0,272
0.275
0.272
0.272
0.269
0.278
WinoGrande
0,512
0,516
0.515
0,513
0.513
0,510
MultiRC
0.521
0.528
0.515
0.529
0.524
0.534
Average Performance
0.424
0.439
0.439
0.441
0.432
0.436
3.7 x 10
3.5 x 10
9.66 x 10
9.66 10t6
Estimated FLOPs
0
(100%)
0
(9.46%)
(0.26%)
(0.26%)
tures and average them as the final data mixture. This trick aligned with previous work (Liu et al., 2024) for more ac- curate and stable results. As stardard practice, each subset of The Pile is trained for at most one epoch. When opti- mizing the mixture ratio r [ri, r2, ... ,m,] other than the trivial restriction Σ r, 1, note that there is another data amount restriction, that is #TotalTokon which is to remove data mixtures which require exceeding tokens in some subsets. Therefore the optimal data mixture predicted by DOMAIN2VEC may vary depending on the number of trained tokens, as well as the size of the model. This restriction is different with Section 4. where each dataset size is seen as unlimited.
FleCC
Fr PaiPapers
Hin Scherter
UiuattuRR
Treetar
DM Mithomatis
Hor lortiee
Farsheri
Mic pedis tun)
Grmuua
Succktechngg
Experimental Results. As shown in Figure 5, our pro- posed DOMAIN2VEC + 2A² and DOMAIN2VEC: REG- MIX significantly improve training efficiency on Pile-CC compared to Human. Specifically, DOMAIN2VEC + DA² and DOMAIN2VEC REGMIX require only about 55. 53.3 and 51.50% of the training steps, respectively, to achieve the same validation loss Human. Compared Human under the same compute budget, DOMAIN2VEC + DA and DOMAIN2VEC + REGMIX reduce validation loss by ap- proximately 4.04% and 4.64%, and improves downstream performance by an average of 1.89% and 2.83%, respec- tively. Table 3, we report the average performance of LMs trained on data mixtures from various baselines across range of downstream tasks. "Pile-CC only" shows 3.54% average accuracy improvement over Human, indicating that training more tokens from Pile-CC enhances downstream performance. Importantly, "Pile-CC only" is good when we regard Pile-CC as validation set. However, a more practical scenario where validation set is somewhat else, we cannot manually find such a golden training set which
Figure 6. Visualization (t-SNE of domain vectors of The Pile.
can 100% correspond to validation set. To this end, we can use our proposed Domain2Vec to get a comparable down- stream performance with lowest cost by mixing datasets from different sources. Notably, DOMAIN2VEC + DA and DOMAIN2VEC REGMIX, using only about 0.26% of the FLOPs required by DoReMi, achieve performance comparable to DoReMi, RegMix, which demonstrates the computational efficiency of DOMAIN2VEC.
Visualization. To investigate further, we employ t- SNE (Van der Maaten & Hinton, 2008) to visualize the domain vectors of each component in The Pile, as shown in Figure This visualization reveals several desirable properties of the learned vectors. The representation space
8

---

Domain2Vec: Vectorizing Datasets to Find the Optimal Data Mixture without Training
exhibits strong clustering behavior where semantically re- lated datasets naturally group together, indicating effective capture of domain-specific characteristics. Related domains such as academic literature (PubMed, arXiv) and techni- cal repositories (GitHub, StackFxchange) demonstrate spa- tial coherence, while maintaining well-defined yet flexible boundaries between different domains. The representation spans diverse domains in The Pile, demonstrating robust generalization capabilities across heterogeneous data types.
computational costs and making it hard to define an ideal reference model. Some works aim model the functional relationship between data mixture and the LM performance. Inspired by scaling laws (Kaplan et al., 2020; Hoffmann et al., 2022a), Ye et al. (2024) introduce Data Mixing I Laws, which describe this relationship using an exponential form. Ge et al. (2024) propose BiMix, a scaling law that considers both compute and data mixture. Que et al. (2024) and Wang et al. (2025) develop scaling laws for continual pretraining, and how mixture ratio as one variable impacts LM loss is modeled. Recently, Liu et al. (2024) propose Linear Re- gression to model the validation loss across different data mixtures, showing a strong and promising performance.
4.3. Discussion on Overfitting
We noticed that some readers interpret our approach DA as a form of "overfitting": optimizing on a selected validation set. We offer the following explanations:
All these prior works face two main issues: 1) Computa- tional Efficiency: For example, the estimated FLOPs for DoReMi and RegMix are high to 3.7 x 10 and 3.5 10 respectively, for calculating less than 22 datasets. More- over, the computational complexity of these methods will grow non-linearly as the number of datasets increases. 2) Lack of Scalability: When the components of the training dataset change (e.g., adding some new datasets), previous methods (Ye et al., 2024; Liu et al., 2024) require resam- pling data mixtures, retraining proxy models, and then re- performing the fitting process. In this paper, we introduce DOMAIN2 VEC to decompose any dataset into a linear com- bination of meta-domains. concepts with prior meta-learning works, such as Jomaa et al. (2021) and Chen et al. (2024), which explore dataset representation in latent spaces. While sharing this concept, DOMAIN2 differs in both purpose and implementation, and we focus on the data mixture in LM pretraining.
The validation set that we define is actually a guide dataset, which is necessary requirement for optimiza- tion data mixture and common setting in related works (see Section 3.
In Section 4.1, we conduct experiments on various validation sets, and the performance demonstrates good stability. In fact, our proposed DA does not even require training, thus "overfitting" is not applicable.
In Section 4.2, we choose Pile-CC as the validation set but ultimately test model performance on bench- marks from 12 downstream tasks, further preventing overfitting risks.
DOMAIN2VEC shares some
5. Related Work
Recent research on optimizing data mixture can be broadly divided into two lines. The first line implicitly adjusts data mixture by down-sampling from various datasets based data quality. For example, Lin et al. (2024) propose RHO- 1, which uses Selective Language Models to select tokens that align the data mixture with the ideal ratio. Instead of token-level selection, Ankner et al. (2024) filter low-quality samples using the perplexity of small reference models. Thakkar et al. (2023) demonstrate that the Influence Score can guide data re- seieethtting, while their subsequent work introduces online data selection method that eliminates the need for reference models.
6. Conclusion
In this work, we introduce DOMAIN2VEC, a novel method to capture the underlying features of datasets by decom- posing datasets into a linear combination of several meta- domains. It enables tion for arbitrary datasets. Building these domain vec- tors, we introduce training-free free approach by Distribu- tion Alignment Assumption (DA7) to identify optimal data mixtures for language model pretraining Furthermore, Do- MAIN2 VEC seamlessly integrates with existing methods, significantly enhancing their efficiency and scalability while establishing direct relationship between model perfor- mance and computed domain vectors-all without requiring retraining when training datasets change. Our experimental results demonstrate that both DOMAIN2VEC+DA? and Do- MAIN2 V ReggR achieve comparable text generation and downstream task performance with reduced compu- tational overhead compared to existing approaches. We believe this work offers valuable insights into optimizing data mixtures for language model pretraining and paves the way for more efficient training strategies.
acquire vectorized representa-
The second line focuses on explicitly adjusting data mix- ture by modeling the relationship between data mixture and language model performance. The simplest approach is to observe the performance of various data mixtures and se- lect the optimal one, done during Gopher training (Rae et al., 2022). This is costly and difficult to scale for larger models. Xie et al. (2023a) propose DoReMi and use a small proxy model to re-weight data from different domains, improving training efficiency for larger models. However, DoReMi still requires pre-trained reference model, adding

---

Deeain2Vec:
Vectorizing Datasets to Find the Optimal Data Mixture without Training
Impact Statement
N., Presser, S., and Leahy, C. The pile: An 800gb dataset of diverse text for language modeling. CoRR, abs/2101.00027, 2021. URL https: /ttp:s:///xxv../////
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Gao, L., Tow, J., Abbasi, B., Biderman, S., Black, S., DiPofi,
A., Foster, C., Golding, L.., Hsu, J., Le Noac'h, h, A., Li,
Acknowledgements
H., McDonell, K., Muennighoff, N., Ociepa, C., Phang,
J., Reynolds, L., Schoelkopf, H., Skowron, A., Sutawika,
This work was supported by the National Natural Sci- ence Foundation of China (No. U24B20181) and Fu- jian Provincial Natural Science Foundation of China (No. 2024108371).
L.,, Tang, E., Thite, A., Wang, B., Wang. K., and Zou, framework for few-shot language model evaluation, 07 2024. URL https: zenodo. org/ records/ 12608602.
References
Ge, C., Ma, Z., Chen, D., Li, Y., and Ding, B. Data mixing made efficient: A bivariate scaling law for lan- guage model pretraining, 2024. URL https: //arxv.. brg/abs/2400 14908.
Ankner, Z., Blakeney, C.,, Sreenivasan, K., Marion, M., Leavitt, M. L., and Paul, M. Perplexed by perplexity: Perplexity-base based data pruning with small reference mod- els, 2024. URL https:/ / wwwpll//.rg/////////e/// 20541.
Gordon, A., Kozareva, Z., and Roemmele, M. SemEval- 2012 task 7: Choice of plausible alternatives: An eval- uation of commonsense causal reasoning. In Agirre, E., Bos, J., Diab, M., Manandhar, S., Marton, Y., and Yuret, D. (eds.), *SEM 2012: The First Joint Confer- ence on Lexical and Computational Semanties Vol- une 1: Proceedings of the main conference and the shared task, and Volume 2: Proceedings of the Sixth International Workshop on Semantic Evaluation (Se- mEval 2012), pp. 394-398, Montréal, Canada, 7-8 June 2012. Association for Computational Linguistics. URL arggs12-1052
Arthur, D. and Vassilvitskii. S. k-means++: The advantages of careful seeding. Technical report, Stanford, 2006.
Bisk, Y., Zellers, R., Bras, R. L.,, Gao, J., and Choi, Y. Piqa: Reasoning about physical commonsense in natural lan- guage. In AAAI Conference on Artificial Intelligence, 2019. URL https: /api. http:/app......a.m.mmmaa..mmaa.e semanticcccoll.. org/Coz Ttg/CorpusIDp 208290939.
Chen, J. Lin, Z., Chen, Q., and Sun, J. Cross-table pretrain- ing towards a universal function space for heterogeneous tabular data, arXiv preprint arXiv:2446.00228 002811 2024.
https: http::://aclaattholl...
Grattafiori, A., Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A..
Clark, P., Cowhey, I., Etzioni, O., Khot, T., Sabharwal, A., Schoenick, C., and Tafjord, O. Think you have solved question answering? try are, the ai2 reasoning challenge. ArXiv, 103.1503 2018. URL https: / api. semanticscholle org http://ccccaaaa.ooooooooorsaaaaa :00p:/ 3922816.
Vaughan, A., al. The llama 3 herd of models. arXiv preprint arXiv:2407.2178.. 2024.
Gu, J. Yang, Z., Ding, C., Zhao, R., and Tan, F. CMR scaling law: Predicting critical mixture ratios for con- tinual pre-training of language models. In Al-Onaizan, Y., Bansal, M., and Chen, Y.N. (eds.), Proceedings of the 2024 Conference on Empirical Methods in Nat- ural Language Processing, pp. 16143-16162, Miami, Florida, USA, November 2024a. Association for Compu- tational Linguistics. doi: 10. 10.18663/v1220044cmml- 903. URL https: http:///caaathhhooceo org/ /204.. emn p m- n. 903/.
DeepSeck-AL. v2: A strong, economical, and efficient mixture-of-experts language model, 2024.
Fan, S., Pagliardini, M.. and Jaggi, M. DOGE: Do- main reweighting with generalization estimation. In Second Agent Learning in Open-Endednen Workshop, wttp:/viwwwwe/ffournrto.
2023. URL
kd=gikgsq
Fei, Z., Shao, Y., Li, L.,, Zeng, Z., He, C., Yan, H., Lin, D., and Qiu, X. Query of Unearthing large scale domain-specific knowledge from public corpora, 2024. URL https::///// aariv / abs/2401. 240.244222..4.
Gu, Y., Dong, L., Wang, H., Hao, Y., Dong, Q., Wei, F,, and Huang, M. Data selection via optimal control for language models. arXiv preprint arXiv: 2410.07064, 2024b.
Gao, L., Biderman, S., Black, S., Golding, L., Hoppe, T.,
Guo, S., Ren, Y., Albrecht, S. V., and Smith, K. Sample rela- tionships through the lens of learning dynamics with label
Foster, C., Phang, J., He, H., Thite, A., Nabeshima,
10

---

Domain2Vec: Vectorizing Datasets to Find the Optimal Data Mixture without Training
information. In First Workshop on Interpolation Regular- izers and Beyond at NeurlPS 2022, 2022. URL https: //openreview.oo net/ forrm (iteruummimmmImm
Lai, G., Xie, Q.. Liu, H., Yang, Y., and Hovy, E. RACE: Large-scale ReAding comprehension dataset from ex- aminations. In Palmer, M., Hwa, R., and Riedel, S. (eds.), Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing, pp- 785-794, Copenhagen, Denmark, September 2017. Association for Computational Linguistics. doi: 10.18653/1//17-1111 org
Hastie, T., Tibshirani, R., Friedman, J. H., and Friedman, J. H. The elements of statistical learning: data mining, inference, and prediction, volume 2. Springer, 2009.
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., de Las Casas, D., Hendricks, L. A., Welbl, J., Clark, A., Hennigan, T., Noland, E., Millican, K., van den Driessche, G., Damoc, B., Guy, A., Osindero, S., Simonyan, K., Elsen, E., Rae, J. W., Vinyals, O., and Sifre, L. Training compute-optinal large language models. 2022a. URL https: /arxiv.. org/ org/abs/ 2203. 15556.
URL https: http:://aclaatttohoo wttp:/gyyyy.gp//1-1100 /tgp111102
Lin, Z.., Gou, Z., Gong, Y., Liu, X., Xu, R., Lin, C.,, Yang. Y., Jiao, J., Duan, N., Chen, W., et al. Not all tokens are what you need for pretraining. Advances in Neural Information Processing Systems, 37:29029-29066 2024.
Liu, J., Cui, L.., Liu, H., Huang, D., Wang, Y., and Zhang, Y. Logiqa: challenge dataset for machine reading com- prehension with logical reasoning. In Proceedings of the Twenty-Nimth International Joinl Conference on Artificial Intelligence, IJ,AI' 2021. ISBN 9780999241
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., de las Casas, D., Hendricks, L. A., Welbl, J., Clark, A., Hennigan, T., Noland, E., Millican, K., van den Driessche, G., Damoe, B., Guy, A., Osindero, S., Simonyan, K., Elsen, E., Vinyals, O., Rae, J. W., and Sifire, L. An empirical analysis of compute-optimal large language model training. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.), Advances in Neural Information Processing Systems, 2022b. URL https: /ooenreview. net foormm? ? 1dd=BBBRUOOAAP
Liu, Q.. Zheng, X., Muennighoff, N., Zeng, G., Dou, L.,, Pang, T., Jiang, J., and Lin, M. Regmix: Data mixture as regression for language model pre-training, 2024. URL https: / wrwaa.ivvvooo.... org/ /ass///////// 240.. 01492.
Loshchilov, and Hutter, Decoupled weight decay regu- larization. arXiv preprint 17Xiv: 2017.
Huber, J. Robust Estimation of a Location Parameter. The Annals of Mathematical Statistics, 35(1):73 101, 1964. doi: 10. 2 //doi. /tt 12111 htt.a//mm///111111/11111
Macqueen, J. Some methods for classification and anal- ysis of multivariate observations. In Proceedings of 5- th Berkeley Symposium on Mathematical Statistics and ProbabilityUUnvvessty of California Press, 1967.
URL https:
Jomaa, H. S., Schmidt-Thieme, L., and Grabocka, J. Dataset2vec: Learning dataset meta-features. Data Min- ing and Knowledge Discoverry33(33)9944 985, 2021.
Mihaylov, T., Clark, P., Khot, T., and Sabharwal, A. Can a suit of armor conduct electricity? a new dataset for open book question answering. In Conference on Empirical Methods in Natural Language Processing, 2018. URL https:/ semanticschhlar. org/ Sts//CorpuaI:::2183555
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language mod- els, 2020. URL https:tt//.. / aaaxiiiiiv v. org/ abs /2011 08361.
OpenAl. Gpt-4 technical report, 2024. URL https: // arxiv. stwii..org/bbs//333..o 08774.
Ke. G., Meng, Q. Finley, T., Wang. T., Chen, W., Ma, W., Ye, Q., and Liu, T.-Y. Lightghm: A highly efficient gradient boosting decision tree. Advances in neural infor- mation processing systems, 30, 2017.
Paperno, D.,, Kruszewski, G., Lazaridou, A., Pham, N. Q.. Bernardi, R., Pezzelle, S., Baroni, M., Boleda, G., and Fernández, R. The LAMBADA dataset: Word pre- diction requiring a broad discourse context. In Erk, K. and Smith, N. A. (eds.), Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Lonng Pupers), pp. 1525-1534, Berlin, Germany, August 2016. Association for Compu- tational Linguistics. doi: 10.18655//1/116-1144 URL wttp:hhowooooooooo11111111
Khashabi, D., Chaturvedi, S., Roth, M., Upadhyay, S., and Roth, D. Looking beyond the surface: A challenge set for reading comprehension over multiple sentences. In Walker, M., Ji, H., and Stent, A. (eds.), Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long Papers), pp. 252- 262, New Orleans, Louisiana, June 2018. Association for Computational Linguistics. doi: 10.18653/v1/N18- 1023. URL https: //ac//accaaacc anthologyyy.ooooo org.Ni1 8-1033.
https:
Penedo, G., Kydlíček, H., allal, L. B., Lozhkov, A., Mitchell, M., Raffel, C., Werra, L. V., and Wolf, T. The fineweb datasets: Decanting the web for the finest text data at

---

Demain2Vec: Vectorizing Datasets to Find the Optimal Data Mixture without Training
scale. In The Tiiht Conference Neural Informa- fion Processing Systems Datasets and Benchmarks Track, 2024. URL https: http:///pppeeeeeeeeeeeeeeeeeeeee wwwpeenwwwwewwweeeeeeeeeeeeeeeee forum? id-n6SCkn2QaG.
data reweighting for language model pre-training. In The 2023 Conference Empirical Methods in Nat- ural Language Processing, 2023. URL https:// openreview net httr:/ooormooii.//AAMMM
Que, H., Liu, J., Zhang, G., Zhang, C., Qu, X., Ma, Y., Duan, F., Bai, Z., Wang, J., Zhang, Y., et al. D-cpt law: Domain-specific continual pre-training scaling law for large language models. Advances in Neural Information Processing Systems, 37:903188 77:90318-9033 18-9454,, 2024.
Touvron, H., Martin, L.,, Stone, K., Albert, P., Almahairi, A., Babaci, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., Bikel, D., Blecher, L., Ferrer, C. C., Chen, M., Cucurull, G., Esiobu, D., Fernandes, J., Fu, J., Fu, W., Fuller, B., et al. Llama 2: Open foundation and fine-tuned chat models, 2023. URL https: //arxiv.. org/ abs / 307 7. 09288.
Qwen Team. Qwen2 technical report, 2024. URL https: //arxiv. org/ abs / 2407 10671.
Van der Maaten, L. and Hinton, G. Visualizing data using t-sne. Journal of machine learning research, 9(11), 2008.
Radford, A. Improving language understanding by genera- tive pre-training. 2018.
Wang, X., Tissue, H., Wang, L., Li, L., and Zeng, D. D. Learning dynamics in continual pre-training for large language models. arXiv preprint arXiv: 2505.07796, 2025.
Rae, J. W., Borgeaud, S., Cai, T., Millican, K., Hoffmann, J., Song, F., Aslanides, J., Henderson, S., Ring, R., Young. S., Rutherford, E., Hennigan, T., Menick, J., Cassirer, A., Powell, R., van den Driessche, G., Hendricks, L. A., Rauh, M., Huang, P.-S., Glaese, A., et al. Scaling lan- guage models: Methods, analysis & insights from train- ing gopher, 2022. URL https: /tt///aaxvvvv/// org/ abs/ 2112. 11446.
Weber, M., Fu, D., Anthony, Q. Oren, Y., Adams, S., Alexandrov, A., Lyu, X., Nguyen, H., Yao, X., Adams, V., Athiwaratkun, B., Chalamala, R., Chen, K., Ryabinin, M., Dao, T., Liang, P., Ré, C., Rish, I., and Zhang, C. Redpajama: open dataset for training large language models. arXiv preprint arXiv: 24/1. 123372, 2024.
Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with a unified text-to-text transformer. J. Mach. Learn. Res., 21(1), jan 2020. ISSN 1532-4435.
Welbl, J., Liu, N. F., and Gardner, M. Crowdsourcing multiple choice science questions. In Derczynski, L., Xu, W., Ritter, A., and Baldwin, T. (eds.), Proceedings of the 3rd Workshop on Noisy User-generated Text, pp. 94-106, Copenhagen, Demmark, September 2017. Association for Computational Lititttiiionnnniiiiiiiiii Linguistics. doi: 10.18653/v1/W11-4415
Sagawa#. S., Koh P. W., Hashimoto, T. B., and Liang, P. Distributionally robust neural networks. In In- ternational Conference on Learning Representations, 2020. URL
URL https: ht//aclaaaaaaa https////ccaaaahhooo..oooo org /117//177444
id=ryxGuJIFvs:
Xiao, S., Liu, Z., Zhang, P., and Muennighoff, N. C-pack: Packaged resources to advance general chinese embed- ding, 2023.
Sakaguchi, K., Bras, R. L., Bhagavatula, C., and Choi, Y. Winogrande: an adversarial winograd schema challenge at scale. Commun. ACM, 64(9):99-106 August 2021. ISSN 0001-0782, doi: 10. 1553447438. URL https: // //doi. wwt.o/gg//ooooo 0111/3344433
Xie, S. M., Pham, H., Dong, X., Du, N., Liu, Lu, Y., Liang, P., Le, Q. V., Ma, T., and Yu, A. W. Doremi: Optimizing data mixtures speeds up language model pre- training. In Thirty-seventh Conference on Neural In- formation Processing Systems, 2023a. URL https: //operreview net/ www.foormmmooo/ Xuw.miidddxuuBUHehh.
Sap, M., Rashkin, H., Chen, D., Le Bras, R., and Choi, Y. Social IQa: Commonsense reasoning about social interactions. In Inui, K., Jiang, J., Ng, V., and Wan, X. (eds.). Proceedings of the 2019 Conference on Empir- ical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pp. 4463-4473, Hong Kong, China, November 2019. Association for Compu-
Xie, S. M., Santurkar, S., Ma, T., and Liang, P. S. Data selection for language models via importance resampling. Advances in Neural Information Processing Systems, 36: 34201-34227, 2023b.
tational Linguistics. doi: 10. 10.18653//1/D19-1144 URL wtt./oogg//////// https: / /acl wtt.aannnooolol...o/011111535550
Ye, J., Liu, P., Sun. T., Zhou, Y., Zhan, J., and Qiu, X. Data mixing laws: Optimizing data mixtures by predicting language modeling performance, 2024. URL https: 16952.
Thakkar, M., Bolukbasi, T., Ganapathy, S., Vashishth, S., Chandar, S., and Talukdar, P. Self-influence guided
//arxiv.org/abs/
12

---

Domain2Vec: Vectorizing Datasets to Find the Optimal Data Mixture without Training
Zellers, R., Holtzman, A., Bisk, Y., Farhadi, A., and Choi, Y. Hellaswag: Can machine ally finish your sentence?
In Anmual Meet-
ing of the Association for Computational Linguisties, 2019. URL https://app /api: / seeanticcchholol............ Org/CorpusID: 159041722.
Zhou, C., Liu, P., Xu, P., Iyer, S., Sun, J., Mao, Y., Ma, X., Efrat, A., Yu, P., YU, L.. Zhang, S., Ghosh, G., Lewis, M., Zettlemoyer, L.,, and Levy, O. LIMA: Less is more for alignment. In seirth-seventh Conference on Neural Information Processing Systems, 2023. URL ht tps: / /ppenrevieww net forum? 1d=KBmoKKX22
13

---

Domain2Vec: Vectorizing Datasets to Find the Optimal Data Mixture without Training
A. Detailed Description of the Distribution Alignment Assumption
In this section, we will introduce the detailed description of the Distribution Alignment Assumption for language model pretraining.
In the scenario of finding the optimal data mixture for language model pretraining, the validation set Deatid is fixed, and should adjust the data mixture to construct the training set Dirain to achieve lower validation loss calculated by Equation 8, where / is parameters of pretrained language model.
X
ExwDsnttd log P(X X
log(P(x 0))
(8)
Typically, pretrain language models via next token prediction (Radford, 2018) like Equation 9.
i arg max Ex X~ log P(X10)
[
(9)
arg max E. X X
Σ
log (og(P(a((zes:)))
That is, we need to find a that maximizes the expected probability of X Dtrain, which is also known Maximum Likelihood Estimation (MLE). When the data distributions of Durin and Daaldd are aligned, the optimization target of language models pretraining process equals find a / that maximizes the expected probability of X -aattdd Therefore, we introduce the Distribution Alignment Assumption for language model pretraining, a novel method to find the optimal data mixture without training.
B. Algorithm
In Algorithm 1, we show the pseudo code for acquiring the domain vector for pretraining datasets.
In Algorithm 2 and 3, we show the pseudo code for how to use DOMAIN2VEC to find the optimal data mixture, including Distribution Alignment Assumption, and applying DOMAIN2VEC to RegMix.
Note that when applying DOMAIN2VEC+DA or DOMAIN2VEC=RegMM one could also average the k-best ratios in the K sampled candidates data mixture. We present top-1 as one example in the pseudo codes. We adopt top-1 for direct comparison in Section 4.1, while adopt top- 100 in Section 4.2, which is aligned with RegMix (Liu et al., 2024).
for getting more stable and accurate results,
Algorithm 1 DOMAIN2VEC
Require: Training datasets Dtrain = (D1, D2..., D validation dataset 'D Dealid meta-domain classifier Classifier
1: Domain vectors Virain =
2: for 1 to do
3: Sample N documents doci...N from D:
4: N j=1 Classifier(docy), where doe; E D
5: Verain (V Vrainn, w.]
6: end for
7: Sample N data points from Dratid
8: Vealid = N Σ}=1 Classifier(doc;) where doc; E D valiid
9: Return: Virain [U1, 02 ..., Um)) Unatid
C. Data Mixture of Different Methods
In this section, we will show the data mixture on The Pile (Gao et al., 2021) of different methods we used in this paper for reproduction. In Table 5, we show the optimal data mixture predicted by DOMAIN2VEC: DA and DOMAIN2 + RegMix. It should be noted that, to avoid the over-fitting problem, any subset of The Pile (Gao et al., 2021) will be only trained at most one epoch. Because we adopt rejection sampling to filter out certain unreasonable data mixtures. The data mixture predicted may change as model sizes change.
14

---

Domain2Vec: Vectorizing Datasets to Find the Optimal Data Mixture without Training
Algorithm 2 DOMAIN2VEC+DA**
Require: Domain vectors of training datasets Virain U1: U2, ..... Um, domain vectors of validation dataset Vealids token
distribution of training datasets Otrain-
1: Sample K candidates data mixture r; from Dirichlet(Atrrint
2: The optimal data mixture r = 71
3: for = 2 to
4: if Dist DiiranV T, Uvalid) < Dist Virain r", Uvalid) then
5: r- ="=
6: end if
7: end for
8: Return: the optimal data mixture r
Algorithm 3 DOMAIN2VFC++RggM
Require: Domain vectors of training datasets Vtrain [1) 12, Um. domain vectors of validation dataset Veatid
[91. 92.-.- 9n. T token distribution of training datasets asrains fitted model for each meta-domain f((-).
1: Sample K" candidates data mixture r; from Dirichlet
2: The optimal data mixture r" T1
3: Def [(r) Σ qi fi(Virain r) 1=1
4: for = 2o K do
5:
if ((r:) < [(r*) then
6:
r' ===10 r;
7:
[(r*)= L(r;)
8: end if
9: end for
Return: the optimal data mixture
D. Experimental Results of Pilot Study
In this section, we report the validation loss on various datasets arXiv, C4, Book3, PG19 from RedPajama (Weber al., 2024). and BookCorpus2, DM Mathematics, Enron Emails, FreeLaw, HackerNews, NIH ExPorter, OpenSubtitles, OpenWeb'Tee PhilPapers, PubMed Abstracts, PubMed Central, USPTO Backgrounds, Ubuntu IRC, Youtube Subtitles from The Pile (Gao et al., 2021) in Figure 3, Figure and Figure 9.
E. Comparative Study on Different Distributional Measures of DA
In Section 3.3, we use Huber Loss measure the similarity of domain vectors. Technically, Huber loss combines the advantages of LI and L2 distance. In Table 4, add the results of different distributional measures. As shown in the Table 4, Huber Loss shows better performance than L1/L2/JS Distance. Additionally, Wasserstein distance is a very great option. However, it would require extra metric space matrix, M, to measure the distance between two domain vectors. In this work, the metric space, M E 260xx660 is actually the "dataset transition cost" between each two meta-domains, and is non-trivial, Each element M. G.j could be estimated via L., the loss at meta-domain j after training on meta-domain i, which requires additional computational resources. Considering that Huber Loss already achieved very positive results, we did not conduct this experiment. We believe that Wasserstein distance can also present positive result (even better) if the metric space well estimated, and we leave this for future work.
15

---

Domain2Vec: Vectorizing Datasets to Find the Optimal Data Mixture without Training
Table 4. Huber Loss shows hetter performance than 1.1/1.2/IS Distance.
Distributional Measure
Pearson
Spearman
Huber Loss
0,5833
0.6657
JS Distance
0,,4277
0.5000
LI Distance
0,4830
0,5400
L2 Distance
0.5720
0.6429
Table 5. The data mixture of The Pile (Gao et al., 2021) from different baselines, which aligns with the data mixture used in Liu et al. (2024).
Data Mixture
Human
DoReMi
Pile-CC Only
RegMix
ArXiv
0.134
0.004
0.0
0.001
FreeLaaw
0.049
0.005
0.0
0.001
NIH ExPorter
0.007
0.008
0.0
0.001
PubMed Central
0.136
0.006
0.0
0.003
Wikipedia (en)
0.117
0.086
0.0
0.016
DM Mathematics
0.025
0.002
0.0
0.0
Github
0.054
0.022
0.0
0.0
PhilPapers
0.003
0.034
0.0
0,0
Stack Exchange
0.118
0.019
0.0
0.0
Enron Emails
0.004
0.009
0.0
0.002
Gutenberg (PG-19)
0.025
0.009
0.0
0.002
Pile-CC
0.142
0.743
1.0
0.87
Ubuntu IRC
0.009
0.011
0.0
0.064
EuroParl
0.005
0.008
0.0
0.0
HackerNews
0.01
0.016
0.0
0.012
PubMed Abstracts
0.107
0.014
0.0
0.024
USPTO Backgrounds
0.053
0.004
0.0
0.002
Figure 7. The Domain Vector of each sub- datastt of The Pile (Gao et al., 2021), where each row corresponds to a sub-datased and each column corresponds
domain. The higher the proportion data belonging particular meta- domain, the closer the color the corresponding cell is to blue). Additionally, since The Pile primarily consists of English texts, we only display the distribution on English meta-domains for clarity.
16

---

Domain2Vec: Vectorizing Datasets to Find the Optimal Data Mixture without Training
Table The optimal data mixture predicted by DOMAIN2VEC DAA and DOMAIN2VEC RegMix. To avoid the over over-fitting problem. any subset of The Pile (Gao al., 2021) will be trained most epoch. And adopt rejection sampling to filter out certain unreasonable data mixtures. Thus, the data mixture predicted may change as model sizes change.
Data Mixture
DOMAIN2VEC+DA"
DOMAIN2VEC+RegMI
106M
290M
595M
IB
106M
290M
595M
1B
ArXiv
0.0131
0.0131
0.0389
0.0431
0.0152
0.0070
0.0114
0.0103
FreeLaw
0.0076
0.0076
0.0316
0.0305
0.0395
0.0267
0.0339
0.0268
NIH ExPorter
0.0008
0.0008
0.0028
0.0023
0.0000
0.0199
0.0000
0.0000
PubMed Central
0.0773
0.0773
0.0519
0.0704
0.0343
0.0576
0.0099
0.0518
Wikipedia (en)
0.2970
0.2970
0.2049
0.2126
0.0847
0.0101
0.1014
0.2577
DM Mathematics
0,0003
0.0003
0.0056
0.0026
0.0177
0.0018
0.0011
0.0008
Github
0.0096
0.0096
0.0290
0.0298
0.0034
0.0538
0.0500
0.0138
PhilPaperss
0,0018
0.0018
0.0093
0.0025
0.0118
0.0005
0.0333
0.0401
Stack Exchange
0.0464
0.0464
0.0661
0.0585
0.0698
0.0430
0.1199
0.0262
Enron Emails
0.0000
0.0000
0.0009
0.0000
0.0018
0.0000
0.0000
0.0000
Gutenberg (PG-19)
0.0217
0.0217
0.0484
0.0370
0.0467
0.0223
0.0007
0.0252
Pi
0,4338
0.4338
0.3191
0.3814
0,5370
0.6323
0.5546
0,4704
Ubuntu IRC
0.0022
0.0022
0.0063
0.0072
1019
0.0123
0.0161
0.0069
EuroParl
0,0003
0.0003
0.0042
0.0040
0.0070
0.0037
0.0116
0.0000
HackerNews
0.0154
0.0154
0.0521
0.0199
0.0028
0.0551
0.0170
0.0673
PubMed Abstracts
0.0596
0.0596
0.0739
0.0532
0.0259
0.0102
0.0190
0,0017
USPTO Backgrounds
0.0130
0.0130
0.0549
0.0449
0.0004
0.0438
0.0201
0,0010
Table The parameters of different models we used in Section 4.1 and Section 4.2. When calculating model parameters, we not take into account the embedding layer and the language model head layer.
Parameter
Text Generation
Downstream Task
83M
1.6B
106M
290M
595M
1B
Hidden Size
768
2,048
768
1,280
1.536
2,048
Hidden Size
2,048
5,504
2,048
3,392
4.096
5,440
Num of Layers
12
24
15
15
21
21
Num of Heads
12
16
12
10
12
32
Max Seq Length
4,096
4,096
4,096
4,096
4,096
4,096
Vocab Size
128.256
128,256
151,936
151,936
151.936
151,936
RoPE Base
10,000
10,000
10,000
10,000
10,000
10,000
17

---

Domain2Vec: Vectorizing Datasets to Find the Optimal Data Mixture without Training
Table Downstream Task Performance of different data mixture on 106M Model. Similar to Liu et al. (2024), Human refers the original data mixture from The Pile. Pile-CC Only refers only training the Pile-CC subset. The data mixture and estimated flops of DoReMi and RegMix are from Liu et al. (2024).
Benchmark
Human
DoReMi
Pile-CC Only
RegMix
DOMAIN2VEC + DOMAIN2VEC RegMix
106M Model Pretrained 28 Tokens
Social IQA
0.340
0.349
0.353
0.356
0.339
0,342
HellaSwag
0.268
0.268
0.269
0.269
0.266
PiQA
0.573
0.584
0,580
0.586
0.579
0,583
OpenBookQA
0.245
0.251
0,249
0.242
0.245
0.249
Lambada
0.065
0.099
0.102
0.091
0.091
0,090
SciQ
0.550
0.520
0.509
0.537
0.549
ARC Fasy
0.329
0.339
0,335
0.337
0.334
0,331
COPA
0.525
0.570
0.572
0.585
0.578
0,557
RACE
0.236
0.254
0,246
0,251
0.240
0,244
LogiQA
0.282
0.280
0.271
0.268
WinoGrande
0.516
0.516
0.502
0.508
0.506
0,499
MultiRC
0.539
0.520
0.515
0.533
0.541
Average Performance
0.372
0.379
0.375
0.381
0.378
0.376
290M Model Pretrained 6B Tokens
Social IQA
0.364
0.373
0.374
0.371
0.371
0.368
HellaSwag
0.295
0.312
0.317
0.315
0.307
0,312
PIQA
0.605
0.639
0.624
0,633
OpenBookQA
0.261
0,271
0.268
Lambada
0.175
0.208
0,206
0,210
0.182
SciQ
0.711
0.682
0.663
0.674
0.670
0.697
ARC Easy
0.395
0.410
0.419
0.420
COPA
0.632
0.660
0.682
0.627
RACE
0.265
0.280
0,280
0.276
0.283
0,281
LogiQA
0.283
0.293
0,296
0.276
0.277
0,292
WinoGrande
0.511
0.506
0,509
0.524
0.498
0,504
MultiRC
0.507
0.555
0.513
0.545
0.521
0.517
Average Performance
0.417
0.432
0.431
0.431
0.421
595M Model Pretrained on 12B Tokens
Social IQA
0.378
0.387
0.390
0.394
0.383
HcllaSwag
0.338
0.377
0,386
0,385
0.355
PiQA
0.624
0.656
0.663
0.651
OpenBookQA
0.273
0.279
0,283
0.288
0,271
Lambada
0.255
0.294
0,332
0.310
0.269
0,292
SciQ
0.777
0.757
0.770
0.791
0.763
0.769
ARC Easy
0.439
0.453
0.478
0.453
COPA
0.642
0.680
0.672
0.668
RACE
0.289
0.309
0.311
0.288
0,303
LogiQA
0.263
0.268
0.252
0.267
0.263
0.267
WinoGrande
0.509
0.515
0,506
0.512
0,503
MultiRC
0.516
0.533
0.522
0.507
0.506
0.527
Average Performance
0.442
0.459
0.464
0.465
0.450
0,456
IB Model Pretrained on 20B Tokens
Social IQA
0.387
0.411
0.406
0.406
0.394
0,401
HellaSwag
0.375
0.427
0.431
0.436
0.410
0,410
PiQA
0.658
0.684
0.693
0.684
0.680
OpenBookQA
0.278
0.298
0.300
0.304
0.299
0,302
Lambada
0.301
0.359
0,348
0.353
0.334
0,339
SeiQ
0.802
0.822
0.809
0.828
0.821
0.818
ARC Easy
0.482
0.508
0.512
0.518
0,500
0.499
CОPA
0.683
0.692
0.713
0.708
0.678
0.698
RACE
0.306
0.319
0.313
0.314
0.305
0,300
LogiQA
0.259
0.258
0.269
0.272
0.268
0.267
WinoGraade
0.513
0.527
0.541
0.512
0.535
0,533
MultiRC
0.523
0.504
0.510
0.530
0.529
0,548
Average Performance
0.464
0.484
0.487
0.489
0.480
0.483
3.7 x 10 (100%)
3.5 x 10
9.66 x 10
9.66 1016
Estimated FL.OPs
(9.46%)
(0.26%)
(0.26%)
18

---

Domain2Vec: Vec: Vectorizing Datasets to Find the Optimal Data Mixture without Training
5516 Hodd not alved z0 225 illos
100 Redo! rictranedd on 128 Hikors
22Rt Nidoo mictrdledd on 128 Filles
1.25 mot
10s weers
1.0
D4:10
02
C4:02 KP:
0.0
04:00 1P:10
110
330
120
1130
130
607
1:50
1150
W00
280
2007
035 Model
153
Model Fistrained 150 Theore
238 Mode Erstralted 150 Thinre
Hotel Pretrateed 152 "Soent
04:08
02
06:06
K7
04:04
01:02 0.
04:02
0.
CK 0.0
3/4
08:0.0 K1:10
04:00
K":10
26
KP:
Model Thlent
Prstriiee Thiers
Fretratted Thinrs
02
219
5.30
200
250
300
2314 Mozia
1.80 Mode: Prstramed 1281 Tokars
Mode Pratiutta 155 Tohura
Husi (h6)
155
C4:10
K7:0.
CE:18 K2:00
1.0
04:08
K"
06:08
04.08
04
04104
02
04.02
0.
04.02 02
0.
0.0
3.5
0.0
04.00 -2. 100
(20
190
(50
199
760
736
160
Modal
Pratacrea Tokurs
Nodd Pratuaraa Tabers
MAU
04:08
0P
U8.0A
04:0222
2/2 KP
04:00 - 1.
CA 2-6 XP:08
2334 Moiid sired 30he
1.08 Mudei Prstuaroa 1381 Tokers
Nudk Prstratooa 138 Talurs
Muiid Phot atred Talrro
15L
352
CE:LD
LE:ON
CE
XP
04:02
92:
120 190
K10
AND
W00
41C
14800
1200 031
R50
XPN
100
010
1120
11000
250
113
302
1002
200
Figure 8. The validation loss on different datasct of models trained using data mixture in Table 1.
19

---

Domain2Vec: Vectorizing Datasets to Find the Optimal Data Mixture without Training
Mocil Prstrained
Modal Prstrannd Tehero
E2M Model Prstruined
Tatera
1.02 Motel Proteined
Takmms
2-
C4:10 K7.00
C4:10 sh-00
-4:08-K0:00
K2:02
K7:04
04:08 K7:94
C4:04 K7:08
C4:0A
6.00
3.0
350
240
CA:10
04:10 KS
64.00
04:08-65.02
CR 0.006-00 10h D4
C6:0. Kh:
CR1 0.6 10 0.6
CR.0A K1:0
04:02-
04.00
04:00 17:10
Phetraired
Tohem
Tooera
1.50 Pretramed
CA 0.0 MP. 3/2
02 XP 2-5
04.03
CE.10 K7:00
X2.02
04:03
C4:10 K7:0
04:03
CR:1. K0:00
0900
6000
SCRO 1250
8000
61.0
0200
Figure 9. The validation loss on different dataset of models trained using data mixture in Table 1.
20


Reviews:
[{"review_id": "VpWfeCUOyq", "reviewer": "Reviewer_ZTKc", "summary": "Domain2Vec introduces a method for vectorizing datasets by decomposing them into linear combinations of Meta-Domains, which enables efficient identification of optimal dataset mixtures for LLM pretraining. They sampled and embedded texts from which predetermined clusters/labels, which they then clustered and trained the labels on with a classifier head. \nDomain2Vec matches dataset distributions by minimizing validation loss differences without the heavy computational cost of training proxy LLMs. \nThe method achieves comparable dataset mixture quality to existing approaches like DoReMi and RegMix at significantly reduced computational expense (only 0.26% of DoReMi’s cost).", "questions": "", "limitations": "", "rating": 3, "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": "Dear Revierwer ZTKc, \n\nThank you for your recognition of the value of our work, as well as for your valuable comments.  We will respond to your questions one by one.\n\n## Q1: Domain2Vec involves much more modelling primitivesand other hyperparameters, as such, demonstrating the method’s robustness is necessary.\n\nA1: In section 4.1, we use a mixture of C4 and the Knowledge Pile as the training set, and the Pile and RedPajama as the validation set. Experiments on different datasets (Sec 4.1: C4, Knowledge Pile, Sec 4.2: The Pile) demonstrate the robustness of Domain2Vec.\n\n## Q2: Specific claim of DoReMi requiring 3.7e19 FLOPS is not clearly substantiated within the reviewed paper or the original DoReMi work. Could you provide explicit calculations or further clarify the claim regarding DoReMi's and RegMix’s computational cost?\n\nA2: Of course, the estimated FLOPS of various baslines are from the Tabel 4 of RegMix paper.\n\n[1] https://arxiv.org/abs/2407.01492\n\n## Q3: However, additional analyses on dataset hyperparameter sensitivity (e.g., varying the number of Meta-Domains) are needed.\n\nA3: Theoretically, increasing the number of Meta-Domains will lead to more accurate Domain2Vec representations of pretraining datasets. Due to limited computational resources, we only experimented with the number of Meta-Domains during clustering (Figure 3). Exploring how varying the number of Meta-Domains affects both the classifier and the final performance is left for our future work.\n\n## Q4:  Distribution Alignment Assumption (or target distribution matching) are shown in prior works but not cited. It’ll be interesting to see simpler methods like the Hashed NGram method in Xie et al [2] compared as well.\n\nA4: First, we will cite [2] in our next version. \n\nSecond, we would like to clarify that although [2] and our proposed Distribution Alignment Assumption share certain similarities:\n\n1) The methods used for feature construction differ. [2] is based on Hashed N-gram Features, whereas Domain2Vec uses a Meta-Domain Classifier to generate Domain Vectors.\n\n2) [2] conducts data selection at the sample level, whereas Domain2Vec performs data mixing at the dataset level across different datasets.\n\n3) [2] is validated on encoder-only language models such as BERT and RoBERTa, while Domain2Vec is validated on autoregressive, decoder-only language models.\n\n[2] https://arxiv.org/pdf/2302.03169\n\n## Q5: Minor presentation improvements, i.e. enlarging plot axes and legends for clarity (Figures 3, 4, 5, and 6).\n\nA5: Thanks for your suggestions, and we will increase the font size in all the figures.", "claims_and_evidence": "The paper claims substantial computational savings compared to baseline methods (DoReMi and RegMix), with evidence supported by experiments demonstrating Domain2Vec's ability to find optimal data mixtures.\n\nHowever, there are two unaddressed claims:\n- DoReMi demonstrated its efficacy on two datasets (The Pile and The Glam), while this paper only focused on The Pile. Domain2Vec involves much more modelling primitives (an embedding model, training a classifier head, etc) and other hyperparameters, as such, demonstrating the method’s robustness is necessary.\n- Specific claim of DoReMi requiring 3.7e19 FLOPS is not clearly substantiated within the reviewed paper or the original DoReMi work.", "methods_and_evaluation_criteria": "Domain2Vec's method involves embedding datasets into a \"Meta-Domain\" space, clustering embeddings, and training a classifier head on top of a small LM to produce domain vectors. Evaluation is performed primarily on the Pile dataset, comparing Domain2Vec’s performance against baselines through validation loss and downstream task accuracy. The evaluation is thorough and well-done, but limited to a single dataset (The Pile).", "theoretical_claims": "None", "experimental_designs_or_analyses": "The paper conducts extensive experiments comparing Domain2Vec against established baselines, showing that Domain2Vec achieves similar downstream task performance at dramatically lower computational cost.\nThe baselines of DoReMi and RegMix are well thought and carried out.\nHowever, additional analyses on dataset hyperparameter sensitivity (e.g., varying the number of Meta-Domains and how that would affect the classifier and the end performance) are needed.", "relation_to_broader_scientific_literature": "The current established works on Data Mixtures are DoReMi, RegMix, Data Mixture Laws (Ye et al). All of these papers require training and using smaller proxy models to inform data mixture for larger models. Domain2Vec effectively bypasses the need to train any small proxy model, by leveraging an embedding model instead. Thus enabling computational savings. \nFurthermore, all the above works do not generalize to new datasets, as they’ll require retraining of the small proxy models, while Domain2Vec works out of the box with its trained classifier. As such, if proven to be robust, this work has a meaningful impact on the field of pretraining data mixtures.", "essential_references_not_discussed": "Distribution Alignment Assumption (or target distribution matching) are shown in prior works but not cited. It’ll be interesting to see simpler methods like the Hashed NGram method in Xie et al [1] compared as well. \n\n[1] Sang Michael Xie and Shibani Santurkar and Tengyu Ma and Percy Liang. Data Selection for Language Models via Importance Resampling. https://arxiv.org/abs/2302.03169\n\n[2] Suchin Gururangan and Ana Marasović and Swabha Swayamdipta and Kyle Lo and Iz Beltagy and Doug Downey and Noah A. Smith. Don't Stop Pretraining: Adapt Language Models to Domains and Tasks. https://arxiv.org/abs/2004.10964", "other_strengths_and_weaknesses": "Strengths: \n\nDomain2Vec significantly reduces computational overhead; strong comparative experiments; clear methodological explanations. I believe it would be a helpful addition to the data curation/mixture/filtering stack, given the transferability of the classifier and the simplicity of the methodology.\n\nWeaknesses: \n\nThe complexity of the embedding and classification steps could benefit from additional robustness demonstrations. Limited hyperparameter ablation (260 meta domains). Validation was performed primarily on one dataset.", "other_comments_or_suggestions": "Minor presentation improvements, i.e. enlarging plot axes and legends for clarity (Figures 3, 4, 5, and 6).", "questions_for_authors": "1. Could you provide explicit calculations or further clarify the claim regarding DoReMi's and RegMix’s computational cost (3.7e19 FLOPS)?", "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": "", "extra_scores": {}}, {"review_id": "TWEGvcSx0s", "reviewer": "Reviewer_HtfL", "summary": "Authors propose a sampling method for training LLMs on multiple sources. Their core idea is as follows: They construct a universal set of real-valued vectors from a large textual corpus--using K-means and doc embeddings. Each vector approximately represents a topical domain of the corpus. Then they take random samples from the training sources of LLM and use a classifier to determine on average which vectors are most similar to the sampled documents. They use these similarity scores to represent the entire multi-source training dataset. They follow the same procedure for the validation set as well. Then they learn a coefficient set that transforms the representation of the training set into the validation set. The obtained weight set is the desired sampling ratio.\n\nConceptually, their idea is similar to what is used in topic-modeling, LSI, or LDA, but the vector entries here are words. So they extract a set of vectors from train and validation sets that best represent these two sets. Then they try to learn a set of coefficients that makes the train matrix similar to the validation matrix. Their argument is that if the model works on validation set, then it will work on unseen sets (the real test set).", "questions": "", "limitations": "", "rating": 1, "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": "Deal Reviewer HtfL, \n\nThanks for your careful review! We will reply to your questions one by one and hope to solve your concern.\n\n## Part1: Clarification of Our Method\n\n### Q1: Line 88 (right column). what does this mean: (where each element vj of v represents the projection (weight) of the dataset D on Dj*). what does it mean to have a weight of a dataset on Dj*? \n\nA1: First, our proposed Domain2Vec can transform a dataset into a domain vector $v$. Next, we propose that any dataset can be viewed as a **linear combination** of the domain vectors for multiple Meta Domains $\\mathcal D_j^*$. And the weight of this linear combination is $v$. Therefore, each element $v_j$ of $v$ represents the projection (weight) of the dataset $\\mathcal D$ on  $\\mathcal D_j^*$.\n\n\n### Q2: Line 143 (left col), what is \"domain vector\"? is it something that represents a document or a dataset? \n\nA2: First, we treat the domain vector as a vector representation of datasets.  \n\nFor a given dataset, we first sample N documents from it.\n\nFor each document $text_j$, $\\mathcal{p}_i$ represents the probability that $text_j$ originates from the i-th Meta-Domain. \nIt should be noted that in our paper, for each document $text_j$, $\\mathcal{p}_i$ is also referred to as the domain vector, since a single text can be viewed as a dataset with a sample size of 1.\n\nTherefore, the domain vector of the given dataset is obtained by averaging the domain vectors of documents sampled from it.\n\n### Q3: The role of validation set is to only validate your ML model, not to use it in the learning algorithm. \n\nA3: First, **in previous studies on data mixture [1][2][3][4], it is common practice to use a validation set to optimize the data mixture, which should not be regarded as overfitting.** For example, Data Mixing Law, RegMix, D-CPT-Law, and BiMix all model the relationship between data mixture and validation loss to identify the optimal mixture that minimizes validation loss.\n\nSecond, the idea of distribution matching has also appeared in previous works[5]. For instance, [5] improves the performance of continued pretraining by selecting a subset of a large raw unlabeled dataset to match a desired target distribution given unlabeled target samples.\n\nThird, our testing is independent of the validation set.\n1) Section 4.1 only evaluates whether Domain2Vec can accurately predict the model’s validation loss.\n2) When evaluating downstream task performance, we use downstream tasks that are entirely independent of the validation set, and our experimental setup follows the approach of RegMix[2].\n\nLast, compared with other baselines, our method significantly reduces computational overhead. Moreover, the data mixture obtained by our method achieves clear performance improvements on downstream tasks compared to the original mixture of The Pile dataset. Therefore, the improvement brought by our method is not \"almost zero.\"\n\n[1] https://openreview.net/pdf?id=jjCB27TMK3\n\n[2] https://openreview.net/forum?id=5BjQOUXq7i\n\n[3] https://openreview.net/pdf?id=JzKFN5fWOk\n\n[4] https://openreview.net/forum?id=JsM46OZix7\n\n[5] https://arxiv.org/pdf/2302.03169\n\n\n## Part2: Comments Suggestions And Typos\n\n### Q4: Line 136 (left column), how did you train the classifier, what is the training data? \n\nA4: First, in lines 158 (left column) through 125 (right column), we have provided a detailed explanation of the meta domain classifier training. Next, we will swap the order of the sections in lines 141–157 (left column) with those in lines 158 (left column) through 125 (right column) to make this part clearer.\n\n### Q5: After the Key assumption please explain what the reader is expecting to see. You suddenly jump into explaining that you are collecting data, the reader would wonder what you would need the data for. \n\nA5: In line 91 (right column), we will add a corresponding transitional sentence for better clarity.\n\n### Q6: Line 156 (left com), when you define V_train the inner vectors should be transposed. \n\nA6: Thanks for your suggestions. We will fix this typo error.\n\n### Q7: Lines 159 and 162 (left col), why is there meta-domain once capitalized and once is not? is there and difference between the two? \n\nA7: There is no difference here. We will capitalize \"meta-domain\" in line 162.\n\n### Q8: Line 160 (left col), what does it mean when you say a \"text belongs to a meta-domain\". As far as I know the verb \"belong\" is used in the set theory and it means membership.\n\nA8: Perhaps \"originate from\" would be a better expression. We will carefully review the usage of certain phrases.", "claims_and_evidence": "Please see the section below", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Yes", "experimental_designs_or_analyses": "Most of them", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "**Strengths:**\n\n-  The topic is timely.\n-  The method is intuitive.\n-  Then analysis is convincing.\n\n**Weaknesses:**\n\n-  Reading Section 2 (which is the core part of the paper) was very painful. I had to read it at least three times, and for some parts of it even more. Here are some of the issues:\n\n   - Line 88 (right column). what does this mean: (where each element vj of v represents the projection (weight) of the dataset D on Dj\\*). what does it mean to have a weight of a dataset on Dj\\*? After the Key assumption please explain what the reader is expecting to see. You suddenly jump into explaining that you are collecting data, the reader would wonder what you would need the data for. Line 136 (left column), how did you train the classifier, what is the training data? Please re-organize the section and put the training explanation before saying how you would use the classifier. Line 143 (left col), what is \"domain vector\"? is it something that represents a document or a dataset? On Line 143 you are using it for a doc, but on line 82 (right col) you are using it for a dataset. Line 156 (left com), when you define V_train the inner vectors should be transposed. Lines 159 and 162 (left col), why is there meta-domain once capitalized and once is not? is there ant difference between the two? Line 160 (left col), what does it mean when you say a \"text belongs to a meta-domain\". As far as I know the verb \"belong\" is used in the set theory and it means membership.\n\n   As you can see the explanations are very vague, the organization of the section is messy, and  the wording is not scientific. \n\n-  The main reason that I oppose accepting this paper is this: the core idea of the authors to do the sampling such that the sampled documents from the training set become similar to the sampled documents from the validation set is a text-book example of overfitting. The role of validation set is to only validate your ML model, not to use it in the learning algorithm. The goal of ML is to develop a model to have a low error rate in an UNSEEN set, and the validation set can be used as an unseen set during the learning stage. But you are using the validation set inside the learning model itself. This is why your improvement is almost zero, and the algorithm is not generalizable.", "other_comments_or_suggestions": "None", "questions_for_authors": "None", "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": "", "extra_scores": {}}, {"review_id": "BKkpWoXW89", "reviewer": "Reviewer_H5YS", "summary": "This paper presents a method for determining the optimal data mixture weights for combining different pre-training datasets to train language models. The authors formulate this as an optimization problem, where the goal is to find the appropriate weights over a set of meta-domains. These meta-domains are constructed by applying K-means clustering to dataset embeddings. A meta-domain classifier is then trained to predict the probability of a dataset belonging to each meta-domain.\nBy representing both the training datasets and validation datasets as a linear combination of these meta-domains, the authors ensure that all datasets exist within the same representational space. This enables them to optimize the dataset mixture weights to minimize validation loss effectively.\nThrough extensive experiments, the authors demonstrate that their approach achieves performance comparable to prior methods such as DoReMi and RegMix while operating at a significantly lower computational cost.", "questions": "", "limitations": "", "rating": 2, "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": "Dear Reviewer H5YS,\n\nThanks for your insightful review and suggestions! We will respond to your questions one by one. \n\n## Q1: The paper should provide a more thorough discussion of its connections to prior related work, i.e, Task2Vec [1].  \n\nA1: Similar to [2] and [3] cited in lines 407–412 (right column), Task2Vec [1] is an efficient way to represent a task or its corresponding dataset as a fixed-dimensional vector.\n\nHowever, Domain2Vec differs from these works in both purpose and implementation, as we focus on data mixture for language model pretraining rather than using a method like Task2Vec to select an expert from a collection (which can improve test performance while adding only minimal overhead to the training process). \n \nLast, we will also include a citation for Task2Vec.\n\n[1] https://arxiv.org/pdf/1902.03545\n\n[2] https://arxiv.org/abs/1905.11063\n\n[3] https://arxiv.org/abs/2406.00281\n\n## Q2: Dataset classification and dataset biases are well-recognized as challenging problems. The authors should elaborate on how their work relates to recent studies tackling these issues, such as the one presented in [4].\n\nA2: We would like to clarify that the target of Domain2Vec is to determine which Meta-Domains a given dataset can be composed of. \n\nIn practice, when applying Domain2Vec, we are actually performing a text classification task rather than classifying an entire dataset. As shown in Figure 6, different datasets can share the same Meta-Domain, which explains why they mutually benefit from training with each other. \n\nLast, We will also add references to [4] and discuss it in future work.\n\n[4] https://arxiv.org/pdf/2403.08632.\n\n## Q3: The paper states that the features used for clustering are derived from the \"bge-small\" model, whereas the meta-domain classifier is trained using Qwen. The rationale behind this decision is not immediately apparent, and the authors should justify why these different models were chosen for these respective tasks. \n\nA3: There are few reasons we choose Qwen as the backbone rather than the bge model.\n\n1) Since the embedding dimensions and semantic feature spaces of the Chinese and English bge models are inconsistent, and no embedding model was used for code data.\n\n2) The Embedding models like bge have a context window limited to 512 tokens, while pre-training data is typically longer than 512 tokens. In contrast, our Meta-Domain Classifier, trained based on Qwen, can handle context lengths of up to 8k tokens or even longer.\n\n3) More importantly, our Meta-Domain classifier can output very specific probability distribution on the meta-domains while the baseline can only output a hard assignment (0 or 1). There are indeed some kNN algorithms that can output soft scores, but the score is indirectly based on the distance to the center of clusters. The efficiency of kNN is also limited because it is a lazy learning algorithm and puts time complexity into inference time.\n\n## Q4：The K-means clustering approach used in the paper does not inherently ensure that clusters are not sufficiently independent as also observed in author’s experiments. To mitigate this, it would have been beneficial to introduce a diversity-enhancing objective within K-means and this could be implemented using Faiss.\n\nA4: First, this is a 260-class classification task, and our meta-domain classifier achieved an accuracy of 74.73% on the validation set. Thus, we believe that the K-means clustering approach effectively ensures clear distinctions among clusters.\n\nSecond, we have utilized FAISS to perform K-means clustering. In the future, we plan to explore introducing a diversity-enhancing objective within the K-means clustering process.\n\n## Q5: A more critical concern relates to the use of linear regression on probabilities, as presented in Equation (6). Probabilities do not exist in a regular Euclidean space but instead lie on a simplex, making standard linear regression an inappropriate model choice. A more suitable approach would be geodesic regression and I would like the authors to comment on this.\n\nA6: Great question! We want to clarify that the domain vector is not merely the probabilities assigned by the Meta-Domain Classifier to a dataset. Its deeper implication is that once we obtain the domain vector of a dataset, the dataset can be regarded as a **linear combination** of various Meta-Domains represented by the domain vector. Therefore, we can apply algorithms such as RegMix (Equation (6)) on Domain2Vec. In future work, we also plan to explore methods like geodesic regression.\n\n## Q7: While I did not verify all numerical claims, I strongly recommend that the authors carefully re-evaluate their calculations to ensure accuracy. \n\nA7: Thank you for your careful review. This is indeed a typo error. We have also verified other numerical claims throughout the paper, and this particular typo does not affect the conclusions presented.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Not applicable", "experimental_designs_or_analyses": "Yes", "relation_to_broader_scientific_literature": "This paper's contribution is relevant to the scientific literature in NLP but lacks discussion on other modalities", "essential_references_not_discussed": "There are a couple, and I note them in the weaknesses sections", "other_strengths_and_weaknesses": "Strengths:\n\nThe paper is easy to follow and the idea to breakdown dataset into a set of “building blocks” called meta-domains is interesting. The results are also impressive. \n\nThe paper should provide a more thorough discussion of its connections to prior related work. Specifically, Task2Vec (https://arxiv.org/pdf/1902.03545) introduces a framework for representing relationships between tasks, which conceptually aligns with the current study. The fundamental principle of meta-learning relies on identifying similar tasks and leveraging that information for training. Consequently, numerous studies have explored this, including those utilizing gradient similarity, such as the work presented in https://arxiv.org/pdf/1911.10600. \nMoreover, dataset classification and dataset biases are well-recognized as challenging problems. The authors should elaborate on how their work relates to recent studies tackling these issues, such as the one presented in https://arxiv.org/pdf/2403.08632. Given these papers it is crucial to situate the current findings within the broader research landscape and highlight their contributions in relation to these prior efforts.\nOne more aspect that requires clarification is the discrepancy between the feature extraction model and the meta-domain classifier. The paper states that the features used for clustering are derived from the \"bge-small\" model, whereas the meta-domain classifier is trained using Qwen. The rationale behind this decision is not immediately apparent, and the authors should justify why these different models were chosen for these respective tasks. A clear explanation would strengthen the coherence of the methodological choices.\nAnother issue pertains to the clustering methodology. The K-means clustering approach used in the paper does not inherently ensure that clusters are not sufficiently independent as also observed in author’s experiments. To mitigate this, it would have been beneficial to introduce a diversity-enhancing objective within K-means and this could be implemented using Faiss, \n\nA more critical concern relates to the use of linear regression on probabilities, as presented in Equation (6). Probabilities do not exist in a regular Euclidean space but instead lie on a simplex, making standard linear regression an inappropriate model choice. A more suitable approach would be geodesic regression and I would like the authors to comment on this. \nAdditionally, there appears to be a discrepancy in the reported calculations in line 377, where the paper states that \"Pile-CC only shows a 4.01% improvement over Human.\" However, based on the values presented in Table 3, the correct computation appears to be: [(0.439-0.424)/0.439]*100=3.5% \nrather than 4.01%. While I did not verify all numerical claims, I strongly recommend that the authors carefully re-evaluate their calculations to ensure accuracy. If my computation is incorrect, clarification would be helpful.\nOverall, while the paper introduces interesting ideas, the identified technical inconsistencies significantly weaken its contributions. Addressing these concerns would considerably improve the rigor and credibility of the study. In its current form, I am unable to recommend acceptance.", "other_comments_or_suggestions": "NA", "questions_for_authors": "Please see weaknesses", "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": "", "extra_scores": {}}, {"review_id": "Xo8ZPEWkx5", "reviewer": "Reviewer_75ct", "summary": "The document introduces DOMAIN2VEC, a technique for optimizing data mixtures in training large language models by decomposing datasets into linear combinations of \"Meta-Domains\" to enable efficient identification of optimal data mixture ratios. \nDOMAIN2VEC uses a meta-Domain classifier to classify any dataset and the Distribution Alignment Assumption (DA2) which suggests that the validation loss is low if the training and validation set are more aligned.", "questions": "", "limitations": "", "rating": 4, "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": "Dear Reviewer 75ct,\n\nThanks for your very valuable review and recognition of our work! We will address your questions point by point.\n\n## Q1 :Why did you use K-means for finding the domains? Why did you choose to represent these domains with 240 Meta-domains?   Also, it is unclear whether these clusters which resulted from a specific dataset, can represent well other datasets.\n\nA1: First, the 5.2TB text data we used to construct the Meta-Domains is unlabeled. We assume there exist distinct characteristics among different Meta-Domains, such as semantic features. In our implementation, we computed the embedding representations of these 5.2TB data to construct the Meta-Domains. Utilizing K-Means clustering on embeddings is an efficient approach under unsupervised conditions.\n\nSecond, we referred to the Elbow method, selecting the number of Meta-Domains based on the point where inertia changes relatively gradually, as shown in Figure 1. Meanwhile, the chosen number of Meta-Domains in this paper is merely an experimental setting.\n\nFinally, the meta-domain classifier achieved an accuracy of 74.73% on the validation set. Given that this is a 260-class classification task, we believe that setting K = 240 effectively ensures clear distinctions among different clusters.\n\n## Q2: The authors use KNN as an embedding-based baseline without providing further details on that. The number of Nearest Neighbours K can have a significant effect on the data classification. Moreover, it seems that KNN performs better than DOMAIN2VEC + RegMix according to table 2.\n\nA2: The details on the KNN baseline are as follows:\n\n```\nFirst, for the training and validation datasets in section 4.1, we sampled 1000 examples from each dataset.\n\nThen, we used bge-small-v1.5 (since the datasets in section 4.1 are in English) to obtain embeddings for samples from each dataset and used mean pooling to get unique embeddings for each dataset.\n\nMeanwhile, we also used bge-small-v1.5 to obtain embeddings for the data in each Meta-Domain.\n\nThen, we set K as 1000 and used KNN (based on Euclidean distance) to obtain probability distributions of training and test datasets belonging to each Meta-Domain.\n\nLast, we treated these probability distributions as new domain vectors. Based on these domain vectors, we implemented the Distribution Alignment Assumption.\n```\n\nSecond, we believe that the superior performance of kNN validates the rationality of our Meta-Domain construction process. **However, it should be noted that DOMAIN2VEC + DA² still significantly outperforms KNN**.\n\nLast, we would like to clarify that: Embedding-based methods, in addition to having more limited context length compared to our method, still have several disadvantages. For different types of data, we used different clustering methods to construct Meta-Domains.\n\n1. For code, we directly identified its programming language without using an embedding model.\n\n2. For Chinese and English data, we used embedding models that output embeddings of different dimensions. Moreover, the semantic meaning of the same dimensions of different models' embeddings is obviously different.\n\nTherefore, our proposed method has better generalizability。\n\n## Q3: However, according to Table 3, the experimental results do not seem to significantly improve performance using DOMAIN2VEC.\n\nA3: We want to clarify that Domain2Vec could provide a universal representation of pre-training datasets, which focuses on the scalability and efficiency of pre-training data mixture experiments. \n\nThe scalability of Domain2Vec is reflected in: Domain2Vec establishes a latent space representation of datasets. Therefore, any pre-training data can be mapped into this space. \n\nExperiments conducted in the latent space remain consistent regardless of changes in pre-training datasets. In contrast, RegMix performs experiments at the dataset level, requiring all previous experimental results to be discarded and new experiments to be conducted when datasets change (such as, adding new datasets, improving the quality of some datasets).", "claims_and_evidence": "The authors have employed a K-means clustering resulting in 240 different Meta-Domain clusters for English and Chinese Data. Similarly, for code data, they classified code using 20 classes based on their programming language. They claim that they can decompose datasets according to these 260 Meta-Domains. However, it is unclear, especially for the English and Chinese text if this number of clusters is the appropriate one and the reason they chose K=240. Also, it is unclear whether these clusters which resulted from a specific dataset, can represent well other datasets.", "methods_and_evaluation_criteria": "The authors use KNN as an embedding-based baseline without providing further details on that. The number of Nearest Neighbours K can have a significant effect on the data classification. Moreover, it seems that KNN performs better than DOMAIN2VEC + RegMix according to table 1.", "theoretical_claims": "All theoretical claims have been checked and seem valid.", "experimental_designs_or_analyses": "The experimental design seems valid. However, according to Table 3, the experimental results do not seem to significantly improve performance using DOMAIN2VEC.", "relation_to_broader_scientific_literature": "Previous works use a proxy model or require resampling of the data mixtures.", "essential_references_not_discussed": "To the best of my knowledge, there are no essential references that are not discussed.", "other_strengths_and_weaknesses": "Strengths\nThe paper is well-written and contains extensive experimental results. \nIt touches on an important topic related to training Large Language Models. \nIt seems to improve computational cost with respect to previous state-of-the-art.\n\nWeaknesses: \nThe experimental results do not demonstrate significant performance improvement.", "other_comments_or_suggestions": "No other comments.", "questions_for_authors": "Why did you use K-means for finding the domains? Why did you choose to represent these domains with 240 Meta-domains?", "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": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes an approach for efficiently identifying optimal dataset mixtures for LLM pretraining by vectorizing datasets through decomposition into meta-domains. This method achieves comparable dataset mixture quality to existing approaches like DoReMi and RegMix but with a significantly reduced computational expense. With more rigorous analysis, this work can add to the literature on understanding datasets.

A significant concern regarding the methodology revolves around the use of the validation set in the learning algorithm for distribution matching. Reviewer HtfL strongly argues that using the validation set to influence the training process by aligning its distribution with the training set is a textbook example of overfitting. It's worth noting that compared to traditional ML, here the validation set's domain vector is being used as a target to align the training data's domain vector. I do not view this issue major enough to warrant rejection, especially given the widespread use of val set reuse and prior empirical evidence that it does not lead to overfitting (e.g., ImageNet V2). However, given the concerns raised, I urge the authors to provide a thorough analysis of potential overfitting effects.
Decision: Accept (poster)