调整数据权重，性能提升6.5%，速度提升2.6倍--Intelligent technology investment network-上海优钧机械设备科技

News Message

调整数据权重，性能提升6.5%，速度提升2.6倍

by wittx 2023-12-09

DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining：DoReMi：优化数据混合加速语言模型预训练
这篇文章方法属于大模型数据工程类的文章，其通过调整大模型训练语料的各部分数据的权重，在Pile数据集合上，效果提升6.5%，达到baseline精度的时间提上2.6倍。fuyao大佬在[2]中也分享了这篇文章，推荐给大家参考下～

一、概述

Title：DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining

论文地址：https://arxiv.org/pdf/2305.10429.pdf

1 Motivation

现有的方法主要依赖于直觉或者依靠下游任务来确定每一类数据的权重，但这些方法可能导致次优的结果。
本文的motivation（动机）在于找到一种方法，能够在不依赖特定下游任务的情况下，优化预训练语言模型的数据混合比例，从而提高模型在各种任务上的性能。

2 Methods

本文提出DoReMi算法，来优化每组数据的权重配比，从而提升模型的最终表现。该算法主要包括以下步骤：

【训练一个小的参考模型】：使用一组参考域权重（每组不同的数据设立一个初始权重）训练一个小的代理模型（例如280M参数）。
【利用DRO算法训练一个代理模型去拿到每个领域数据的权重】：使用这个代理模型来引导一个更小的模型（例如280M参数），在多个域上进行分布鲁棒优化（Distributionally Robust Optimization, DRO）。这里使用了Group DRO优化器，它根据每个域的损失动态更新域权重，从而调整训练目标。DRO优化的目标是最小化最坏情况的额外损失（相对于参考模型的损失）。
【利用新的权重训练一个更大的模型】：使用在DRO训练过程中得到的域权重来训练一个更大的模型（例如8B参数）。

DRO算法的详细步骤：

初始化代理模型的权重（θ0）和域权重（α0）。在实验中，作者将域权重初始化为相等的值（1/k1，其中k是域的数量）。
对于每个训练步骤t（从1到T）：
a. 从多个域的样本中随机抽取一个minibatch B，包含大小为b的样本。
b. 计算每个域i的额外损失（λt[i]）。额外损失是代理模型在示例x上的损失（ℓθt−1,j(x)）与参考模型在示例x上的损失（ℓref,j(x)）之差，经过归一化处理。额外损失用于衡量代理模型相对于参考模型在每个示例上的改进空间。
c. 更新域权重（α′t）。使用指数函数根据额外损失更新域权重：α′t = αt−1 * exp(ηλt)。这会使得具有较高额外损失的域获得更大的权重，从而在后续训练中获得更多关注。
d. 对更新后的域权重进行归一化和平滑处理：αt = (1 - c) α′t + c u，其中u是一个全1向量，c是一个平滑参数（例如c = 1e-3）。这一步确保域权重之和为1，同时避免了权重剧烈变化。e. 使用标准优化器（如Adam或Adafactor）根据目标函数L(θt−1, αt)更新代理模型的权重θt。目标函数是代理模型在特定域权重下的损失。
在所有训练步骤结束后，返回平均域权重（1/T * ΣTt=1 αt）。这些权重将用于训练更大的模型。

通过以上步骤，DRO优化方法能够动态地调整域权重，使得代理模型在训练过程中关注那些具有较大改进空间的域。这有助于找到一组能够在各种任务上表现良好的域权重，从而提高预训练语言模型的性能。

3 Conclusion

DoReMi能够在不依赖下游任务的情况下，优化预训练语言模型的数据混合比例，从而提高模型在各种任务上的性能。
实验结果表明，DoReMi在The Pile和GLaM数据集上都能显著提高模型的困惑度和下游任务的准确性。
此外，DoReMi还能够在不知道下游任务的情况下，找到与基于下游任务调整的域权重相媲美的权重。

二、详细内容

1 实验设置

Pile数据集：一个 800GB 的文本数据集，包含 22 个域。默认域权重是启发式确定的。这里使用 The Pile 数据集中的默认域权重来训练基线，并作为 DoReMi 中的参考域权重。

GLaM数据：集包含来自 8 个域的文本。为了进行比较，GLaM 域权重（下游调整）根据在每个域上训练的模型的下游性能和每个域的大小进行调整。我们使用统一的域权重来训练 DoReMi 的基线和参考域权重。

2 下游任务实验表现

Pile数据集表现：下游任务提升6.5%个点，达到baseline精度的速度提升2.6倍。

GLaM数据集表现：

DoReMi算法在不知道下游任务的情况下，找到了一组域权重，使得预训练语言模型在下游任务上的表现与使用基于下游任务调整的域权重相当。这表明DoReMi能够在不依赖特定下游任务的情况下，找到一组具有较好泛化能力的域权重。
与其他数据集的对比：在The Pile数据集上表现更好，可能是因为GLaM数据集的域数量较少（8个域），相较于The Pile数据集（22个域），数据的多样性较低，因此DoReMi在这种情况下的优势不如在The Pile数据集上明显。

3 所有领域的困惑度都下降

发现所有任务的数据集都有下降，尽管为某些域分配了较低的权重，DoReMi 仍显著降低了所有域的基线上的困惑度。在使用DoReMi算法优化预训练语言模型的数据混合比例后，所有领域的困惑度都降低，这可能有以下几个原因：

