Domain2Vec: Vectoriseren van Datasets om de Optimale Data Mix te Vinden zonder Training

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

June 12, 2025

Auteurs: Mozhi Zhang, Howe Tissue, Lu Wang, Xipeng Qiu

cs.AI

Samenvatting

We introduceren~Domain2Vec, een nieuwe aanpak die elke dataset ontbindt in een lineaire combinatie van verschillende meta-domeinen, een nieuw concept ontworpen om de belangrijkste onderliggende kenmerken van datasets vast te leggen. Domain2Vec onderhoudt een vocabulaire van meta-domeinen en gebruikt een classifier om elke gegeven dataset te ontbinden in een domeinvector die overeenkomt met een verdeling over dit vocabulaire. Deze domeinvectoren maken het mogelijk om de optimale datamix voor het vooraf trainen van taalmodel (LM) te identificeren op een trainingsvrije manier onder de \textbf{Distributie-Uitlijningsaanname} (DA^{2}), die suggereert dat wanneer de datadistributies van de trainingsset en de validatieset beter zijn uitgelijnd, een lagere validatieverlies wordt bereikt. Bovendien kan Domain2Vec naadloos worden geïntegreerd in eerdere werken om de relatie tussen domeinvectoren en LM-prestaties te modelleren, waardoor de efficiëntie en schaalbaarheid van eerdere methoden aanzienlijk worden verbeterd. Uitgebreide experimenten tonen aan dat Domain2Vec helpt bij het vinden van de datamix die de prestaties van downstream taken verbetert met minimale rekenkundige overhead. Specifiek behaalt Domain2Vec hetzelfde validatieverlies op Pile-CC met slechts 51,5% van de benodigde berekeningen bij het trainen op de originele mix van de Pile-dataset. Onder een gelijkwaardig rekenbudget verbetert Domain2Vec de downstream prestaties gemiddeld met 2,83%.

English

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