Domain2Vec: Vectorizing Datasets to Find the Optimal Data Mixture without Training

Dear Reviewer tRxw,

Thanks for your insightful review! We will respond to your questions one by one.

Q1: How do we construct the vocabulary of meta-domains given a set of sources?

In the "Construct the Vocabulary of Domain2Vec" paragraph in Section 2 (lines 127-146), we explain how we construct several Meta-Domains from our collected pretraining datasets. Specifically, we used K-means to cluster on our collected text data of over 5TB size, where texts belonging to the same cluster form a Meta-Domain. Finally, we obtained 260 Meta-Domains. We use the term, "vocabulary" to denote a set composed of Meta-Domains, as an analogy of Word2Vec.

Q2: How do we choose N, i.e., the number of points needed to get a confident probability vector？

Theoretically, a larger N would be better. In this paper, we sampled 1,000 points from each validation set to calculate the domain vector. Our experimental results show that 1,000 points are sufficient for good performance. In practice, one can determine the sample size by adding samples until the calculated domain vector becomes very stable.

Q3: When training the meta-domain classifier, is an arbitrary text belonging to a single meta-domain or can the arbitrary text also belong to multiple meta-domains? Is the meta-domain classifier a multi-class or multi-label problem? How do you handle precision/recall over long-tail meta-domains?

1. During the training phase, we train the meta-domain classifier with a single-label problem. The classifier predicts which meta-domain a text belongs to. During inference time, the meta-domain classifier outputs a probability distribution over meta-domains, rather than a one-hot vector.
2. When training the Meta-Domain classifier, we use the same sample size for different Meta-Domains, which helps mitigate possible class imbalance problem. Moreover, there is no real long-tail meta domains. even the smallest Meta-Domain contains approximately 2GB texts as shown in the table below.

Q4: Rather than fitting a meta-domain classifier, a simple baseline would be to leverage learned feature embeddings to automatically cluster the data into a set of meta-domains. How well would this perform versus the meta-domain vocabulary and the classifier?

Great question! Technically, your mentioned baseline is totally feasible. However, we show the reasons for our practice from the following aspects:

1. 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.

2. 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 $k$-Nearest Neighbors algorithms that can output soft scores, but the score is indirectly based on the distance to the center of clusters. The efficiency of $k$-NN is also limited because it is a lazy learning algorithm and puts time complexity into inference time.

Q5: Ultimately, $DA^2$ is trying to match two categorical probability vectors. Why use a Huber loss rather than, e.g., distributional measures (divergence, Wasserstein distance, etc.)?

Another great question! Technically, Huber loss combines the advantages of L1 and L2 distance. We add the results of different distributional measures based on Table 2 in our paper. As shown in the table, Huber Loss shows better performance than L1/L2/JS Distance.

Additionally, Wasserstein distance is a very great option, to be honest. However, it would require an extra metric space matrix, $\boldsymbol{M}$, to measure the distance between two vectors. In this work, the metric space $\boldsymbol{M} (260 \times 260)$ is actually the "dataset transition cost" between each two Meta-Domains, and is non-trivial. Each element in $\boldsymbol{M}$, $c_{ij}$ could be estimated via $L_{ij}$, the loss at Meta-Domain $j$ after training 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 a positive result (even better) if the metric space is well estimated, and we leave this for future work.

Dear Reviewer 9xLo,

Thanks for your very valuable review and recognition of our work! We will address your questions point by point.

Q1: The authors should discuss the relationship and differences between their work and these studies (e.g., [1], [2]).

Thank you very much for your specific suggestions! Dataset2Vec, XTFORMER, and our proposed Domain2Vec indeed share some similar ideas, establishing a latent space to model features in a domain. Thus, we supplemented the citations and discussions of prior studies (line 520 - line 523) to our paper and modified the corresponding statements. However, We believe that Domain2Vec is significantly different from the mentioned related works based on the following aspects.

