摘要

本文提出了一种名为 REGMIX 的方法，用于自动识别大型语言模型预训练的高性能数据混合。该方法通过将数据混合问题表述为回归任务，并训练小型模型来预测不同数据混合的影响，从而能够高效地识别最佳混合，然后将其推广到大规模模型训练中。

原理

REGMIX 的工作原理主要包括以下四个关键步骤：

1. 生成随机数据混合并训练小型代理模型：使用基于狄利克雷分布的方法生成各种数据混合，并使用这些混合来训练小型代理模型。
2. 拟合回归模型：使用训练好的小型代理模型的性能和其对应的混合作为输入，拟合一个回归模型，以预测其他数据混合的性能。
3. 模拟和预测：使用训练好的回归模型，对整个可能的数据混合空间进行模拟和预测，以快速找到产生最佳目标值的输入。
4. 大规模模型训练：使用模拟得到的最佳数据混合，对大规模模型进行训练。

流程

具体的工作流程如下：

1. 首先，根据数据的 token 分布，使用狄利克雷分布生成各种稀疏和近乎均匀的数据混合。
2. 然后，使用这些数据混合训练小型代理模型，并在训练过程中评估这些模型在不同领域或基准上的性能，得到目标值。
3. 接下来，使用这些目标值和数据混合作为输入，拟合一个回归模型，以预测其他数据混合的性能。
4. 最后，使用训练好的回归模型，对整个可能的数据混合空间进行模拟和预测，找到产生最佳目标值的输入，并使用该输入对大规模模型进行训练。

应用

REGMIX 可以应用于大型语言模型的预训练，以提高模型的性能和效率。此外，该方法还可以应用于其他领域，如数据挖掘、机器学习等，以优化数据混合和提高模型的性能。