用向量表示数据集,零训练找最优数据组合
Domain2Vec: Vectorizing Datasets to Find the Optimal Data Mixture without Training
- 将数据集分解为元领域向量,构建数据分布的线性表达
- 仅需51.5%计算量即可达到原数据混合的验证损失水平
- 适合大规模预训练数据优化,提升效率与可扩展性
我们提出~\textsc{Domain2Vec},一种新方法将任意数据集分解为若干\emph{元领域}的线性组合,这些元领域旨在捕捉数据集的核心特征。\textsc{Domain2Vec}维护一个元领域词汇表,并利用分类器将任意数据集分解为对应于该词汇表分布的域向量。在\emph{\textbf{D}istribution \textbf{A}lignment \textbf{A}ssumption}(DA$^{2}$)假设下,该方法无需训练即可识别语言模型预训练的最优数据混合。实验表明,\textsc{Domain2Vec}在使用\textsc{Pile-CC}数据时,仅需原\textsc{The Pile}数据混合51.5%的计算量,即能达到相同验证损失;在等量计算预算下,下游任务性能平均提升2.83%。该方法可无缝集成至已有工作,显著提升数据混合优化的效率与可扩展性。
原文摘要 · Abstract (English)
We introduce~\textsc{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. \textsc{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{\textbf{D}istribution \textbf{A}lignment \textbf{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, \textsc{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 \textsc{Domain2Vec} helps find the data mixture that enhances downstream task performance with minimal computational overhead. Specifically, \textsc{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, \textsc{Domain2Vec} improves downstream performance by an average of $2.83\%$.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。