DoReMi:Optimizing Data Mixtures Speeds Up Language Model Pretraining

什么是困惑度：https://zhuanlan.zhihu.com/p/114432097
简单说明自回归的深度学习模型中，困惑度的计算就是：
$ppl=P(w_1,w_2,...,w_N{)}^{-\frac{1}{N}}=[P(w_1)P(w_2|w_1)...P(w_N|w_1...w_{N-1}){]}^{-\frac{1}{N}}$

论文背景：

数据域由不同来源的数据分布组成。例如维基百科，书籍，网文等。而数据域分布极大的影响了语言模型的性能。论文希望对于域权重的调整能够具有更好的适应性，即在不知道下游任务的情况下，通过最小化域的最坏情况损失来找到域权重，从而使得模型在所有域上都表现良好。于是文章提出了一种计算域权重的算法。使用域上的**组分布鲁棒优化 (Group DRO) **训练一个小型代理模型，以在不知道下游任务的情况下生成域权重（混合比例）。然后论文使用得到的域权重对数据集进行重新采样，训练一个更大的模型。
域权重：不同域数据集在整个数据集的占比

---

实现流程：

符号说明：假设有k个域，每个域i记为$D_i$
step1：基于初始化的一组权重，训练一个reference model $p_{ref}$
step2：训练一个小Proxy model $p_{\theta }$
优化目标：${\min }_{\theta }{\max }_{\alpha \in {\Delta }^k}L(\theta ,\alpha )={\min }_{\theta }{\max }_{\alpha \in {\Delta }^k}{\sum }_{i=1}^k{\alpha }_i\cdot \left[\frac{1}{{\sum }_{x\in D_i}|x|}{\sum }_{x\in D_i}{\ell }_{\theta }(x)-{\ell }_{\mathrm{ref}}(x)\right]$
${\ell }_{\theta }(x)-{\ell }_{\mathrm{ref}}(x)$称为超额损失excess loss，

组分布鲁棒优化（group DRO）：Group DRO 优化器通过在训练步骤 t 上交织域权重 αt 的取幂梯度上升更新和代理模型权重 θt 的梯度更新来工作。

---

算法流程：
Require：数据域 $D_1,...,D_K$，训练步数 T，批量大小 b，步长 η，平滑参数 c ∈ [0, 1]（例如，在论文的实现中，c = 1e-4）。
初始化代理权重${\theta }_0$
初始化域权重${\alpha }_0=\frac{1}{k}$
for t in 1 to T do:
从p中采样批次大小为b的minibatch B = { x 1 , . . . , x j } B=\left \{ x_1,...,x_j \right \} B={x1​,...,xj​}
$|x|(|x|)\le L$表示样本x的token长度
计算每个域 i ∈ { 1 , 2 , . . . , k } i \in \left \{ 1, 2,...,k \right \} i∈{1,2,...,k}的excess loss： λ t [ i ] ← 1 ∑ x ∈ B ∩ D i ∣ x ∣ ∑ x ∈ B ∩ D i ∑ j = 1 ∣ x ∣ max ⁡ { ℓ θ t − 1 , j ( x ) − ℓ r e f , j ( x ) , 0 } \lambda_{t}[i] \leftarrow \frac{1}{\sum_{x \in B \cap D_{i}}|x|} \sum_{x \in B \cap D_{i}} \sum_{j=1}^{|x|} \max \left\{\ell_{\theta_{t-1}, j}(x)-\ell_{\mathrm{ref}, j}(x), 0\right\} λt​[i]←∑x∈B∩Di​​∣x∣1​∑x∈B∩Di​​∑j=1∣x∣​max{ℓθt−1​,j​(x)−ℓref,j​(x),0}
更新权重：${\alpha }_t^{{}^{\prime }}\leftarrow {\alpha }_{t-1}exp(\eta {\lambda }_t)$
重新归一化和平滑域权重：${\alpha }_t^{{}^{\prime }}\leftarrow (1-c)\frac{{\alpha }_t^{{}^{\prime }}}{{\sum }_{i=1}^k{\alpha }_t^{{}^{\prime }}[i]}+cu$
利用$L({\theta }_{t-1},{\alpha }_t)$ 更新模型权重${\theta }_t$
end for
retrun $\frac{1}{T}{\sum }_{t=1}^T{\alpha }_t$

---

**step3：**使用新的数据域权重训练大模型

---

具体实验：

数据集介绍：

1、The Pile dataset （Gao 等人，2020）
Pile是一个 800GB 的文本数据集，包含 22 个域。默认域权重是启发式确定的。我们使用 The Pile 数据集中的默认域权重来训练基线，并作为 DoReMi 中的参考域权重。

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

对于 The Pile，训练每个模型 20万步；对于 GLaM 数据集，每个模型30 万步训练。

实验结果：

1、Pile的下游精度提高
图左显示了 The Pile 上基线模型和 DoReMi (280M→8B) 模型的平均下游性能。 DoReMi 将下游精度提高了 6.5%，并在 75k 步内实现了基线精度，比基线（200k 步）快了 2.6 倍。因此，DoReMi 可以显着加快训练速度并提高下游性能。

2、GLaM的下游调整权重的性能不变
论文在 GLaM 数据集上使用迭代 DoReMi 超过 3 轮。我们发现第二轮和第三轮的域权重几乎相同（表2）。图右显示了前两轮迭代 DoReMi 的一次性结果。第一轮之后，DoReMi 主模型的下游精度与基线（统一域权重）相当。第二轮之后，DoReMi 主模型实现了与下游调整域权重相当的下游精度。总体而言，域重新加权对 GLaM 的影响较小，可能是因为与 The Pile 中的 22 个域相比，只有 8 个域。
3、困惑度降低
图 4 显示了Pile 8B 模型的每个域的困惑度。尽管为某些域分配了较低的权重，DoReMi 仍显着降低了所有域的基线上的困惑度。
直观上，熵最低和最高的域可以被降低权重，而不会太大影响困惑度。从统计上看，最低熵域需要很少的样本来学习。最高熵域的令牌分布接近常见的统一先验——例如，随机初始化的模型倾向于输出统一的下一个令牌分布。因此，我们需要更少的样本来适应这些领域。从分配更多样本到中等熵域的正向转移可以改善所有域的困惑度。

4、多尺度消融分析

5、优化器Group DRO消融分析

6、Proxy model规模对较大主模型性能的影响
、

总结：

论文提出了 DoReMi，一种用于训练语言模型的数据域重新加权算法。 DoReMi 能够在小型模型上运行，并将优势转移到 30 倍大的模型上，只需更改域上的采样概率，即可将 Pile 上的训练速度提高 2.6 倍。
DRO 式训练旨在产生在所有领域都具有良好性能的模型，可能会对公平性产生有利影响，但这需要进一步研究。