3.2.1 Step 1 训练一个参考模型#

首先，本文在基于经验的情况下得到了一组 domain weight来采样数据，训练 T 步，批大小为 b，得到参考模型 $p_{ref}$，这个模型仅为 280M。

3.2.2 Step 2 优化 domain weight#

为了优化并获取最优的 domain weight，本文基于 DRO-LM 框架，并使用 Group DRO 优化器来训练一个小的代理模型 $P_{\theta}$ ，其中 $\theta$ 为代理模型的权重。 这个框架通过优化最坏情况下的损失来训练一个 robust 模型，这个训练目标等价于以下 minimax 目标：

$\min_{\theta} \max_{\alpha \in \Delta^k} L(\theta, \alpha) := \sum_{i=1}^{k} \alpha_i \cdot \left[ \frac{1}{\sum_{x \in D_i} |x|} \sum_{x \in D_i} \left( \ell_{\theta}(x) - \ell_{\text{ref}}(x) \right) \right]$

这个公式看着很会唬人，实际上阅读起来非常简单，我们来拆解一下：

1. $\ell_{\theta}(x) = -log p_\theta(x)$ 和 $\ell_{\text{ref}}(x) = - log p_{ref}(x)$ 这俩分别是代理模型和参考模型的负对数似然，这俩做减法得到的差值是用来衡量代理模型相对于参考模型，在当前样本上还有多大的提升空间。

2. $|x|$ 是 $x$ 中的 tokens 数量，这里内部的求和是为了求第 $i$ 个域上的所有 token 的总数然后做公平的 token 级归一化，防止长样本总损失高的问题。

3. 最外层的求和相当于是总的加权域损失，把所有域的平均超额损失，按域权重进行加权，得到总损失 $L(\theta,\alpha)$。

4. 最外层的 $\min_{\theta} \max_{\alpha \in \Delta^k}$：

   • 内层的 $max$ 为的是给域权重 $\alpha$ 做优化，放大短板域的影响，核心的是，哪个域的平均超额损失越高，就把这个域的权重调的越高，让其在总损失上的占比更大。

   • 外层的 $min$ 为的是给代理模型权重 $\theta$ 做优化，把放大后的短板损失给压下去。

此外，这里的 $\Delta^k$ 代表了 $k$ 维标准单纯形，它代表了一个概率分布的集合，如果一个向量 $\alpha$ 属于 $\Delta^k$，那么它必须满足：

1. 非负性：向量里每一个元素都大于等于 0（$\alpha_i \geq 0$）;

2. 归一性：所有元素的加和必须等于 1，即（$\Sigma_{i=1}^k \alpha_i = 1$）。

直观理解它就是一个权重分配，比如我们有 3 个不同的数据领域（维基百科、代码、书籍)，$\alpha$ 就代表了这三个领域分配的权重系数：

1. $\alpha = [0.3, 0.5, 0.2]$，这是一个合法的 $\Delta^k$，因为其非负，且和为 1 。

2. $\alpha = [0.8, 0.4, -0.2]$，这是一个非法的 $\Delta^k$，因为其虽然和为 1 ，但存在负数权重。

实际上，这个 3 维的单纯形约束为：$\alpha_1 + \alpha_2 + \alpha_3 = 1$（$\alpha_i \geq 0$），也就是一个等边三角形平面区域：

• 如果在三个顶点位置，说明权重分配的极其不公平比如：$\alpha = [1, 0, 0]$

• 如果在三条边上，代表其中一个权重为0，只在另外两个域上分配

• 如果在中心点 $(\frac{1}{3}, \frac{1}{3},\frac{1}{3})$，则代表权重分配完全公平，三个领域完全一样。

3.2.3 Step 2 中的一些细节#

我们来看一下这个伪代码，初始化的 domain weight 我们全部均匀初始化，$1$ 为一个全 1 的 $k$ 维向量，$\alpha_0$ 就是一个 $k$ 维的单纯形。

每一轮迭代都会完成：「采样 → 计算 excess loss → 更新 domain weight → 更新模型」，其中：

1. 采样训练的批次：

   • 我们可以看到，在优化的过程中，采样的分布 $P_u$ 是一个均匀分布，所有域的权重都是均等的，这也是 Group DRO 算法的标准要求，其内部的经验风险评估必须在均匀的先验分布上进行，然后通过权重去放大那些表现差的样本对梯度的贡献，目的是公平的计算每个 domain 的超额损失。

   • 如果用之前的 weight 进行采样，很可能超额损失高但 weight 较小的 domain 根本抽不到样本。

   • DoReMi 的本质就是解一个 Minimax 问题，在这个零和博弈中，$max$ 的任务是寻找超额损失最高的 domain，要做到这一点，其必须持续拥有对所有领域 Loss 的全局视野，不然其无法收敛。

2. 计算样本的 token 长度 $|x|$，为了防止不同样本和不同域的文本长度差异问题，只有按 token 归一化后才能公平对比 excess loss 。

