预训练数据域（如维基百科、书籍、网络文本）的混合比例极大地影响了语言模型（LM）的性能。在本文中，我们提出了具有Minimax优化的域重新加权（DoReMi），它首先在域上使用组分布鲁棒优化（group DRO）来训练一个小型代理模型，以在不知道下游任务的情况下产生域权重（混合比例）。然后，我们用这些域权重对数据集进行重新采样，并训练一个更大的全尺寸模型。在我们的实验中，我们在280M参数代理模型上使用DoReMi来找到域权重，以便更有效地训练8B参数模型（30倍大）。在The Pile上，DoReMi改善了所有领域的困惑，即使它降低了一个领域的权重。与使用The Pile的默认域权重训练的基线模型相比，DoReMi将平均少杆下游准确度提高了6.5%，并通过减少2.6倍的训练步骤达到基线准确度。在GLaM数据集上，对下游任务一无所知的DoReMi甚至与使用对下游任务进行调整的域权重的性能相匹配。

1简介

用于训练语言模型（LM）的数据集通常从许多领域的混合物中采样（Brown等人，2020，Chowdhery等人，2022，Du等人，2021，Gao等人，2020）。例如，The Pile（Gao et al.，2020）是一个大型的公开数据集，由24%的网络数据、9%的维基百科、4%的GitHub等组成。1预训练数据的组成极大地影响了LM的有效性（Du et al.，2021，Hoffmann et al.，2022，Xie et al.，2023）。然而，目前尚不清楚每个领域要包含多少，才能生成一个在各种下游任务中表现良好的模型。现有的工作通过使用直觉或一组下游任务来确定领域权重（每个领域的采样概率）。例如，The Pile使用启发式选择的域权重，这可能是次优的。另一方面，现有的LM，如PaLM（Chowdhery et al.，2022）和GLaM（Du et al.，2021），基于一组下游任务来调整域权重，但需要在不同的域权重上训练潜在的数千个LM，并有过度适应特定下游任务集的风险。

我们的方法不是基于一组下游任务来优化域权重，而是通过最小化域上最坏情况下的损失来找到域权重，从而使模型在所有域上都表现良好。一个天真的最坏情况方法是用

1本文中基于令牌计数的域权重因令牌化器而异；见附录C。

图1：给定一个具有一组域的数据集，使用Minimax优化的域重新加权（DoReMi）优化域权重，以改进在数据集上训练的语言模型。首先，DoReMi使用一些初始参考域权重来训练参考模型（步骤1）。参考模型用于指导使用域上的组分布鲁棒优化（group DRO）来训练小型代理模型（Nemirovski等人，2009，Oren等人，2019，Sagawa等人，2020），我们将其调整为输出域权重，而不是鲁棒模型（步骤2）。然后，我们使用调整后的域权重来训练大型模型（步骤3）。

引用域

使用重新加权的数据集训练大型语言模型

大型语言模型

大多数噪声数据，因为每个域都有不同的最优损失（也称为熵）。为了使领域困惑具有可比性，我们遵循Mindermann等人（2022），Oren等人（2019），并优化最坏情况下的超额损失，即正在评估的模型和预训练的参考模型之间的损失差距。

这激发了我们的算法，即具有Minimax优化的域重新加权（DoReMi），该算法利用分布式鲁棒优化（DRO）在不了解下游任务的情况下调整域权重（图1）。首先，DoReMi以标准方式训练小型参考模型（例如280M个参数）。其次，我们训练了一个小型分布式鲁棒语言模型（DRO-LM）（Oren等人，2019），该模型在所有领域中最大限度地减少了最坏情况下的超额损失（相对于参考模型的损失）。值得注意的是，我们采用DRO训练产生的域权重，而不是使用鲁棒LM。最后，我们在由这些域权重定义的新数据集上训练大的（8B）LM。我们的方法采用DRO-LM框架（Oren等人，2019）来优化域权重，而不是生成稳健的模型。为此，我们使用了来自Group DRO的基于在线学习的优化器（Nemirovski et al.，2009，Sagawa et al.，2020），该优化器根据每个域的损失动态更新域权重，以重新调整训练目标，而不是像Mindermann et al.（2022），Oren et al.（2019）那样从迷你批次中细分示例。最后，DoReMi采用DRO训练步骤的平均域权重。

