DoReMi：优化数据混合加速语言模型预训练
计算与语言
2023-11-22 v4 机器学习
摘要
预训练数据域（例如维基百科、书籍、网络文本）的混合比例极大影响语言模型（LM）性能。在本文中，我们提出基于极小极大优化的域重加权（DoReMi），其首先通过对域的群分布鲁棒优化（Group DRO）训练一个小型代理模型，以在不知晓下游任务的情况下产生域权重（混合比例）。然后我们依据这些域权重对数据集重新采样，并训练一个更大的全尺寸模型。在我们的实验中，我们使用一个 280M 参数的代理模型上的 DoReMi 来为更高效地训练一个 8B 参数模型（大 30 倍）设定域权重。在 The Pile 上，DoReMi 改善了所有域的困惑度，即使当它降低某个域的权重时也是如此。DoReMi 在使用 The Pile 默认域权重训练的基线模型上，将平均少样本下游准确率提高了 6.5 个百分点，并以少 2.6 倍的训练步数达到基线准确率。在 GLaM 数据集上，DoReMi 虽不知晓下游任务，甚至匹配了使用针对下游任务调优的域权重的性能。
引用
@article{arxiv.2305.10429,
title = {DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining},
author = {Sang Michael Xie and Hieu Pham and Xuanyi Dong and Nan Du and Hanxiao Liu and Yifeng Lu and Percy Liang and Quoc V. Le and Tengyu Ma and Adams Wei Yu},
journal= {arXiv preprint arXiv:2305.10429},
year = {2023}
}
备注
NeurIPS 2023