DoReMi（DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining，NeurIPS 2023，Google DeepMind + Stanford）是预训练数据配比优化的早期工作之一，后续 RegMix、DoGE、Online Data Mixing 等方法都在某种程度上回应或改进了它提出的框架。核心思路：用 Group DRO（Group Distributionally Robust Optimization，Sagawa et al. 2020 提出的最小化最差组损失的在线优化算法）在 280M 的小代理模型上动态调整域权重，取训练轨迹上的平均权重，再用这组权重去训练 8B 的大模型。整个配比搜索（reference + proxy 两个 280M 模型）的额外计算量只占 8B 主模型训练的约 8%，但主模型下游准确率提升约 6.5 个点，达到基线最终准确率所需的训练步数减少到约 1/2.6。

论文链接： arxiv.org/abs/2305.10429

问题与动机

预训练语料由多个域混合而成（Wikipedia、书籍、网页、代码等），各域的采样比例直接影响最终模型的质量。The Pile 的域权重是启发式设定的；PaLM/GLaM 的权重则是在下游任务上调出来的，代价是需要训大量模型做网格搜索，且容易过拟合到特定评测集。

DoReMi 想解决的问题：不借助任何下游任务标注，自动找到一组域权重，使得大模型在所有域上都表现良好。优化目标不是"平均 loss 最低"，而是"最差域的 excess loss 最低"，即一种 minimax（worst-case）准则。

三步流程

上图是 DoReMi 的整体流程。三步串联：先用初始域权重训一个小的 reference 模型，再用 Group DRO 训一个同尺寸的 proxy 模型并提取域权重，最后用这组优化后的域权重重采数据去训练大模型。

训练 reference 模型。用初始域权重（如 The Pile 原始权重或 GLaM 的均匀权重）训练一个 280M 参数的小模型，作为后续 DRO 的基线。这个 reference 模型的作用是给出每个样本的"基础难度"：reference loss 高的样本，说明该样本本身就高熵（比如随机文本），不需要额外投入；reference loss 低的样本，说明这个内容是"可学的"。

用 Group DRO 训练 proxy 模型并提取域权重。核心步骤。Group DRO 的基本思想是：把训练数据按域分组，不优化所有组的平均 loss，而是优化最差组的 loss。具体做法是维护一组域权重 $\alpha$，训练时交替做两件事：（1）对 $\alpha$ 做梯度上升，把权重集中到当前 loss 最高的域；（2）对模型参数做梯度下降，降低加权后的 loss。这个过程的效果是迫使模型在所有域上都不能太差，否则那个最差的域就会被 $\alpha$ 放大。

在 DoReMi 中，proxy 模型的训练目标不是标准的 NLL，而是 minimax 形式的 excess loss：

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

其中 $\ell_\theta(x)$ 和 $\ell_{\text{ref}}(x)$ 分别是 proxy 和 reference 模型在样本 $x$ 上的 NLL，$|x|$ 是样本的 token 数。内层 max 对域权重 $\alpha$ 做指数梯度上升，把权重集中到 excess loss 最大的域；外层 min 用标准优化器更新 proxy 模型参数。每一步都交替更新 $\alpha_t$ 和 $\theta_t$。训练结束后取整条轨迹上 $\alpha_t$ 的时间平均作为最终域权重。

实现细节：每步从均匀分布采 minibatch（而非按当前域权重采），逐 token 计算 excess loss 并 clip 到非负（满足 Group DRO 对损失非负的要求），指数梯度更新步长 $\eta = 1$，平滑参数 $c = 10^{-3}$。

用新域权重训大模型。把第二步得到的权重当成新的采样分布，重采数据训练 8B 参数的大模型，训练过程本身与标准预训练无异。

Excess loss 的设计逻辑

为什么不直接优化 proxy 模型的 raw loss？论文做了消融：只用 proxy loss（倾向于给模型最难的域加权）或只用负 reference loss（倾向于给 reference 模型最容易的域加权），都不如完整的 excess loss。

excess loss = proxy loss $-$ reference loss，度量的是 proxy 模型相对于 reference 的"改进空间"。高 excess loss 通常意味着样本可学（reference loss 不高），但 proxy 还没学好（当前训练不够充分）。低 excess loss 可能来自两个原因：样本本身高熵不可学（reference loss 也高），或者样本太简单 proxy 已经学好了。两种情况下都不需要额外加权。

这种设计使 DRO 不会把权重浪费在本质不可学的噪声域上，也不会浪费在已经学饱和的简单域上。

The Pile 上的实验

用 280M proxy 得到的域权重训练 8B 模型，在 The Pile 的 22 个域上评估。

上图是 8B 模型在 The Pile 上的下游 one-shot 准确率随训练步数的变化。DoReMi 在整个训练过程中持续领先基线，最终准确率高约 6.5 个点；基线需要 200k 步才达到的准确率，DoReMi 在约 75k 步就已经达到。