在第3节中，我们在280M代理和参考模型上运行DoReMi，以优化The Pile（Gao et al.，2020）和GLaM数据集（Du et al.，2021）（用于PaLM（Chowdhery et al.，2022））上的域权重。DoReMi域权重用于训练8B参数LM（大于30倍）。在The Pile上，DoReMi在所有域上都比基线域权重减少了困惑，即使它降低了域的权重。DoReMi在生成的少镜头任务中，将基于The Pile默认域权重训练的基线模型的平均下游精度提高了6.5%，并将基线下游精度提高2.6倍（图2）。我们发布了the Pile的调谐域权重，以改进使用the Pile训练的未来LMs（表1）。在第4节中，我们发现当改变代理模型和用优化域权重训练的主模型的大小时，DoReMi一致地改进了LM训练。在GLaM数据集上，其中域权重

图2:DoReMi用一个小模型（280M个参数）优化域权重，并使用这些域权重来训练一个更大的模型（8B个参数，30倍大）。这里，优化域权重（训练一个小模型两次）需要训练大模型的8%的计算量。DoReMi在the Pile上进行预训练时，将单杆下游的平均精度提高了6.5%，并将基线精度提高了2.6倍。

在下游任务上进行调优是可用的，DoReMi甚至在下游任务性能上的性能与调优域权重相当。

使用Minimax优化的2域重新加权（DoReMi）

在本节中，我们定义了DoReMi，这是一种使用小型代理模型来优化语言建模数据集的域权重的算法，然后改进了大型模型的训练。

设置。假设我们有k个域（例如，维基百科、GitHub），其中对于每个域i，我们有一组示例Di。域权重a e Ak指定了k个域上的概率分布，因此指定了训练数据上的分布：Pa=Pk=1 ai•unif（Di），其中unif（D）=| dD | Px<D$x是D中示例上的均匀分布，如果x=x，则8x（x）为1，否则为0。

DoReMi。DoReMi的输入是数据D1，Dk、参考域权重aref和用于大型全尺寸模型的训练超参数（训练步骤数T和批量大小b）。DoReMi返回优化的域权重a，并最终返回在Pa上训练的大型模型。

步骤1：获取一个小型参考模型。我们首先在T个步骤、批量大小b的一些参考域权重aref（例如，默认情况下为统一域权重）上训练模型pref。该模型用作步骤2的参考模型，并捕获每个示例/域的基线难度水平。参考模型可以是一个相对较小的模型（在我们的实验中有280M个参数）。

步骤2：使用Group DRO训练代理模型以获得域权重。为了获得域权重，我们在分布式鲁棒语言建模（DRO-LM）（Oren et al.，2019）框架中使用Group DRO优化器（Sagawa et al.，2020）训练一个小型代理模型p$，其中0是代理模型的权重。该框架通过优化域上的最坏情况损失来训练鲁棒模型，这相当于以下最小最大目标：

其中损失'e（x）=-log pe（x）和'ref（x）=-log pref（x）分别是本文中代理模型和参考模型的负对数似然性，|x|是示例x的令牌长度。目标是最小化跨域的最坏情况下的超额损失，因为a上的内部最大化将所有权重放在具有最高超额损失的域上。

直观地说，超额损失（'e（x）-'ref（x））测量了代理模型相对于参考模型在示例x上改进的净空。具有较高超额损失的示例是参考模型实现低损失（使得示例是“可学习的”）但代理模型仍然具有高损失的示例。具有低过量损失的示例可以是非常高的熵（即，最优损失高，因此参考损失高）或非常低的熵（例如，易于学习，因此代理损失低）。Group DRO优化器通过在训练步骤t上将域权重的指数梯度上升更新与代理模型权重Gt的梯度更新交错来工作。优化器在更新到具有高超额损失的向上八个域，这在这些域的示例上放大了代理模型的梯度更新。根据Nemirovski等人（2009），我们返回训练轨迹上的平均权重a=T PTL1，作为在步骤3中使用的优化域权重。

步骤3：用新的域权重训练大型模型。调整后的域权重a定义了一个新的训练分布Pa。我们使用标准训练程序在来自这个新分布的数据上训练一个主模型（比参考/代理模型大）。

算法1 DoReMi域重加权（步骤2）

需要：域数据D1，。。。，Dk，训练步骤的数量T，批量大小b，步骤大小n，平滑参数c e[0,1]（例如，在我们的实现中，c=1e-4）。

