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

Mozhi Zhang, Howe Tissue, Lu Wang, Xipeng Qiu
Proceedings of the 42nd International Conference on Machine Learning, PMLR 267:76139-76158, 2025.

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 Alignment Assumption (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$%.

Cite this Paper

BibTeX
@InProceedings{pmlr-v267-zhang25bu, title = {{D}omain2{V}ec: Vectorizing Datasets to Find the Optimal Data Mixture without Training}, author = {Zhang, Mozhi and Tissue, Howe and Wang, Lu and Qiu, Xipeng}, booktitle = {Proceedings of the 42nd International Conference on Machine Learning}, pages = {76139--76158}, year = {2025}, editor = {Singh, Aarti and Fazel, Maryam and Hsu, Daniel and Lacoste-Julien, Simon and Berkenkamp, Felix and Maharaj, Tegan and Wagstaff, Kiri and Zhu, Jerry}, volume = {267}, series = {Proceedings of Machine Learning Research}, month = {13--19 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v267/main/assets/zhang25bu/zhang25bu.pdf}, url = {https://proceedings.mlr.press/v267/zhang25bu.html}, 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 Alignment Assumption (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$%.} }
Endnote
%0 Conference Paper %T Domain2Vec: Vectorizing Datasets to Find the Optimal Data Mixture without Training %A Mozhi Zhang %A Howe Tissue %A Lu Wang %A Xipeng Qiu %B Proceedings of the 42nd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2025 %E Aarti Singh %E Maryam Fazel %E Daniel Hsu %E Simon Lacoste-Julien %E Felix Berkenkamp %E Tegan Maharaj %E Kiri Wagstaff %E Jerry Zhu %F pmlr-v267-zhang25bu %I PMLR %P 76139--76158 %U https://proceedings.mlr.press/v267/zhang25bu.html %V 267 %X 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 Alignment Assumption (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$%.
APA
Zhang, M., Tissue, H., Wang, L. & Qiu, X.. (2025). Domain2Vec: Vectorizing Datasets to Find the Optimal Data Mixture without Training. Proceedings of the 42nd International Conference on Machine Learning, in Proceedings of Machine Learning Research 267:76139-76158 Available from https://proceedings.mlr.press/v267/zhang25bu.html.

Related Material