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은 메타 도메인의 어휘를 유지하며, 주어진 데이터셋을 이 어휘에 대한 분포에 해당하는 도메인 벡터로 분해하기 위해 분류기를 사용한다. 이러한 도메인 벡터는 \textbf{분포 정렬 가정}(DA^{2}) 하에서 학습 없이도 언어 모델(LM) 사전 학습을 위한 최적의 데이터 혼합을 식별할 수 있게 해준다. 이 가정은 훈련 세트와 검증 세트의 데이터 분포가 더 잘 정렬될수록 더 낮은 검증 손실이 달성된다는 것을 제안한다. 또한, Domain2Vec은 이전 연구에 원활하게 통합되어 도메인 벡터와 LM 성능 간의 관계를 모델링함으로써 이전 방법의 효율성과 확장성을 크게 향상시킬 수 있다. 광범위한 실험을 통해 Domain2Vec이 최소한의 계산 오버헤드로 하위 작업 성능을 향상시키는 데이터 혼합을 찾는 데 도움이 된다는 것을 입증하였다. 구체적으로, Domain2Vec은 Pile 데이터셋의 원래 혼합으로 훈련할 때 필요한 계산량의 51.5%만 사용하여 Pile-CC에서 동일한 검증 손실을 달성한다. 동일한 계산 예산 하에서 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%.