步骤2的详细信息。算法1为步骤2提供了伪代码。算法1的主要结构是一个训练循环，它在T个步骤上更新代理模型。在每一步中，我们都遵循Sagawa等人的方法。（2020）并对具有统一域权重的迷你批次进行采样（不考虑参考域权重aref，这只影响参考模型）。然后，我们计算每个域的超额损失，通过每个域中的令牌总数进行归一化，并在每一步使用它们来更新域权重。我们首先计算每个令牌级别的每个域的超额损失，然后进行聚合，其中索引j处的令牌级别损失为'et-1，j（x）=-log pg（xj|x1，…，xj-1，这对于Group DRO优化器的基本保证是必要的（Sagawa et al.，2020），在实践中很少需要，但这意味着我们没有严格优化最小-最大目标。最后，我们使用标准优化器更新目标L（dt-1，at）的代理模型，如Adam（Kingma和Ba，2015）或Adafactor（Shazeer和Stern，2018）。在我们的所有实验中，我们将域权重更新步长设置为n=1，将平滑参数设置为c=1e-4，并且没有广泛地调整这些超参数。

重复的DoReMi。我们通过运行DoReMi多轮来扩展它，将下一轮的参考域权重aref设置为上一轮的a。我们称之为迭代DoReMi。整个迭代过程仍然只使用小模型来调整域权重。当域权重收敛时，我们停止迭代，我们将其定义为当任何域权重的最大变化||a-arefk^小于1e-3时。根据经验，在GLaM数据集上，这只需要3轮（第3.2节）。

3 DoReMi提高LM训练效率和表现

在本节中，我们使用使用280M参数代理模型优化的DoReMi域权重来训练8B参数主模型（大30倍）。我们考虑两个数据集，The Pile（Gao et al.，2020）和GLaM训练集（Du et al.，2021）。在The Pile上，DoReMi显著减少了每个领域的困惑，将生成一次性任务的平均下游精度提高了6.5%，并将基线精度提高了2.6倍。在GLaM数据集上，在下游数据集上调整的域权重可用，DoReMi发现域权重的性能与下游调整的域权值相当。

3.1实验设置

Pile数据集。The Pile（Gao et al.，2020）是一个800GB的文本数据集，包含22个域（表1）。默认的域权重是启发式确定的。我们使用the Pile数据集中的默认域权重来训练基线，并在DoReMi中作为参考域权重（见附录C）。

GLaM数据集。GLaM数据集（Du等人，2021）（也用于训练PaLM（Chowdhery等人，2022））包括来自8个域的文本（表2）。为了进行比较，GLaM域权重（下游调整）是根据在每个域上训练的模型的下游性能和每个域的大小进行调整的（Du等人，2021）。我们使用统一的域权重来训练基线，并使用参考域权重aref来训练DoReMi。

培训设置。我们训练Transformer（Vaswani et al.，2017）解码器专用LMs，其标准下一个令牌语言建模丢失。我们通过均衡计算量来进行受控比较，计算量由训练期间处理的令牌数量来衡量。对于The Pile，我们对每个模型进行200k步的训练；对于GLaM数据集，我们对每个模型进行300k步的训练。所有型号都使用512的批量大小和1024的最大令牌长度。代理模型和参考模型具有280M个参数。所有模型都是从头开始训练的（其他超参数见附录C）。

评价我们使用持有的验证数据来衡量每个领域的困惑。对于下游评估，我们使用GPT-3论文（Brown et al.，2020）中的生成一次性任务：TriviaQA（Joshi et al.，2017）、NaturalQuestions（Kwiatkowski et al.，2019）、WebQuestions（Berant et al.，2013）、SQuADv2（Rajpurkar et al.，2018）和LAMBADA（Paperno et al.，2016）。我们对这些数据集使用标准的精确匹配精度度量。

用于优化域权重的计算。我们训练两个280M模型（参考模型和代理模型）来优化域权重。这是训练主8B模型所需FLOP的8%。所有FLOP都来自标准的向前和向后传球。

DoReMi中型号尺寸的符号。我们将参考/代理模型的大小（在我们的实验中总是相同的大小）和用DoReMi域权重训练的主模型的大小表示为“DoReMi（参考/代理的大小^主模型的尺寸）”：例如，DoReMi（280M^8B）。当我们讨论独立于主模型的优化域权重时，我们只包括一个数字（例如，DoReMi（280M）），它指的是参考/代理模型的大小。

3.2 DoReMi提高困惑和下游准确性

我们表明，相对于各自的基线域权重，DoReMi显著提高了在the Pile和GLaM数据集上训练的8B模型的困惑性和下游准确性。

