Domain2Vec: Vektorisierung von Datensätzen zur Ermittlung der optimalen Datenmischung ohne Training

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

June 12, 2025

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

cs.AI

Zusammenfassung

Wir stellen~Domain2Vec vor, einen neuartigen Ansatz, der jeden Datensatz in eine lineare Kombination mehrerer Meta-Domänen zerlegt, ein neues Konzept, das entwickelt wurde, um die wesentlichen zugrunde liegenden Merkmale von Datensätzen zu erfassen. Domain2Vec verwaltet ein Vokabular von Meta-Domänen und verwendet einen Klassifikator, um einen gegebenen Datensatz in einen Domänenvektor zu zerlegen, der einer Verteilung über dieses Vokabular entspricht. Diese Domänenvektoren ermöglichen die Identifizierung der optimalen Datenmischung für das Vortraining von Sprachmodellen (LM) auf trainingsfreie Weise unter der \textbf{Verteilungsausrichtungsannahme} (DA^{2}), die besagt, dass bei besserer Ausrichtung der Datenverteilungen des Trainings- und Validierungssatzes ein geringerer Validierungsverlust erreicht wird. Darüber hinaus kann Domain2Vec nahtlos in frühere Arbeiten integriert werden, um die Beziehung zwischen Domänenvektoren und der Leistung von Sprachmodellen zu modellieren, wodurch die Effizienz und Skalierbarkeit früherer Methoden erheblich verbessert wird. Umfangreiche Experimente zeigen, dass Domain2Vec die Datenmischung findet, die die Leistung bei nachgelagerten Aufgaben mit minimalem Rechenaufwand verbessert. Insbesondere erreicht Domain2Vec denselben Validierungsverlust auf Pile-CC mit nur 51,5 % des Rechenaufwands, der beim Training auf der ursprünglichen Mischung des Pile-Datensatzes erforderlich ist. Bei gleichem Rechenbudget verbessert Domain2Vec die nachgelagerte Leistung im Durchschnitt um 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%.