Domain2Vec: Vectorizing Datasets to Find the Optimal Data Mixture
  without Training
Authors: name: Mozhi Zhang, name: Howe Tissue, name: Lu Wang, name: Xipeng Qiu

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 the optimal data mixture for language model (LM) pretraining
in a training-free manner under the \textbf{Distribution
Alignment Assumption} (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, 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 computation 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%.