The Pile的下游精度有所提高。图3（左）显示了桩上基线和DoReMi（280M^8B）模型的平均下游性能。DoReMi将下游精度提高了6.5%，并在75k步以内达到基线精度，比基线（200k步）快2.6倍。因此，DoReMi可以显著加快训练速度并提高下游性能。

（a） The Pile

图3：5项任务的平均单次下游精度（精确匹配），在The Pile（左）和GLaM数据集（右）上训练8B参数模型。在The Pile上，DoReMi将下游精度提高了6.5%，并将基线精度提高了2.6倍（与图2相同）。在GLaM数据集上，迭代的DoReMi（第2轮）获得了与下游任务调整的域权重相当的性能。

（b） GLaM数据集

图4:The Pile上8B模型的逐域对数困惑。尽管降低了一些域的权重，但DoReMi改善了所有域的日志困惑。

表1：桩上的域权重。基线域权重是根据默认的Pile数据集计算的。DoReMi（280M）使用280M代理模型来优化

DoReMi可以在没有权衡的情况下减少所有领域的困惑。图4显示了the Pile上8B模型的全域困惑。DoReMi显著减少了所有领域对基线的困惑，尽管为某些领域分配了较低的权重。这是怎么发生的？直观地说，熵最低和最高的域可以被降权，而不会对困惑产生太大影响。熵最小的域在统计上只需要很少的样本来学习。最高熵域的令牌分布接近常见的一致先验——例如，随机初始化时的模型倾向于输出一致的下一个令牌分布。因此，我们需要更少的样本来适应这些领域。从分配更多样本到中等熵域的正迁移可以改善所有域的困惑。在附录D中，我们提供了一个简单的例子，其中重新加权域可以改善所有域的困惑，DoReMi在模拟中发现了这样的域权重。

迭代的DoReMi在GLaM数据集上实现了下游调整权重的性能。我们在GLaM数据集上使用了3轮迭代的DoReMi。我们发现第二轮和第三轮域权重几乎相同（表2）。图3（右）显示了前两轮迭代DoReMi的一次性结果。在第一轮之后，DoReMi主模型的下游精度与基线相当（统一的域权重）。第二次之后

表2:GLaM数据集中的域权重。迭代DoReMi（280M）在3轮内收敛，总体模式与下游任务上调整的域权重相似。

DoReMi主模型实现了与下游调谐域权重相当的下游精度。总体而言，域重新加权对GLaM的影响较小，可能是因为与The Pile中的22个域相比，只有8个域。

检查DoReMi域权重。表1和表2给出了the Pile和GLaM数据集的DoReMi域权重。当在280M代理模型（DoReMi（280M））上运行DoReMi时，大多数权重都放在不同的Pile CC web文本域上。请注意，与基线相比，维基百科的权重有所降低，但DoReMi仍然提高了维基百科任务的下游准确性（例如，TriviaQA，附录表5）。在GLaM数据集上，DoReMi权重与下游调谐域权重具有相同的通用模式。DoReMi能够通过从统一的初始参考域权重开始恢复类似的域权重集，而不使用任何下游数据。

4消融和跨标度分析

之前在第3节中，我们展示了DoReMi使用280M模型找到域权重，这可以改进8B模型的训练。在本节中，我们对DoReMi进行了分析，其中我们改变了代理模型相对于主模型的规模，并消除了超额损失目标的组成部分。

DoReMi在不同规模上持续改进LMs。我们考虑在一个简单的设置中使用相同大小的代理和主模型来分析DoReMi的行为，而不需要域权重在尺度上进行泛化。特别是，我们在The Pile上运行DoReMi（X^X），其中X为280M、510M、760M或1B。图5显示，DoReMi在基线基础上持续提高了2%的下游精度，并在各个尺度上平均快4倍达到基线精度，而且这种提高不会随着模型尺寸的增大而缩小。DoReMi改善了所有量表上最坏情况下的困惑，并平均改善了22个个体领域困惑中的18个（附录表6）。

代理模型的性能不如主模型，尤其是在较大的规模下。回想一下，DoReMi使用Group DRO来训练代理模型，该模型用域权重重新加权目标。相反，主模型是通过对来自DoReMi的域权重进行重新采样来训练的。当代理模型和主模型大小相同时，哪一个模型更好？表3b显示，在这种情况下，代理模型通常表现不如主模型。代理模型和主模型之间的差距随着规模的扩大而增加，因为1B代理模型不仅不如1B主模型，而且不如1B基线模型，而在19/22域上，280M代理模型比280M基线模型实现了更好的困惑。尽管1B代理模型的质量相对较差，但域权重仍然允许1B主模型实现基线