1. The research questions and motivations are different and our paper mainly focuses on data mixture of LLMs pre-training.
2. The specific implementation of Domain2Vec is completely different from prior works (as we state in detail in section 2).
3. By obtaining representations of pre-training datasets in the same latent space, we can not only align datasets in the latent space (Domain2Vec +$DA^2$) but also model the relationship between latent space representations and model performance (Domain2Vec + RegMix).

[1] Cross-Table Pretraining Towards a Universal Function Space for Heterogeneous Tabular Data.

[2] Dataset2Vec: Learning Dataset Meta-Features.

Q2: What if the domain vector of the training datasets is a subset of the domain vector of the validation datasets, or vice versa?

We do not well understand this question and we will discuss this issue from two aspects.

1. A domain vector is a vector, rather than a set, thus it should be not possibly a "subset" of another domain vector.
2. Assuming your question is, what would the domain vector of the training dataset be if training dataset is a subset of the validation dataset? Our answer is: their domain vectors would be very similar (i.e., $\boldsymbol{\mathcal{v}}_1 \cdot \boldsymbol{\mathcal{v}}_2 \approx 1$) if domain 1 is randomly sampled (thus a subset) from domain 2 because the goal of Domain2Vec is just for calculating the features of a domain. When the sampling number becomes larger, $\boldsymbol{\mathcal{v}}_1$ becomes of less variance and finally stable as $\boldsymbol{\mathcal{v}}_2$.

Thanks for your time to read our response! We have updated our paper and fixed some typos. And we look forward to a further discussion with you.

Best regard,

Authors of Paper 9648

Dear Reviewer 9xLo,

This is a kindly reminder that the discussion period for our submission will end on Nov 26, 2024 (AOE). We have carefully addressed all concerns and questions raised in your review.

Please let us know if there’s anything more we can clarify or elaborate on.

Thank you so much for your insightful review and recognition of our work!

Best regards,

Authors of Paper 9648

W1: The notation is not clearly defined, for example, the dimensionalities, what is a vector and what is a scalar, can only be guessed from the equations rather than being clearly defined.

A1: In this paper, we strictly follow the recommended notation provided in the ICLR 2025 template.

