Domain2Vec: Vetorizando Conjuntos de Dados para Encontrar a Mistura Ótima de Dados sem Treinamento

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

Resumo

Apresentamos o Domain2Vec, uma abordagem inovadora que decompõe qualquer conjunto de dados em uma combinação linear de vários meta-domínios, um novo conceito projetado para capturar as características subjacentes essenciais dos conjuntos de dados. O Domain2Vec mantém um vocabulário de meta-domínios e utiliza um classificador para decompor qualquer conjunto de dados fornecido em um vetor de domínio que corresponde a uma distribuição sobre esse vocabulário. Esses vetores de domínio permitem a identificação da mistura de dados ideal para o pré-treinamento de modelos de linguagem (LM) de maneira livre de treinamento, sob a \textbf{Hipótese de Alinhamento de Distribuição} (DA^{2}), que sugere que, quando as distribuições de dados do conjunto de treinamento e do conjunto de validação estão melhor alinhadas, uma menor perda de validação é alcançada. Além disso, o Domain2Vec pode ser integrado de forma contínua a trabalhos anteriores para modelar a relação entre vetores de domínio e o desempenho do LM, aumentando significativamente a eficiência e a escalabilidade dos métodos anteriores. Experimentos extensivos demonstram que o Domain2Vec ajuda a encontrar a mistura de dados que melhora o desempenho em tarefas subsequentes com um custo computacional mínimo. Especificamente, o Domain2Vec alcança a mesma perda de validação no Pile-CC utilizando apenas 51,5% da computação necessária ao treinar na mistura original do conjunto de dados The Pile. Sob um orçamento computacional equivalente, o Domain2Vec melhora o desempenho em tarefas subsequentes em uma média 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%.