图5：在4个模型尺度（280M、510M、760M、1B）上的平均单次下游精度，其中DoReMi的参考/代理模型与使用DoReMi域权重训练的主模型大小相同。DoReMi不断提高各尺度的下游精度，在大多数尺度（510M除外）的200k步长下，精度差距相似，为3%。DoReMi在各种尺度上平均以4倍的速度达到基线精度。

性能提高了2倍以上。这表明DoReMi对极小极大优化过程中的任何次优性都是相当鲁棒的。然而，我们假设代理和主模型训练之间的不匹配（损失重加权与重采样）解释了它们的性能差异，因此基于重采样的Group DRO优化器可以改进较大代理模型的DoReMi。

代理模型规模对大型主模型性能的影响。我们考虑了DoReMi代理模型的70M、150M、280M和1B比例，同时将主模型大小固定为8B（DoReMi（X^8B））。从70M到280M，增加代理模型大小可提高8B的下游精度（图6左）。我们假设1B代理模型不会继续这种趋势，因为组DRO优化器在更大范围内更差（表3b）。虽然DoReMi（280M^8B）在8B时得到了最大的改善，但DoReMi（150M^8B。这表明DoReMi对代理模型规模是稳健的。

选择最简单或最难的领域是不够的。我们通过在示例x上仅使用代理模型pe的损失来运行DoReMi，即“e（x）”（更喜欢代理模型的最硬域）或仅使用引用的负损失，来消除超额损失度量“e（x）--ref（x）”的分量

图6：在The Pile上训练的模型的平均下游精度。（左）在DoReMi中将参考/代理模型的大小从70M增加到280M可以提高8B主模型的下游精度，但1B代理模型的趋势不会继续。我们假设组DRO优化器对于较大的代理模型更差。右）针对最难或最简单的域而不是过度损失（两者结合）进行优化，并不能实现与DoReMi（280M模型）相同的平均下游精度。

表3：The Pile上每个领域的困惑汇总（共22个领域）。平均困惑是每个领域困惑的未加权平均值。

-’ref（x）（对于参考模型，更喜欢最简单的域）。图6（右）显示，单独的超额损失分量都不足以实现DoReMi的收益。

5相关工作

为LMs绘制预训练数据。最密切相关的是GLaM数据集（Du et al.，2021）（也用于训练PaLM（Chowdhery et al.，2022）），其具有使用下游数据调整的域权重。优化下游任务的域权重可能很昂贵，并且可能需要搜索/零阶优化（Snoek等人，2012）、RL（Zoph和Le，2016），或者对域之间的正/负转移如何工作的启发式假设。示例级别的过滤也为LM培训带来了好处。C4数据集（Raffel等人，2019）通过启发式数据清理方法显示了对CommonCrawl的优势。杜等人（2021），谢等人（2023）表明，在示例级别过滤数据，以获得看起来像维基百科和书籍的高质量文本，可以显著提高LMs的下游性能。与这些工作相反，DoReMi只需2次小的LM训练就可以自动设置域权重，并且不会对喜欢的数据类型做出假设（类似维基百科等）。

一般数据选择方法。Moore-Lewis选择（Axelrod，2017，Feng et al.，2022，Moore和Lewis，2010）选择在目标数据和原始数据上训练的语言模型之间具有高交叉熵差（类似于过度对数困惑）的例子。Coleman等人（2020）基于用于主动学习的小型代理模型的不确定性来选择示例。Mindermann等人（2022）通过根据超额损失在迷你批次中选取前k个示例，以在线方式选择示例。许多其他关于数据选择的工作都在设想中（Kaushal等人，2019，Killamsetty等人，2021a，b，c，Mirzasoleiman等人，2020，Paul等人，2021，Sener和Savarese，2018，Sorscher等人，2022，Wang等人，2020和Wei等人，2015）。总的来说，这些方法不涉及预训练的数据选择，其中下游数据分布可能与预训练分布非常不同。DoReMi旨在解决DRO的预训练/下游分配转变问题。

