Domain Reweighting with Minimax Optimization 简介

Domain Reweighting with Minimax Optimization是由Google与斯坦福大学研究人员共同提出的一种通过更改训练数据混合比例的方法提供大语言模型训练速度和性能的方法。

大语言模型的训练通常都是通过从多个不同领域数据进行抽样进行的。以Pile数据集为例，它包含24%的web数据集、9%的维基百科数据集、4%的GitHub数据集等。但是，这样的混合比例不一定对LLM的训练是最优的，不同领域的训练数据集的比例不同也会影响模型的训练效果。

当前大多数方法都采用直觉或者根据下游任务情况选择数据集比例。DoReMi方法则是通过提出一个小规模的代理模型先从数据集上训练，寻找最优的数据混合比例，然后再将这个比例用于大模型的训练。

作者做了一个实验，用2.8亿参数的代理模型做数据混合比例的寻优测试，再将该比例数据用于训练一个80亿参数规模的大模型，结果发现，训练速度提升2.6倍，效果提升6.5%。

DoReMi模型的特点、相关参数和训练细节如下：

特点：

1. DoReMi模型是一种用于训练语言模型的数据域重新加权算法。
2. DoReMi能够在小模型上运行，并将优势转移到30倍大的模型上，仅通过改变域上的采样概率，就能在Pile上实现2.6倍的训练加速。
3. DoReMi在不同的模型规模上都能提高下游任务的准确性。

相关参数：

1. DoReMi的域重新加权步骤涉及到T次迭代模型更新，其中T=500。
2. 在训练过程中，使用的小批量大小为1，训练样本数为500，域权重更新率为0.5。
3. 在简单的例子上，DoReMi返回的域权重为[0.39, 0.61, 0.0]，这些权重对应我们的直觉：第一个域（非噪声）增加了一小部分，第三个域（噪声）减少到0权重，大部分权重分配给第二个域。

训练细节：

1. DoReMi使用一些初始域权重来训练参考模型，然后使用组分布鲁棒优化（Group DRO）在域上训练一个小的代理模型，我们将其适应为输出域权重而不是鲁棒模型。
2. DoReMi使用这些调整的域权重来训练一个大模型。
3. 在训练过程中，对于来自域z的样本x，我们会增加当前域权重对应的伪计数。
4. 我们使用一个固定的、独立的评估集来计算算法1中的每个域的过度对数困惑度，而不是使用小批量中的样本（因为它们的大小只有1，不能代表所有域）。
5. DoReMi在Pile和GLaM数据集上进行了实验，这两个数据集分别有22个和8个域。
6. 所有模型都使用512的批量大小和1024的最大令牌长度进行训练，所有模型都是从头开始训练的。
7. 我们使用保留的验证数据来测量每个域的困惑度，对于下游评估，我们使用了GPT-3论文中的生成一次性任务。