平衡数据分布：DoReMi通过优化域权重来平衡不同领域之间的数据分布。这有助于确保模型在训练过程中关注到各种类型的数据，从而提高其在各个领域的泛化能力。通过调整数据分布，模型能够更好地捕捉不同领域之间的潜在语义关系，从而在各个领域上都取得更好的性能。
减少过拟合：通过优化域权重，DoReMi可以在一定程度上减少模型对某些领域的过拟合。过拟合可能导致模型在特定领域的训练数据上表现良好，但在其他领域的数据上泛化能力较差。通过调整数据分布，DoReMi使得模型在各个领域上都有足够的训练样本，从而降低过拟合的风险。
有效利用数据：DoReMi算法通过最小化最坏情况的损失来优化域权重，这使得模型能够在各个领域上都取得较好的性能。这种优化方法有助于确保模型在训练过程中充分利用各个领域的数据，从而提高其在各个领域的困惑度。
正向迁移：在调整数据分布时，DoReMi可能会在某些领域上分配更多的权重，从而使得模型在这些领域上获得更好的性能。这种正向迁移效应可能会对其他领域产生积极影响，使得模型在所有领域的困惑度都得到降低。
鲁棒性：DoReMi通过分布鲁棒优化（Distributionally Robust Optimization, DRO）来优化域权重，这使得模型在面对不同领域数据时具有更好的鲁棒性。DRO训练方法使得模型能够在各个领域上都取得稳定的性能，从而降低所有领域的困惑度。

总之，DoReMi算法通过优化数据混合比例，平衡数据分布，减少过拟合风险，有效利用数据，实现正向迁移和提高鲁棒性等方式，使得预训练语言模型在所有领域的困惑度都得到降低。这有助于提高模型在各种任务上的性能。

4 不同尺寸模型上都能带来提升

比较了280M、510M、760M、1B不同模型尺寸下本文DoReMi方法的影响，发现在大多数尺度上（510M除外），在20k步处有大约的3%的精度差距，同时DoReMi达到基线精度的速度也提升了4倍。

5 代理模型尺寸和优化目标的影响

Left：在The Pile数据集上，使用不同规模的代理模型（70M, 150M, 280M, 1B）进行DoReMi优化后，对下游任务性能的影响。图中的x轴表示代理模型的规模，y轴表示下游任务的平均准确率。

从70M到280M，随着代理模型规模的增加，DoReMi优化后的下游任务准确率逐渐提高。这表明，使用较大规模的代理模型进行DoReMi优化能够更好地找到适用于更大模型的域权重，从而提高下游任务的性能。当代理模型规模达到1B时，尽管代理模型本身的表现优于其他规模的模型，
但使用1B规模的代理模型进行DoReMi优化后的下游任务准确率并没有继续提高，反而略低于280M规模的代理模型。这可能说明，当代理模型规模过大时，DoReMi优化过程可能受到一些负面影响，如过拟合或计算效率降低等。总体来看，DoReMi算法在不同规模的代理模型上都表现出了对下游任务性能的提升。在实验中，280M规模的代理模型在提高下游任务准确率方面表现得最好，同时在计算效率上也相对较高。因此，作者建议在实际应用中使用280M规模的代理模型进行DoReMi优化。
Figure 6的结果表明，DoReMi算法能够在不同规模的代理模型上找到合适的域权重，从而提高预训练语言模型在下游任务上的性能。同时，选择适当规模的代理模型对于实现最佳性能和计算效率至关重要。

Right：不同数据选择策略的方法的影响（DOReMi、Hardest和Easiest）。它们的主要区别在于如何根据损失函数为数据分配权重。以下是对这三种方法的解释：

DOReMi（Domain Reweighting with Minimax Optimization）：这是本文提出的主要算法。它使用分布鲁棒优化（Distributionally Robust Optimization, DRO）来优化域权重，从而平衡不同领域之间的数据分布。DOReMi通过最小化最坏情况的损失来找到一组较好的域权重，使得预训练语言模型在各个领域上都取得较好的性能。这有助于提高模型在各种任务上的性能。
Hardest：这种方法优化的是各个领域中损失最大的样本。它试图通过关注那些对模型最具挑战性的样本来提高性能。然而，这种方法可能会导致过拟合，因为它过于关注那些难以学习的样本，而忽略了其他相对容易学习的样本。
Easiest：与Hardest方法相反，Easiest方法优化的是各个领域中损失最小的样本。它试图通过关注那些容易学习的样本来提高性能。然而，这种方法可能会导致模型在某些领域的性能较好，但在其他领域的性能较差，因为它忽略了那些具有挑战性的样本。
通过比较这三种方法，作者发现DOReMi在多个任务上的性能优于Hardest和Easiest方法。这表明，DOReMi通过平衡不同领域之间的数据分布，能够在各种任务上实现更好的性能。同时，它避免了过拟合和忽略具有挑战性样本的问题。

6 模型困惑度的变化

（a）不同尺寸代理模型下最终8B模型困惑度的变化情况：280M->8B效果看起来最好，推荐使用这个。1B可能有点过拟合了。

（b）比较了使用DoReMi方法得到的主模型（main model）与相同规模的代理模型（proxy model）的困惑度情况：

尽管1B参数规模的代理模型的质量相对较差，但通过这个代理模型得到的域权重（domain weights）仍然能够改善主模型的性能。
这里强调了DoReMi方法的鲁棒性，即使在代理模型质量不佳的情况下，它仍然能够找到一组有效的域权重来优化主模型的性能。这表明DoReMi方法在实际应用中具有较好的稳定性和可靠性。