分布式鲁棒优化。在深度学习的DRO方法中（Ben Tal et al.，2013，Oren et al.，2019，Sagawa et al.，2020，Sinha et al.，2018），我们针对一种被称为组转移的受限形式的转移（Duchi et al.，2017，Oren等人，2019，萨加瓦等人，2020），其中测试分布可以是组（域）的未知混合。我们遵循DRO-LM（Oren et al.，2019），该方法在分组轮班设置中为LMs使用DRO。DRO-LM也使用基线损失，但使用简单的二元参考模型。DoReMi使用与代理模型大小和架构相同的参考模型，以确保可以比较损失。在优化过程中，DRO-LM采用每个小批量的最坏情况子集来更新模型，而我们使用不需要子选择的Group DRO优化器（Sagawa et al.，2020）。子选择是昂贵的，因为它需要在大型迷你批次上评估模型，而只选择一小部分来更新模型。总的来说，与这些旨在生成稳健模型的DRO方法相比，我们使用DRO来优化数据，以便更有效地训练更大的模型。

6讨论和限制

依赖于训练算法。理想情况下，DoReMi将独立于训练算法，但DoReMi运行训练算法来训练参考/代理模型。尽管如此，DoReMi在某些方面实现了算法独立性：在第3.2节中，我们展示了DoReMi域权重跨尺度传递增益。通常，我们预计DoReMi发现的域权重会在广泛的模型规模、计算预算和其他训练超参数之间传递。

使用外推法将计算保存在DoReMi中。在第2节中，我们运行DoReMi来计算将用于训练最终模型的训练步骤的数量，这可能是不必要的昂贵。节省计算的未来方向将是在早期步骤停止运行DoReMi，并外推所需步骤数的域权重，因为我们发现DoReMi运行期间域权重的大多数变化似乎发生在训练开始时（附录图8）。

参考模型的选择。参考模型的选择可能会影响DoReMi找到的域权重。例如，迭代DoReMi（第3节）通过使用参考模型来提高性能，该参考模型是在上一轮DoReMi的调谐域权重上训练的。进一步的方向包括改变参考模型大小和使用专门的参考模型来优化特定应用领域的域权重。

什么是域？我们在实验中通过数据来源来定义域，但这只能实现粗粒度控制。使用细粒度域可以提高DoReMi的收益。例如，DoReMi在The Pile（22个域）上比GLaM数据集（8个域）更有效。

更广泛的影响。我们希望提高训练效率，减少训练大型LMs对环境的影响（Lacoste等人，2019，Ligozat等人，2021，Patterson等人，2021；Strubell等人，2019）。然而，大型LMs也被充分证明具有风险和偏见（Abid等人，2021，Blodgett和OConnor，2017，Bommasani等人，2021；Gehman等人，2020；Nadeem等人，2020）。DRO风格的训练旨在产生一个在所有领域都具有良好性能的模型，可能会对公平性产生有利影响，但这需要进一步调查。

7结论

我们介绍了DoReMi，一种用于训练语言模型的数据域重新加权算法。DoReMi能够在小模型上运行，并将优势转移到30倍大的模型上，只需改变域上的采样概率，即可在桩上实现2.6倍的训练加速。我们希望对以数据为中心的方法进行更多的研究，以提高语言模型的训练效率。

鸣谢

我们感谢陈翔宁、戴、邹宾·加赫拉马尼、巴拉吉·拉克什米纳拉亚南、保罗·米歇尔、吴永辉、郑、陈竺和广大谷歌巴德团队成员的深入讨论和建议。

图7：在4个模型尺度上平均一次下游精度，其中DoReMi的参考/代理模型与使用DoReMi的域权重训练的最终模型大小相同。该图中的所有模型都是在GLaM数据集上训练的。DoReMi不断提高跨尺度的下游精度。

GLaM数据集上的跨尺度结果

图7显示了GLaM数据集上不同尺度（280M、510M、760M、1B）的结果，其中代理/参考模型与使用DoReMi域权重训练的主模型大小相同。在所有尺度上，DoReMi与基线（统一）域权重和下游调谐域权重相当或更好。有趣的是，对于280M规模的迭代DoReMi，当用于训练280M模型时，第二轮权重实现的下游精度略低于第一轮权重，但更好地转移到训练8B模型。

B桩的详细结果

8B模型的每域困惑。表4显示了在桩上训练的8B模型的每个领域的困惑。在这种情况下，参考/代理模型是70M、150M、280M和1B。与基线域权重相比，DoReMi改善了每个域的困惑。

表4：在The Pile上训练的8B模型的每个领域的困惑，其中参考/代理模型是或更小的尺寸（70M、150M、280M、1B）。使用

表5：生成一次性任务的每个任务的精确匹配精度。所有DoReMi模型都比基线域权重显著提高了下游性能。

8B型号的每任务精度。表5显示了从70M到1B的各种参考/代理模型大小的一次性生成任务的精度。所有DoReMi模型在基线基础上显著提高了下游性能。

