Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance

📄 arXiv: 2403.16952v2 📥 PDF

作者: Jiasheng Ye, Peiju Liu, Tianxiang Sun, Jun Zhan, Yunhua Zhou, Xipeng Qiu

分类: cs.CL, cs.AI, cs.LG

发布日期: 2024-03-25 (更新: 2025-03-20)

备注: accepted by ICLR2025, camera ready version


💡 一句话要点

提出数据混合法则以优化语言模型的数据混合比例

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 数据混合法则 语言模型 性能预测 训练优化 机器学习

📋 核心要点

  1. 现有方法依赖启发式或定性策略来调整数据混合比例,缺乏定量预测能力,导致模型性能优化困难。
  2. 论文提出数据混合法则,通过拟合混合比例与模型性能的函数关系,提前预测未见混合的模型性能。
  3. 实验表明,该方法在RedPajama中优化了1B模型的训练混合,性能与训练48%更多步骤的模型相当,显示出显著提升。

📝 摘要(中文)

大型语言模型的预训练数据由多个领域(如网络文本、学术论文、代码)组成,其混合比例对模型性能有重要影响。现有方法依赖启发式或定性策略调整比例,而我们发现模型性能与混合比例之间存在定量可预测性,称为数据混合法则。通过对样本混合进行拟合,可以在实际运行前预测模型在未见混合上的性能,从而指导理想数据混合的选择。此外,我们提出嵌套使用训练步骤、模型规模和数据混合法则的缩放法,以便在小规模训练下预测在各种混合下的大型模型性能。实验结果验证了我们的方法有效优化了在RedPajama中训练的1B模型的训练混合,性能可与在默认混合下训练48%更多步骤的模型相媲美。

🔬 方法详解

问题定义:本论文旨在解决大型语言模型预训练数据混合比例对模型性能影响的优化问题。现有方法多依赖于经验和定性分析,缺乏系统的定量预测能力,导致模型性能提升的效率低下。

核心思路:论文提出了数据混合法则,认为模型性能与混合比例之间存在可预测的函数关系。通过对样本混合进行拟合,可以在实际训练前预测模型在未见混合上的性能,从而指导数据混合的选择。

技术框架:整体方法包括数据混合法则的建立、样本混合的拟合以及性能预测模块。首先,通过对不同混合比例的样本进行性能评估,建立混合比例与模型性能的函数关系;然后,利用该关系预测未见混合的性能;最后,结合训练步骤和模型规模的缩放法,优化训练过程。

关键创新:最重要的技术创新在于提出了数据混合法则,系统地将模型性能与数据混合比例的关系量化,突破了以往依赖经验的局限。这一方法使得在小规模训练下也能有效预测大型模型的性能。

关键设计:在模型训练中,关键参数包括混合比例的选择、样本数量的设置以及损失函数的设计。通过对不同混合比例的样本进行系统评估,确保拟合函数的准确性和可靠性,从而提升模型性能的预测能力。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示,采用数据混合法则优化的1B模型在RedPajama中训练100B tokens,性能与在默认混合下训练48%更多步骤的模型相当,验证了该方法的有效性和优越性。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、机器翻译和对话系统等。通过优化数据混合比例,可以显著提升大型语言模型的训练效率和性能,具有重要的实际价值。此外,研究成果为动态数据调度提供了理论基础,未来可能推动更智能的训练策略发展。

📄 摘要(原文)

Pretraining data of large language models composes multiple domains (e.g., web texts, academic papers, codes), whose mixture proportions crucially impact the competence of outcome models. While existing endeavors rely on heuristics or qualitative strategies to tune the proportions, we discover the quantitative predictability of model performance regarding the mixture proportions in function forms, which we refer to as the data mixing laws. Fitting such functions on sample mixtures unveils model performance on unseen mixtures before actual runs, thus guiding the selection of an ideal data mixture. Furthermore, we propose nested use of the scaling laws of training steps, model sizes, and our data mixing law to enable predicting the performance of large models trained on massive data under various mixtures with only small-scale training. Moreover, experimental results verify that our method effectively optimizes the training mixture of a 1B model trained for 100B tokens in RedPajama, reaching a performance comparable to the one trained for 48% more steps on the default mixture. Extending the application of data mixing laws to continual training accurately predicts the critical mixture proportion that avoids catastrophic forgetting and outlooks the potential for dynamic data schedules