(see page 4 in https://www.overleaf.com/latex/templates/template-for-iclr-2025-conference-submission/gqzkdyycxtvt.pdf)

For example, in our paper, $p$ is a scalar, $\boldsymbol{p}$ is a vector, and $\boldsymbol{V}_{train}$ is a Matrix.

If you have any more questions about the notation somewhere, could you please point out which specific notation?

W2: Please consider clearly defining what is a vocabulary, features, and semantic features. Is a feature a vector? Is it an abstract concept? Is vocabulary a set of features? What is a semantic feature?

We are drawing an analogy between Domain2Vec and Word2Vec to help people in relevant fields better understand.

Vocabulary: a set composed of Meta-Domains.

(Semantic) Features(of a domain): a probability distribution on the vocabulary, which is computed by our Domain2Vec classifier. We use the term "semantic" here because our classifier is based on semantic embeddings.

Thanks for your suggestions and we will add the definitions in our paper to help readers better understand our ideas.

W3: Issues on the dependence on well-trained reference model.

We think we have a strong response. There is a big difference between our embedding model and their reference models.

Prior works trains a reference model (which is typically a smaller-size of the final language model) to help them find more valuable domains to learn during training. When the training dataset changes, these reference models used in previous works need to be retrained.

Conversely, Domain2Vec adopts BGE and Qwen-2 models to compute domain vectors of different datasets, and they do not participate in the model training process. The parameters of the models are frozen. And the models could be viewed as a tool to compute embeddings and they can be easily replaced by any other embedding models (such as a simple FastText).

Overall, the procedure in our work does not reply such a reference model, and provides a much simpler and more efficient way to compute the optimal mixture ratio.

W4: The improvement in validation loss seems not to translate to any improvement in downstream performance. No better than “Pile-CC only” that is a baseline with zero additional cost.

We clarify this issue from the following aspects:

1. Our work does not intend to research on the relationship between validation loss and downstream performance. Instead, we followed RegMix's experimental design and baselines (improvement in validation loss). We would like to emphasize that our contributions mainly lie in efficiency and scalability: in Table 3, Domain2Vec has achieved comparable downstream performance with only 0.26% of DoReMi's computational cost.

2. In RegMix's work, it shows that validation loss on Pile-CC has the strongest correlation with the benchmarks, thus we follow and consider Pile-CC as validation set in this work. Interestingly, the better downstream performance achieved by "Pile-CC Only" actually confirms the correctness of our proposed Domain Alignment Assumption ($DA^2$).

3. Yes, "Pile-CC only” is good when we regard Pile-CC as validation set. However, in a more practical scenario where validation set is somewhat else, we cannot manually find such a golden training set which can 100% correspond to validation set. However, we can use our proposed Domain2Vec to get a comparable performance with lowest cost.

W5: For Figure 5, please consider adding more baselines such as DoReMi, a simple 50-50 mix of Pile and C4, and a baseline that deduplicates this simple mix.

1. Figure 5 aims to demonstrate that the data mixture found by Domain2Vec can reduce the loss on Pile-CC more quickly compared to the original data mixture in The Pile. In Table 3, we have indeed compared downstream performance with baseline, DoReMi. The results show that we achieved comparable downstream performance using only 0.26% of DoReMi's computational cost.

2. The experiments corresponding to Figure 5 only involve adjusting data mixture within The Pile's sub-datasets, so adding C4 doesn't seem to be a reasonable baseline.

3. Finally, data deduplication is fundamentally orthogonal to our proposed method. Typically, data deduplication and optimal mixture ratio discovery are two adjacent procedures in LLM pre-training. We believe that data deduplication can absolutely improve the performance and it does NOT conflict with our proposed method.

Thank you very much for your careful review! We have updated the paper and corrected the typos you mentioned. We look forward to receiving your response!

Best regards,

Authors of Paper 9648

Dear Reviewer me3R,

Thanks for your insightful review! Q1 – Q3 are our responses to your questions and W1 – W5 are our responses to the weaknesses you have mentioned.

Q1: How one would infer the claim on lines 172-173 from this figure that is “our classifier could reasonably distinguish some base features from different datasets.”?

We clarify this question from the two following aspects:

Firstly, in Lines 168-169, we reported that the Meta-Domain Classifier achieved a classification accuracy of 74.73% on the test set. It suggests that our classifier can distinguish the differences among the datasets and get a accurate result.

Secondly, in Figure 2, we presented the decomposition results of different datasets by the Meta-Domain Classifier. For example, FreeLaw contains a large amount of data from Meta-Domain 232, while DM Mathematics contains a large amount of data from Meta-Domain 142. This indicates that FreeLaw and DM Mathematics have very different compositions, which aligns with our prior knowledge.

Q2: My understanding of Section 3.1 is that the proposed method relies on a predefined split of the dataset into domains. Is that correct?

There might be some misunderstandings. Section 3.1 describes the basic task formulation, which is to find the optimal mixture ratio from a given list of datasets (e.g. Github + ArXiv + Wikipedia). The task formulation is independent on any methods and keeps consistent across many other previous works like [1], [2], [3].

However, Our method itself never relies on a predefined split of the dataset, we explain this question from training and inference stages:

1. In training stage, we train a Domain2Vec classifier based on the clustering result of an unsupervised method, kmeans.
2. In inference stage, we can adopt our Domain2Vec classifier to automatically identify an unseen domain, thus generate a vector of the domain.

[1] Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance

[2] RegMix: Data Mixture as Regression for Language Model Pre-training

[3] DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining

Q3: Please clarify the $DA^2$ assumptions and requirements in the paper and particularly this sentence in the abstract.

Thank you for pointing this issue out! We have updated the abstract to better clarify our assumptions. Specifically, the revised sentence becomes "... under the Distribution Alignment Assumption (DA2), which suggests that when the data distribution of the training set and the validation set is more aligned, a lower validation loss is achieved." (line 16 - line 20)

Q4: The notation of the loss in section 3.1 does not seem correct. Shouldn’t there be an additional inner argmin over given the mixing ratio.

Thanks for your careful review！We have revised the notation of the loss in section 3.1 and other related notations in the updated version of our paper.

Dear Reviewer me3R,

We are writing to kindly remind you that the discussion period for our submission is approaching its conclusion on Nov 26, 2024 (AOE). We have thoroughly responded to the comments and concerns you’ve raised.

And if there are any points you’d like to discuss further or additional input you’d like from us, we would be delighted to assist before the deadline.

We truly appreciate your time and effort in reviewing our work!

Best regards,

Authors of Paper 9648

Dear Reviewer s3iG,

Thanks for your careful review! We will reply to your questions one by one and hope to solve your concern.

Q1: The proposed method seems to offer little to no gains empirically.

1. Domain2Vec achieved comparable downstream performance with only 0.26% of DoReMi's computational cost. Extremely high efficiency while maintaining good performance is Domain2Vec's advantage. This has been acknowledged by Reviewers me3R, 9xlo, and tRxw.

Reviewer me3R: The proposed method is ~3x more efficient than prior work, DoReMi. The method can efficiently adapt and scale as the number of domains within a dataset increases over time.

Reviewer 9xlo: Experimental results indicate that DOMAIN2VEC can achieve better text generation capabilities and downstream task performance with lower computational costs.

Reviewer tRxw: The results demonstrate that the proposed method is effective.

2. Our method's efficiency becomes particularly significant when the training dataset changes and scales up. For example, when we add a new training dataset, RegMix needs to give up everything, and re-collect approximately $|\mathcal{D}|^2$ ($|\mathcal{D}|$ denotes the number of datasets) data mixture to fit the linear regression model. In contrast, Domain2Vec only needs to use the Meta-Domain classifier to convert the new dataset into a vector representation. Note that in practical scenarios, the number of datasets could be significantly more than 20 (e.g. 200). It's clear that RegMix's computational cost increases quadratically with the number of datasets, while Domain2Vec increases linearly.

3. Our method provides a novel perspective on optimizing data mixture, and Domain2Vec can naturally integrate with prior work (e.g., Domain2Vec + RegMix), this further suggest the scalability of our proposed method.

Q2: The paper does not compare to state-of-the-art pre-training datasets such as RefinedWeb, FineWeb, and DCLM.

The main research scope of our work is to find the optimal data mixture among multiple pretraining datasets, rather than constructing a better pretraining dataset. Our proposed method is orthogonal to dataset sources. For example, Domain2Vec can also adjust the datat mixture of different domains within RefinedWeb. In this paper, to align with prior works, such as DoReMi, RegMix, we primarily selected The Pile, C4, and Knowledge-Pile as our pre-training datasets.

We appreciate your thorough review and insightful questions. We hope we address each point you mentioned in detail. And we look forward to your further feedback.

Best regard,

Authors of Paper 9648

Dear Reviewer s3iG,

Thank you again for your time and effort in reviewing our work! As a reminder, the discussion period for our submission will conclude on Nov 26, 2024 (AOE). We have carefully addressed the questions and concerns raised in your review and sincerely look forward to your valuable feedback.

If you have any additional questions or require further clarification, we would be more than happy to engage and provide prompt responses during this period.

Best regards,

Authors of Paper 9648

Dear Reviewer s3iG,

Thanks for your feedback of our response.

Q1: I agree that Domain2Vec is more efficient than DoReMi, but is this the most relevant comparison? The results in the submission show that RegMix is more efficient than DoReMi. When comparing to prior work, the comparison to the best prior work is particularly relevant.

In Table 3, we also compared with RegMix. Correspondingly, Domain2Vec's computational cost is only 2.76% of RegMix's. Additionally, as we mentioned in our initial response to your review, "RegMix's computational cost increases quadratically with the number of datasets, while Domain2Vec increases linearly. This demonstrates that our method is more efficient than RegMix. It is also notable that, while RegMix requires training language models, Domain2Vec only needs to perform inference. This gives Domain2Vec significant advantages in terms of code implementation complexity and GPU memory consumption. Please refer to the following Q2 for more explanations and comparisons with our method and RegMix.

Q2: As I mentioned in my original review, it is not clear to me that the computational advantage over RegMix is relevant because the computational cost, at least based on flops count, is more than 10x smaller than training the model. Is this correct? If so, the training mix optimization with RegMix is not the bottleneck here.

No, it is not correct. Taking training a 1B parameter model for 25B tokens as an example (following RegMix), the Flops required to train this model is approximately 1.5e+20 (N = 1B, D=25B). The following table shows how the FLOPs of Domain2Vec and RegMix change with the number of training datasets. When the number of training datasets reaches 200, RegMix's FLOPs already exceed those required for training the model itself.

Note that over 1K training datasets in LLM training is quite possible. For example, your mentioned Fineweb[1] has over 10K dumps with inconsistent distributions. Another example is Qwen-2.5 [2] and Llama-3 [3], which are trained on over 15 trillion tokens—approximately 100 times the size of The Pile dataset.

Our method provides a universal representation of pre-training datasets, which focuses on the scalability and efficiency of pre-training data mixture experiments. The 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. Experiments 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 a new datasets, improving the quality of some datasets).

