Domain2Vec : Vectorisation des jeux de données pour identifier le mélange de données optimal sans entraînement

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

Résumé

Nous présentons~Domain2Vec, une approche novatrice qui décompose tout ensemble de données en une combinaison linéaire de plusieurs méta-domaines, un nouveau concept conçu pour capturer les caractéristiques sous-jacentes clés des ensembles de données. Domain2Vec maintient un vocabulaire de méta-domaines et utilise un classifieur pour décomposer tout ensemble de données donné en un vecteur de domaine correspondant à une distribution sur ce vocabulaire. Ces vecteurs de domaine permettent l'identification du mélange de données optimal pour le pré-entraînement de modèles de langage (LM) de manière sans entraînement sous l'\textbf{Hypothèse d'Alignement des Distributions} (DA^{2}), qui suggère que lorsque les distributions de données de l'ensemble d'entraînement et de l'ensemble de validation sont mieux alignées, une perte de validation plus faible est obtenue. De plus, Domain2Vec peut être intégré de manière transparente dans les travaux précédents pour modéliser la relation entre les vecteurs de domaine et la performance des LM, améliorant ainsi considérablement l'efficacité et l'évolutivité des méthodes antérieures. Des expériences approfondies démontrent que Domain2Vec aide à trouver le mélange de données qui améliore la performance des tâches en aval avec un surcoût computationnel minimal. Plus précisément, Domain2Vec atteint la même perte de validation sur Pile-CC en utilisant seulement 51,5% du calcul requis lors de l'entraînement sur le mélange original de l'ensemble de données The Pile. Avec un budget de calcul équivalent, Domain2Vec améliore la performance en aval en moyenne de 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%.