Domain2Vec: Vettorizzazione di Dataset per Trovare la Miscela Ottimale di Dati senza Addestramento

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

June 12, 2025

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

cs.AI

Abstract

Introduciamo Domain2Vec, un approccio innovativo che scompone qualsiasi dataset in una combinazione lineare di diversi meta-domini, un nuovo concetto progettato per catturare le caratteristiche fondamentali sottostanti dei dataset. Domain2Vec mantiene un vocabolario di meta-domini e utilizza un classificatore per scomporre qualsiasi dataset dato in un vettore di dominio che corrisponde a una distribuzione su questo vocabolario. Questi vettori di dominio consentono l'identificazione della miscela di dati ottimale per il pre-addestramento di modelli linguistici (LM) in modo privo di addestramento, sotto l'**Assunzione di Allineamento della Distribuzione** (DA^{2}), che suggerisce che quando le distribuzioni di dati del set di addestramento e del set di validazione sono meglio allineate, si ottiene una perdita di validazione inferiore. Inoltre, Domain2Vec può essere integrato senza soluzione di continuità in lavori precedenti per modellare la relazione tra vettori di dominio e prestazioni del LM, migliorando notevolmente l'efficienza e la scalabilità dei metodi precedenti. Esperimenti estensivi dimostrano che Domain2Vec aiuta a trovare la miscela di dati che migliora le prestazioni nei task downstream con un sovraccarico computazionale minimo. Nello specifico, Domain2Vec raggiunge la stessa perdita di validazione su Pile-CC utilizzando solo il 51,5% del calcolo richiesto quando si addestra sulla miscela originale del dataset The Pile. Con un budget computazionale equivalente, Domain2Vec migliora le prestazioni downstream in media del 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%.