Last but not least, our method does not replace RegMix but can enhance its scalability and efficiency using our Domain2Vec. Please refer to the "Domain2Vec + RegMix" section in our paper.

[1] The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale

[2] Qwen2.5: A Party of Foundation Models

[3] The Llama 3 Herd of Models

Q3: If I understand the results correctly, at the 1B scale (Table 3), Domain2Vec offers no advantage in aggregate model performance over training on a single data source only (Pile-CC). Is this correct? Training on Pile-CC alone also requires no cost for optimizing the training mix. It would be helpful to see results in a regime where Domain2Vec results in a better model.

We would like to clarify this issue from the following aspects:

1. Our work does not intend to research on the relationship between validation loss and downstream performance. Therefore, we followed RegMix's experimental design and baselines.
2. RegMix finds that validation loss on Pile-CC has the strongest correlation with the benchmarks they evaluates. Thus, we follow and consider Pile-CC as validation set in our experiments. Interestingly, the better downstream performance achieved by "Pile-CC Only" actually confirms the correctness of our proposed Domain Alignment Assumption ($DA^2$).
3. Most importantly, "Pile-CC only” is good when we regard Pile-CC as validation set. However, in a more practical scenario where validation set is somewhat else, we cannot manually find such a golden training set which can 100% correspond to validation set.. However, we can use our proposed Domain2Vec to get a comparable downstream performance with lowest cost by mixing datasets from different sources.