3. 计算超额损失：$\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_{\text{ref},j}(x), 0 \right\}$：

   • $\max\left\{ \ell_{\theta_{t-1},j}(x) - \ell_{\text{ref},j}(x), 0 \right\}$:

     • 这里是针对一句话 $x$ 里面的第 $j$ 个词（Token）进行计算。

     • 前者是代理模型预测这个词的 Loss（负对数似然），Loss越大说明模型越不会预测这个词。

     • 后者是参考模型预测同一个词的 Loss 。

     • 相减是衡量代理模型比参考模型落后了多少，以及过滤一些困难样本或毫无规律的乱码样本，因为这种样本在两个模型上的 Loss应该都会很大。

     • 这里，也不用担心一旦代理模型超过参考模型之后会无法增加权重的情况，因为此时重心会放到那些依然落后 baseline 的 domain 上，这是一种动态的负载均衡。并且，我们也不关心代理模型到底有多强，它的唯一作用就是用来找到哪些 domain 上的数据最值得学习。

     • 此外，作者也考虑到了如果一开始参考模型太弱或者初始权重太差，会把优化方向带偏，作者提出了 Iterated DoReMi 进行补充，相当于优化一次得到权重后，再基于这些权重训练参考模型，再进行优化得到权重，反复左脚踩右脚的操作。

     • 外侧使用了 $ReLU$ 截断 $max \{ \cdots , 0\}$，因为 Group DRO 算法要求损失必须是非负的，从直观上理解，如果代理模型在某个词上已经做的比参考模型好了，那么我们不应该给它奖励，我们目的是找最薄弱的环节。

   • $\sum_{j=1}^{|x|}$ 这个是句子级别的求和， $|x|$ 是这句话的总 Token 数量，把这句话里面的每一个词被截断后的超额损失都加起来，就得到了这句话的总超额损失。

   • $\sum_{x \in B \cap D_i}$ 其中，B 代表当前均匀采样出的这批数据，$D_i$ 代表是第 $i$ 个领域，$B \cap D_i$ 代表把当前 Batch 中，属于 $D_i$ 域的所有句子都挑出来。

   • $\frac{1}{\sum_{x \in B \cap D_i} |x|}$ 这里是最终的归一化操作，分母表示，在这个 Batch 中，属于 $D_i$ 域的总 Token 数有多少，最终完成 Token 级别的归一化。

   • 整个式子计算出了 $D_i$ 域在当前 Batch 下的平均每个 Token 的超额损失。

4. 更新 domain weight， $\alpha_t' = \alpha_{t-1} \cdot exp(\eta\lambda_t)$，这一步是指数梯度上升的更新操作，$\eta$ 是步长，如果 $\lambda_t[i]$ 大，则 $exp(\eta\lambda_t[i])$ 就会是一个大于 1 的数，乘以上一步的权重 $\alpha_{t-1}[i]$ 后，领域 $i$ 的权重就被放大了，反之，如果超额损失为 0，则其相对比重就会在后续的归一化操作中缩小。

5. 归一化与平滑操作，$\alpha_t = (1-c)\frac{\alpha_t'}{\Sigma\alpha'_t} + c \cdot u$：

   • 这里，将所有的权重等比例缩小，使其总和为 1（满足单纯形要求）

   • 引入 $c \cdot u$ 是一种平滑策略（探索机制），强制给每个领域保留一个极小的基础概率，这样防止某个领域的权重归零后永远无法被更新到，相当于一种正则化技巧。

   • 实际上这个就是一个插值操作，点 A 纯粹由数据驱动算出来的最新概率分布，位于单纯形的某处（可能非常靠近边缘），点 B (u) 是均匀分布，位于单纯形的绝对几何中心，插值参数 c，决定了在 A 和 B 之间连一条线段，最终的点 $\alpha_t$ 取在这条线段上的什么位置。

   • 这样保证了指数梯度上升算法在理论上能以更优的速率收敛（不会在边界附近疯狂震荡）

   • 提供了一个下界保障，保证每个领域在优化的过程中不会饿死

   • 将整个搜索空间从整个单纯形，稍微往里收缩了小一点，防止了出现绝对0的情况，避免了过拟合到单一领域，也改善了优化地形（指数梯度上升算法如果太靠近边缘，容易步子迈得太大，导致不同域上反复横跳）

6. 更新代理模型，使用最新计算出来的 $\alpha_t$ 作为加权系数，对代理模型的 Loss 求加权和：$L(\theta_{t-1}, \alpha_t) = \Sigma_{i=1}^k\alpha_t[i] \cdot \mathrm{Loss}_i$，使用常规的优化器（Adam/Adafactor）进行反向传播，更新代理模型参数 $\theta_t$，因为刚才我们调高了薄弱领域的权重，优化器在更新参数时，就会更加侧重于优化这些 domain 。

3.2.4 Step 3 使用新的 domain weight 训练大模型#

在第二步得到的一组 domain weight $\overline{\alpha}$，其定义了一个新的数据分布 ${P_{\overline{\alpha}}}$，我们从这个新的分布中重新采样数据，以标准的训练过程来训练一个主模型。