不同尺度的困惑结果汇总。表6显示了DoReMi在4个尺度（280M、510M、760M、1B）上的每个领域困惑的汇总。这里，参考/代理模型与使用DoReMi域权重训练的主模型大小相同。平均而言，DoReMi在The Pile的22个领域中提高了18.25个领域的困惑。相对于基线域权重，最坏情况下的困惑总是减少（或在510M的情况下相当）。

消融带来困惑。表7显示了DRO物镜消融的困惑。我们更改DRO目标，并使用这些目标来调整280M参考/代理模型的域权重。然后使用这些调谐的域权重来训练主280M模型。最困难是指在没有参考模型基线的情况下优化域级日志困惑。最简单是指在参考模型下对对数困惑度最低的域进行优化。二者都

表6：在The Pile上训练的280M、510M、760M和1B模型的每个领域困惑汇总，其中参考/代理模型的大小相同。DoReMi在所有情况下都改善了基线域权重的最坏情况和平均困惑。平均而言，DoReMi改善了22个领域中18个领域的困惑。

表7：DRO目标消融的困惑结果汇总（超额损失）。单独的组成部分（分别喜欢最难和最容易的领域）并不能减少对基线的困惑。

消融不会改善基线以上任何领域的困惑。为“最难”的领域进行优化实际上并不能改善最坏情况下的困惑，这支持了Oren等人的结果。（2019），该等人也使用DRO进行语言建模，但有基线损失。

域权重的轨迹。图8显示了DoReMi运行期间域权重的指数移动平均值（平滑参数0.99）。在这两种情况下，都有一些域最初的重量非常高，但重量下降得非常快（在50k步以内）。由于我们通过逐步积分这些曲线并进行归一化来计算最终的域权重，这表明如果我们的计算预算较小，这些域可能会变得更加重要——这突出了混合权重对计算预算的依赖性。同时，域权重往往在50k步后迅速稳定，这表明对于较大的计算预算，最佳域权重应该相似。我们还可以在50k步后利用这种稳定性，以较小的步数运行DoReMi，并外推域权重以节省计算。

图8:280M和1B参考/代理模型的DoReMi运行过程中域权重的指数移动平均值。在跑步开始时，域权重变化很快，然后在50k步后变得更加稳定。这表明，1）较小的计算预算可能需要截然不同的域权重，2）我们可以通过在50k步后外推域权重来节省计算。

280M和1B的域权重比较。表8显示了280M和1B代理模型下the Pile的DoReMi域权重。不同的代理模型大小会导致不同的域权重，这表明在域权重空间中可能存在多个局部极小值。对于280M代理模型，大部分权重放在Pile CC web文本域上，而对于1B代理模型的DoReMi，大部分权重则放在OpenWebText2上。其余域的域权重的总体模式是相似的。

C培训详细信息

数据预处理。对于所有数据集，我们通过将长度为1024个的示例分块为256k词汇大小的句子片段标记器来对数据进行预处理。示例按域分离，以便于分层采样（首先根据一些域权重对域进行采样，然后随机从该域中对示例进行采样）。减少

20

表8：桩上的域权重。基线域权重是根据默认的Pile数据集计算的。对于不同的代理模型大小，DoReMi（280M）和DoReMi（1B）导致不同的域权重。尽管存在差异，但除了哪个域名的权重最大之外，定性模式是相似的。

填充令牌，我们努力将示例（可能来自不同的域）打包到同一序列中。当进行这样的打包时，我们在DoReMi中计算每个令牌级别的域困惑。

The Pile的基线域权重。The Pile的基线域权重是根据The Pile数据集计算的，每个域的历元数在Gao等人（2020）中给出。在将1024个示例分块后，我们计算每个域中的示例数，并乘以Gao等人（2020）中指定的该域的历元数。然后，我们对这些计数进行归一化以获得基线域权重。

培训设置。对于所有训练运行（包括DRO运行），我们使用512的批量大小、1e-3的初始学习率、1e-2的权重衰减和到范数1的梯度剪裁进行训练。我们将学习率呈指数衰减，直到在训练结束时达到最小1e-4，线性预热占总训练步骤的6%。我们在The Pile上训练了200k步，在GLaM数据集上训练了300k步。1B参数下的模型用TPUv3加速器训练，而1B和8B模型用TPU4训练。

模型架构。表9显示了论文中使用的模型大小的架构超参数。我们使用的所有型号都是香草变压器解码器，只有256k人声大小的型号。

