RegMix: Data Mixture as Regression for Language Model Pre-training

本文研究了大型语言模型预训练中数据混合对性能的影响，但如何确定有效的数据混合仍不清楚。作者提出了RegMix，将其作为回归任务来自动识别高性能的数据混合。RegMix涉及使用不同的数据混合训练一组小模型，并拟合回归模型以预测它们各自混合的性能。通过拟合的回归模型，我们模拟出排名靠前的混合，并将其用于使用更多计算的大规模模型的训练。为了验证RegMix，作者训练了512个具有不同混合的1M参数模型，以拟合回归模型并找到最佳混合。使用此混合，我们训练了一个1B参数模型，用于25B个标记（即比其他混合的64个1B参数模型大1000倍，长25倍），并发现其在性能上表现最佳。此外，我们的方法表现优于人工选择，并实现了与DoReMi相匹配或超越其结果，同时只利用了10％的计算预算。我们的实验还表明：（1）数据混合对性能有显著影响，单任务性能变化高达14.6％; （2）与被认为是高质量的数据（如维基百科）相比，Web语料库与下游性能之间具有最强的正相关性; （3）领域以复杂的方式相互作用，常常与常识相矛盾，因此需要像RegMix这样的自动方法; （4）数据混合效应超越了缩放定律，我们的方法通过考虑所有领域来捕捉复杂性。我们的代码可在https://github.com/sail-sg/regmix上找到。

论文旨在解决语言模型预训练中数据混合对性能的影响问题，提出了一种自动化的数据混合方案RegMix。

RegMix将数据混合的性能作为回归任务进行预测，通过训练一组小模型来确定一个高性能的数据混合方案，并使用该方案训练一个大规模模型。

实验结果表明，数据混合对性能影响显著，单任务性能变化高达14.6%；RegMix方法在性能和计算成本方面都优于人工选择和DoReMi方法；研究发现，与Wikipedia等高质量数据相比，Web语料库对下游性能的影响更强；研究还发现，不同领域之间的数据混合效应相互影响，需要自动化的方法。研究代码已开源。

在语言模型预训练领域，最近的相关研究包括MegatronLM、GShard、Turing-NLG等。