Domain2Vec: Векторизация наборов данных для поиска оптимальной смеси данных без обучения

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

June 12, 2025

Авторы: Mozhi Zhang, Howe Tissue, Lu Wang, Xipeng Qiu

cs.AI

Аннотация

Мы представляем~Domain2Vec, новый подход, который разлагает любой набор данных на линейную комбинацию нескольких метадоменов — нового концепта, разработанного для захвата ключевых базовых характеристик наборов данных. Domain2Vec поддерживает словарь метадоменов и использует классификатор для разложения любого заданного набора данных в вектор домена, который соответствует распределению по этому словарю. Эти векторы доменов позволяют идентифицировать оптимальную смесь данных для предварительного обучения языковой модели (LM) без необходимости обучения в рамках \textbf{Предположения о Согласовании Распределений} (DA^{2}), которое предполагает, что при лучшем согласовании распределений данных обучающего и валидационного наборов достигается меньшая ошибка на валидации. Более того, Domain2Vec может быть легко интегрирован в предыдущие работы для моделирования взаимосвязи между векторами доменов и производительностью LM, значительно повышая эффективность и масштабируемость предыдущих методов. Многочисленные эксперименты демонстрируют, что Domain2Vec помогает найти смесь данных, которая улучшает производительность на последующих задачах с минимальными вычислительными затратами. В частности, Domain2Vec достигает той же ошибки на валидации на Pile-CC, используя только 51,5% вычислений, необходимых при обучении на исходной смеси данных набора The Pile. При эквивалентном вычислительном бюджете Domain2Vec улучшает производительность на последующих задачах в среднем на 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%.