Lastly, thank you for understanding that our method and datasets are orthogonal. Due to computational resource constraints, and to ensure fair comparison with prior works, we chose to use the Pile. However, in Table 1, we also conduct experiments on the Knowledge Pile (A recent high quality pretraining dataset). We consider experiments on higher-quality datasets as our future work.

We hope that our response can address your concern, and we look forward to your further feedback.

Best regards,

Authors of Paper 9648

Dear Reviewer s3iG,

We are very grateful for the opportunity to improve our work through your valuable feedback. As the discussion period is near the end, this is a kind reminder for your further feedback.

If you think our response has addressed your concerns and questions, we sincerely hope you might consider raising your score. Should you have any further inquiries or require additional clarification, we warmly welcome your questions and are eager to provide comprehensive responses.

Best regards,

Authors of Paper 9648

Dear Reviewer s3iG,

We are very grateful for the opportunity to improve our work through your valuable feedback. As the discussion period is near the end, this is a kind reminder for your further feedback.

If you think our response has addressed your concerns and questions, we sincerely hope you might consider raising your score. Should you have any further inquiries or require additional clarification, we warmly welcome your questions and are eager to provide comprehensive responses.

Best regards,

Authors of Paper 9648

Dear Reviewer s3iG,

We greatly appreciate the opportunity to enhance our work based on your invaluable feedback. As the discussion period draws to a close, this is a kind reminder for your further feedback.

If our response has addressed your concerns and questions, we sincerely hope you might consider raising your score. Should you have any further inquiries or require additional clarification, we warmly welcome your questions and are eager to provide further responses.

Warm regards,

Authors of Paper 9648