21

表9：论文中使用的各种模型规模的架构超参数。所有型号都是香草变压器解码器专用型号，使用256k大小的词汇表。

D数据重新加权没有权衡的简单示例

受第3.2节研究结果的启发，我们提出了一个简单的语言建模示例，其中重新加权来自不同领域的训练数据可以改善所有领域的困惑。该示例表明，DoReMi降低了熵极高或极低的域的权重。

设置。假设文本p*的基本真值分布是k个域上的混合，其中每个域z e｛1，…，k｝由m个令牌上的不同unigram分布p*（x|z）定义。给定n个训练样本的预算，目标是选择域权重p（z）（加1的k个标量）来对训练数据进行采样，这样我们就可以很好地学习从1到k的所有z的unigram分布p*（|z）的参数。值得注意的是，我们的目的不是估计各领域的基本事实混合比例。

数据给定一些域权重p（z），我们对训练数据进行分层采样：首先，我们通过从k个可能性上的多项式分布中提取概率，确定每个域z的样本数量nz，概率由p（z）和n个总试验定义。然后，对于每个域z，我们从p*（|z）中采样nz个令牌，形成长度为nz的令牌Xz的向量。

模型对于每个域z，我们考虑在unigram分布参数0 e Am上具有Dirichlet先验p（0|z；d）的unigram分布p（x|z；0）的贝叶斯模型。Dirichlet先前具有超参数d e Rm，其可以被视为每个令牌的“伪计数”。对于每个域z，我们通过计算以数据为条件的后验分布的平均值来估计参数0z：

对于域z，我们可以将该估计器的参数误差写成预测下一个令牌的“难度Hz”和先前Az的“质量”的函数，定义如下。

引理1。对于具有nz个样本的域索引z，参数错误为

没有折衷的例子。假设有3个域z e{1,2,3}，并且m＝3个词汇表标记x e{1,2，3}。我们使用对称的Dirichlet先验（优选均匀的令牌分布），其中对于所有令牌x和域z，Az（x）=1/3。在这个设置中，我们展示了一组域权重的参数误差严格低于基线，在基线中，我们从每个域中采样相同数量的令牌：所有域z的nz都相等。

假设unigram分布的基本真值参数为

这个不等式在这种情况下成立，因为H2＜1并且s2＝1。因此，参数误差在样本数量n2中减小。

因此，将示例从域3重新分配到域1和2的任何域权重都减少了所有域的参数误差。

什么样的域名被降低了权重？直观地说，我们可以下调非常嘈杂（高熵/难度）的域3，因为初始化与基本事实完全匹配。这允许我们将样本重新分配到其他域1和2。在这两者之间，域1需要较少的额外样本，因为参数误差随着样本数量ni而非常迅速地减小（难度Hi为零）。因此，最简单的域也应该获得相对较小的权重。在实践中，域之间的正向转移（这里没有捕捉到）也可能导致重新加权不会导致跨域权衡的情况。

使用DoReMi进行模拟。我们考虑在上述简单示例的无权衡实例上运行DoReMi，该示例具有等式14中的基本真值unigram分布。请注意，DoReMi的域重加权步骤（步骤2，算法1）涉及T迭代模型更新的循环，而方程2的估计器是以闭合形式计算的。为了使估计器适用于DoReMi，我们考虑了一种迭代版本，其中平均值是以在线方式计算的。在域权重更新率n=0.5的n=500个训练示例上，我们使用迷你批量大小1对T=500个步骤运行DoReMi。对于在步骤t对来自域z的示例x进行的模型更新，我们将伪计数Qz（x）增加对应于域z的当前域权重。我们不使用小批量中的示例（其大小仅为1，不代表所有域），而是使用30个示例的固定、独立评估集来计算算法1中的每个域的超额日志困惑。

我们将DoReMi与使用基线域权重训练的模型进行比较，基线域权重在3个域上是一致的。所有模型都是在n=500个训练示例上进行训练的。我们使用基本真值unigram分布参数来评估封闭形式的每个域上的模型的对数困惑性。

在这个简单的例子中，DoReMi在四舍五入到小数点后返回域权重[0.39，0.61，0.0]。这些权重对应于我们的直觉——第一个域（无噪声）增加了少量，第三个域（有噪声）减少到0权重，大部分权重分配给第二个域。我们使用这些域权重来生成一个包含500个示例的新数据集。用这个新数据集训练的模型在所有领域都比基线模型有所改进。