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

2025-06-12

AI Summary

To address the computationally expensive trial-and-error training required for optimizing data mixing ratios in language model pretraining, this paper proposes a training-free, zero-shot data mixture optimization method. The core innovation introduces the concept of “meta-domains” and domain vectors, grounded in the Distribution Alignment Assumption (DA²). Leveraging classifier-driven domain decomposition, meta-domain vocabulary construction, and zero-shot data vectorization, the method enables pre-training distribution alignment modeling without any model training. It is plug-and-play and fully compatible with existing pretraining pipelines. On the Pile-CC dataset, our method achieves the same validation loss as uniform mixing using only 51.5% of the computational budget; under equal compute constraints, it yields an average 2.83% improvement across downstream tasks.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Natural Language Processing: Learning & Optimization for NLP
• Search and Optimization: Mixed Discrete/Continuous Search

Application Category

• Web Mining and Content Analysis: Large pretrained models with web data
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions
• Search and Retrieval-Augmented AI: Vertical and domain-specific search

Abstract

We introduce~ extsc{Domain2Vec}, a novel approach that decomposes any dataset into a linear combination of several emph{meta-domains}, a new concept designed to capture the key underlying features of datasets. extsc{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 emph{ extbf{D}istribution extbf{A}lignment extbf{A}ssumption} (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, extsc{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 extsc{Domain2Vec} helps find the data mixture that enhances downstream task performance with minimal computational overhead. Specifically, extsc{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, extsc{Domain2Vec} improves downstream performance by an average of $2.83%$.

Problem

• Finds optimal data mixture for LM pretraining without training
• Decomposes datasets into meta-domains to capture key features
• Enhances downstream task performance with minimal computation

Innovation

• Decomposes datasets into meta-domains for analysis
• Uses domain vectors for training-free data mixture optimization
• Enhances efficiency and scalability of language model pretraining