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

TL;DR

Domain2Vec is a new method that decomposes datasets into meta-domains to identify optimal data mixtures for language model pretraining without additional training, improving efficiency and performance.

Contribution

It introduces the concept of meta-domains and a domain vector decomposition approach for training-free data mixture selection in language model pretraining.

Findings

Achieves same validation loss with 48.5% less computation on Pile-CC.

Improves downstream task performance by 2.83% on average under same compute budget.

Enhances scalability of data mixture optimization methods.

Abstract

We introduce~\textsc{Domain2Vec}, a novel approach that decomposes any dataset into a linear combination of several \emph{meta-domains}, a new concept designed to capture the key underlying features of datasets. \textsc{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 the optimal data mixture for language model (LM) pretraining in a training-free manner under the \emph{\textbf{D}istribution \textbf{A}lignment \textbf{A}ssumption} (DA$^{2}$), which suggests that when the data distributions of the training set and the validation set are better aligned, a lower validation loss is achieved. Moreover, \textsc{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 \textsc{Domain2Vec} helps find the data mixture that enhances downstream task performance with minimal computational overhead. Specifically, \textsc{Domain2Vec} achieves the same validation loss on Pile-CC using only $51.5\%$ of the computation required when training on the original mixture of The Pile dataset. Under equivalent compute budget, \textsc{Domain2Vec} improves downstream performance by an average of $2.83\%$.