Domain2Vec: Vectorización de Conjuntos de Datos para Encontrar la Mezcla Óptima de Datos sin Entrenamiento

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

June 12, 2025

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

cs.AI

Resumen

Presentamos Domain2Vec, un enfoque novedoso que descompone cualquier conjunto de datos en una combinación lineal de varios meta-dominios, un nuevo concepto diseñado para capturar las características subyacentes clave de los conjuntos de datos. Domain2Vec mantiene un vocabulario de meta-dominios y utiliza un clasificador para descomponer cualquier conjunto de datos dado en un vector de dominio que corresponde a una distribución sobre este vocabulario. Estos vectores de dominio permiten la identificación de la mezcla óptima de datos para el preentrenamiento de modelos de lenguaje (LM) de manera libre de entrenamiento bajo la **Suposición de Alineación de Distribución** (DA²), la cual sugiere que cuando las distribuciones de datos del conjunto de entrenamiento y del conjunto de validación están mejor alineadas, se logra una pérdida de validación más baja. Además, Domain2Vec puede integrarse sin problemas en trabajos previos para modelar la relación entre los vectores de dominio y el rendimiento del LM, mejorando significativamente la eficiencia y escalabilidad de los métodos anteriores. Experimentos extensivos demuestran que Domain2Vec ayuda a encontrar la mezcla de datos que mejora el rendimiento en tareas posteriores con un mínimo sobrecarga computacional. Específicamente, Domain2Vec logra la misma pérdida de validación en Pile-CC utilizando solo el 51.5% del cómputo requerido al entrenar con la mezcla original del conjunto de datos The Pile. Bajo un presupuesto computacional equivalente, Domain2Vec mejora el rendimiento en tareas posteriores en un promedio 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%.