• 所有域的 perplexity 都下降。即使被大幅降权的域（如 ArXiv、PubMed Central、StackExchange），perplexity 也优于基线。这是论文最反直觉的发现。
• 下游 5 个 generative few-shot 任务（TriviaQA、NaturalQuestions、WebQuestions、SQuADv2、LAMBADA）平均准确率提升约 6.5 个点。
• 达到基线最终准确率只需约 75k 步，基线需要 200k 步，加速约 2.6 倍。

DoReMi 的域权重把 Pile-CC（Common Crawl 子集）从基线的 11.2% 上调到 60.6%，并大幅降低了 ArXiv（10.5% → 0.4%）、PubMed Central（10.7% → 0.5%）、StackExchange（9.3% → 1.5%）等域的权重。

上图是 8B 模型在 The Pile 各域上的 log-perplexity 对比。红色是基线，蓝色是 DoReMi。每一个域 DoReMi 都低于基线，包括被大幅降权的 ArXiv、PubMed Central 等。

为什么降权后 perplexity 反而下降？论文给出的假设：最低熵的域（内容简单重复）和最高熵的域（接近均匀分布）都不需要大量样本就能学好。把更多预算分配给中等熵的域（如网页混合体），产生的正向迁移可以惠及所有域。

GLaM 数据集上的实验

GLaM 只有 8 个域，且已有人工在下游任务上调好的域权重（downstream-tuned weights）。DoReMi 从均匀初始权重出发，跑 3 轮 iterated DoReMi（每轮把上一轮的输出域权重作为新 reference 的训练权重），最终域权重与 downstream-tuned weights 的整体模式接近：Filtered Webpages 权重最高（0.51 vs downstream-tuned 0.42），Conversations 和 Books 次之，Forums 和 News 被压低。

在下游准确率上，iterated DoReMi 第二轮的 8B 模型与 downstream-tuned 域权重训练的模型表现相当。

跨规模迁移与代理模型尺寸

论文考察了 proxy 模型规模对最终大模型的影响。固定大模型为 8B，proxy 尺寸从 70M 扫到 1B：

280M 是最优点。1B proxy 反而比 280M 差，论文将原因归为 Group DRO 优化器在较大模型上表现退化：1B proxy 模型本身的质量明显低于同尺寸的标准训练基线（19/22 域不如基线），而 280M proxy 模型的质量是优于基线的（19/22 域优于基线）。尽管如此，即使 1B proxy 自身质量不佳，它产出的域权重仍然能让 8B 主模型达到基线准确率 2 倍的加速。

另一组实验把 proxy 和 main 设为同尺寸（280M/510M/760M/1B），DoReMi 在所有规模上都一致提升约 2 个点下游准确率，达到基线准确率约 4 倍加速。

局限性

域定义的粒度问题。DoReMi 的效果与域的划分方式直接相关：The Pile 有 22 个域，效果明显；GLaM 只有 8 个域，提升幅度小得多。如果数据没有现成的域标签（如一个大型 Common Crawl dump），DoReMi 无法直接使用。论文讨论了通过聚类自动发现细粒度域的可能性，但未实验验证。

DRO 在大 proxy 模型上退化。从 280M 到 1B proxy，域权重质量下降。论文推测原因是 DRO 的 loss reweighting 与主模型的 resampling 之间存在 mismatch：proxy 训练时每个 batch 里不同域的梯度被加权，而主模型训练时是直接按权重重采样数据。论文观察到这种 mismatch 在 1B 规模上比 280M 更严重，但未给出完整的理论解释。

计算开销仍然可观。虽然 reference + proxy 的总训练量只占大模型计算量的约 8%，但 proxy 需要跑满大模型的全部训练步数（200k 步）。论文观察到域权重在训练早期变化最大，后期趋于稳定，提出可以提前截止并外推，但未做严格验证。

单次 DRO 可能陷入局部最优。280M proxy 和 1B proxy 得到的域权重模式不同（前者重 Pile-CC，后者重 OpenWebText2），论文认为域权重空间可能存在多个局部最优解。iterated DoReMi 在 GLaM 上有效（3 轮收敛），但在 The Pile 上的多轮迭代效果未报告。

小结

DoReMi 把预训练配比优化表述为一个对 excess loss 做 minimax 的问题，用 Group DRO 在小模型上求解，把得到的域权重迁移给大模型。后续工作基本都沿用了这个"小模型代理大模型做配比搜索"的范式（RegMix 用回归代替 DRO，DoGE 用梯度信号代替 loss，Online Data Mixing 做在线调整）。DoReMi 确立的实验事实是：280M 模型上优化出的域权重可以直接加速 8B 模型的训练，且在 The Pile 实验中改善是全域覆盖的，22 个域的 perplexity 全部优于基线。