解耦搜索与训练:通过模型合并扩展大语言模型预训练数据混合 Decouple Searching from Training: Scaling Data Mixing via Model Merging for Large Language Model Pre-training
用模型合并代替重复训练代理模型,高效搜索最优预训练数据混合比例
前置知识
数据混合(Data Mixing)
在大语言模型预训练阶段,需要将来自不同领域(通用文本、数学、代码、多语言等)的数据按一定比例混合成训练语料。不同的混合比例会显著影响模型在各类任务上的最终表现。例如,增加数学数据比例可提升数学推理能力,但可能损害通用语言理解。找到最优的混合比例是预训练配方设计的核心难题之一,传统做法依赖大规模代理实验,成本极高。
本文的核心目标就是优化数据混合比例,理解这个概念是理解整篇论文的基础。
代理模型(Proxy Model)
为了探索最优数据混合比例,研究者通常训练一系列较小规模的模型来近似评估不同混合配置的效果。这些用于评估的模型被称为代理模型。传统方法如 RegMix 和 CLIMB 需要训练大量(如 112 个)小规模代理模型,每个模型使用少量 token(如 2B)训练,然后通过回归预测器外推到大规模训练场景。然而,小规模代理与目标模型之间的能力差距导致预测不可靠。
DeMix 的创新点在于用模型合并生成免费的代理模型,彻底改变了代理模型的构建方式。
模型合并(Model Merging)
模型合并是一种免训练的技术,通过算术运算(如加权平均)将多个具有相同架构的模型组合成新模型。假设模型参数为 $\Theta$,合并后的模型参数可以表示为 $\Theta_{merged} = \sum_{i=1}^{N} \alpha_i \Theta_i$,其中 $\alpha_i$ 是权重。近年来,研究发现在参数更新幅度较小的情况下(即 $\delta \ll 1$),分别在不同数据集上训练的模型的权重增量近似满足可加性:$\Delta(D_i \cup D_j) \approx \Delta(D_i) + \Delta(D_j)$。
这是 DeMix 方法的理论基础——通过合并已训练的组件模型来模拟在混合数据上训练的效果,从而免去重复训练的开销。
Spearman 秩相关系数(Spearman's ρ)
一种用于衡量两个变量之间单调关系强度的非参数统计指标,取值范围 [-1, 1]。在本文中,Spearman's ρ 用于衡量代理模型与参考模型在不同数据混合比例下的排名一致性。例如,若代理模型认为混合比例 A 优于 B,而真实训练的参考模型也认为 A 优于 B,则两者排名一致。ρ 越接近 1,说明代理模型越能准确预测不同混合配置的相对优劣。
这是评估 DeMix 代理模型质量的核心指标,直接衡量搜索结果的可靠性。
研究动机
在大语言模型预训练中,确定最优数据混合比例是一项关键但极具挑战性的任务。现有方法面临严重的效率-准确性权衡困境:一种策略是训练中等规模的代理模型(如 8B 参数),使用较大 token 预算(如 100B)进行实验,虽然信号相对准确,但计算成本高昂,无法系统性地探索混合空间;另一种策略如 RegMix 和 CLIMB,通过大量极小规模代理实验(如 2B token 训练的小模型)来自动化搜索过程,但这些轻量级代理的有效性受到质疑——由于代理与目标模型之间存在巨大的能力差距,这些自动化策略在数学推理和代码生成等复杂任务上往往失败。论文指出,在 112 个代理、每个代理 2B token 的设置下,传统方法在数学任务上的 Spearman's ρ 仅为 -0.20(负相关),说明代理完全失效。
本文的目标是本文的核心目标是打破「充分性(足够的代理数量)、准确性(代理信号可靠)、效率(计算成本可控)」三者之间的矛盾,提出一种能够同时实现三者的数据混合搜索框架。具体而言,DeMix 旨在:以固定且较低的训练预算(约 200B token),生成无限数量的代理模型,这些代理模型在排名上与真实训练的参考模型高度一致(Spearman's ρ > 0.8),从而发现能够平衡通用语言理解、数学推理和代码生成能力的最优数据混合比例。
与已有工作不同的是,DeMix 抓住了一个被已有工作忽视的关键观察:传统方法将「搜索」和「训练」耦合在一起——每采样一个新的混合比例就需要重新训练一个代理模型,导致搜索成本与代理数量线性增长。然而,如果能将搜索过程从训练中解耦出来,即通过某种免训练的方式生成代理模型,就能以固定成本获得无限数量的代理,从而进行更充分的搜索。DeMix 的独特视角在于利用模型合并的理论基础——当参数更新幅度较小时,加权合并独立训练的组件模型等价于在混合数据上训练——将组件模型训练一次后,通过改变合并权重即可生成任意数量的代理模型。
核心方法
DeMix 的核心思想可以用一个类比来理解:传统方法像是为每道菜都重新准备食材并烹饪(每种混合比例都重新训练模型),而 DeMix 则是预先准备好几种核心食材(组件模型),然后通过不同的配比组合来试菜(通过模型合并生成代理模型),既节省时间又能尝试无限多种组合。具体技术路线分为四个阶段:首先,对大规模异构数据进行严格清洗和分类,得到 N 个候选数据集;其次,在每个候选数据集上独立训练一个组件模型;然后,通过加权线性合并这些组件模型来生成代理模型,其中合并权重就是目标数据混合比例;最后,基于代理模型的评估结果,使用回归预测器迭代优化混合权重,找到最优比例。
DeMix 与 RegMix、CLIMB 等已有方法的本质区别在于:它将「搜索混合比例」与「训练模型」彻底解耦。传统方法中,每采样一个混合比例 $\alpha = \{\alpha_1, ..., \alpha_N\}$,就需要从头训练一个代理模型 $M_\alpha = T(D_\alpha, \Theta_{base})$,搜索成本与代理数量成正比。DeMix 的理论基础是:当模型参数更新幅度 $\delta = \frac{\sum|T(D, \Theta_{base}) - \Theta_{base}|}{\sum|T(D, \Theta_{base})| + \sum|\Theta_{base}|} \ll 1$ 时,有近似等式 $\Delta(D_i \cup D_j) \approx \Delta(D_i) + \Delta(D_j)$。这意味着,只要在各候选数据集上训练好组件模型 $\Theta_i = \Theta_{base} + \Delta(D_i)$,就可以通过加权合并 $M_{mix} = \sum_{i=1}^{N} \alpha_i \Theta_i$ 来近似在混合数据上训练的效果。在本文实验中 $\delta \approx 10\%$,满足小更新假设。这样,N 个组件模型训练完成后,生成代理模型的边际成本几乎为零,可以采样任意数量的混合比例进行评估。
方法步骤详情
DeMix 的完整流程包括四个阶段。第一阶段是数据预处理:从多个来源(通用语料、数学数据集、代码集合)收集大规模数据,执行去重(精确去重和 MinHash 模糊去重,目标相似度阈值 90%)、困惑度过滤(使用 Qwen3-0.6B 作为评分模型)、FastText 质量分类等清洗操作,最终将清洗后的语料分为 7 个候选数据集类别。第二阶段是组件模型准备:先在一个通用数据集 $D_{base}$(50B token)上从头训练一个基础模型,然后在每个候选数据集 $D_i$ 上混合 50% 通用数据继续训练(总 token 预算 $50 \times 7 = 350B$,实际使用的组件模型训练预算为 $30 \times 7 = 211B$),得到 7 个组件模型 $\{\Theta_1, ..., \Theta_7\}$。第三阶段是模型合并作为代理:对每个采样的混合比例 $\alpha^j = \{\alpha^j_1, ..., \alpha^j_7\}$,通过加权线性合并 $M^j_{mix} = \sum_{i=1}^{7} \alpha^j_i \Theta_i$ 构建代理模型,然后在 benchmark suite 上评估其平均排名 $r_j$。第四阶段是迭代优化混合权重:收集 $(\alpha^j, r_j)$ 对后训练 LightGBM 回归预测器,用预测器对新采样的混合比例打分,选择排名靠前的比例进入下一轮迭代(共 3 轮:64+32+16 个样本),最终从 top-128 混合比例中取平均得到最优配置。
技术新颖性
DeMix 的技术新颖性体现在多个层面。首先,在方法论上,它首次将模型合并理论应用于预训练数据混合搜索,实现了搜索与训练的解耦。虽然之前有 Merge-to-Mix 等工作使用模型合并进行数据选择,但它们仅用于微调阶段的二元子集选择,而 DeMix 处理的是连续值混合权重且面向预训练,挑战性更大。其次,在效率上,DeMix 在 212B token 总预算下达到的代理准确性(macro avg. ρ = 0.81)需要传统方法 1344B token(6.4 倍)才能匹配。第三,在工程贡献上,DeMix 发布了 DeMix Corpora——一个包含 22T token(15T 原始 + 混合扩展)的高质量预训练数据集,涵盖通用、多语言、数学、代码领域,并附带经过验证的混合比例,填补了该领域缺乏可直接复用的大规模混合语料的空白。
实验结果
DeMix 在代理一致性和混合质量两个维度上都取得了显著成果。在代理一致性方面(Table 2),当使用 30B token 训练的组件模型时,DeMix 达到 macro avg. Spearman's ρ = 0.81、Top 25% ρ = 0.59,而同等预算(~212B token)的传统训练方法仅为 0.53 和 0.20。传统方法需要 1344B token(12B token/代理 × 112 代理)才能达到类似水平,成本是 DeMix 的 6.4 倍。DeMix 的能力恢复率(Capability Recovery Rate)最高达 0.85,说明合并后的代理模型能够保留参考模型约 85% 的绝对性能。在混合质量方面(Table 3),使用 224 个合并代理的 DeMix 取得最优平均排名 24.00(在 96 个参考模型中),优于 RegMix 最佳配置的 28.00 和 CLIMB 最佳配置的 27.67。值得注意的是,DeMix 的代理数量从 56 增加到 224 时,排名持续改善(29.33→25.67→24.00),但到 448 时略有下降(27.67),表明过多代理可能引入过拟合噪声的风险。在 8B 模型上的迁移验证(Table 6)显示,DeMix 产生的混合比例同样优于基方法,证明其有效性不局限于小模型。消融实验表明:线性合并是最简单有效的合并策略;候选数据集中混入 50% 通用数据对代理准确性至关重要(降至 25% 时 ρ 从 0.787 降至 0.667)。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 通用语言理解(ARC-E, HellaSwag, PIQA, SIQA, WinoGrande) | 平均排名(在 96 个参考模型中) | DeMix: 46 (224 proxies) | RegMix 最佳: 2, CLIMB 最佳: 2 | DeMix 在通用任务上排名略逊,但在综合排名上最优 |
| 代码生成(HumanEval, MBPP) | 平均排名 | DeMix: 12 (224 proxies) | RegMix 最佳: 29, CLIMB 最佳: 25 | 排名提升约 58-60% |
| 数学推理(GSM8K, MATH) | 平均排名 | DeMix: 14 (224 proxies) | RegMix 最佳: 21, CLIMB 最佳: 11 | 与 CLIMB 最佳接近,优于 RegMix |
| 综合排名(General + Code + Math 宏平均) | 宏平均排名 | DeMix: 24.00 (224 proxies) | RegMix 最佳: 28.00, CLIMB 最佳: 27.67 | 综合排名最优,提升 13-15% |
| 代理准确性(与 96 个参考模型排名一致性) | Spearman's ρ (macro avg.) | DeMix: 0.81 (30B component) | 训练代理 1344B: 0.82 | 仅用 1/6.4 的训练预算达到相当准确性 |
局限与改进
论文存在几个值得关注的局限性。首先,DeMix 的理论基础依赖于参数更新幅度 $\delta \ll 1$ 的假设,实验中 $\delta \approx 10\%$,但当组件模型训练 token 数增加时(如 50B),$\delta$ 增至约 10.5%,代理一致性开始下降(代码任务 Consistency 从 0.97 降至 0.75),这限制了组件模型的训练规模。其次,论文仅在 1.7B 和 8B 两个模型规模上验证,未探索更大规模(如 70B、400B)下的适用性,而大规模预训练正是数据混合最核心的应用场景。第三,DeMix 假设候选数据集数量 N 较小(本文 N=7),搜索空间维度受限;当候选数据集更细粒度划分(如 15 个类别)时,虽然代理准确性保持,但搜索效率下降,且小类别的混合比例优化可能不可靠。第四,论文的基准比较对象主要是 RegMix 和 CLIMB,缺乏与 DoReMi、Rho Loss 等其他数据混合方法的直接对比(作者以「这些方法依赖评估损失而非代理」为由排除)。最后,DeMix Corpora 的 22T token 规模虽然庞大,但数据来源主要基于开源语料,对于拥有私有高质量数据的大厂而言,该方法的实际增量价值可能有限。
独立分析的弱点
从独立分析的角度,DeMix 存在以下几个值得关注的弱点。第一,组件模型的多样性问题:DeMix 仅训练 N=7 个组件模型,每个组件模型在其对应数据集上训练时混入 50% 通用数据,这可能导致组件模型之间的差异度不够大,合并后的代理模型可能无法充分探索混合空间的边界情况。改进方向可以是探索自适应的组件模型数量,或引入更多样化的训练策略(如不同的学习率、不同的基础模型)。第二,线性合并的局限性:论文虽然比较了多种合并方法,但所有方法都假设参数空间是「平坦」的,即权重可以线性插值。然而,深度神经网络的损失景观高度非线性,当组件模型之间差异较大时,线性合并可能产生位于损失盆地之外的无效模型。可以探索基于低秩适应(LoRA)的合并方式或使用感知器级合并来缓解此问题。第三,搜索策略的随机性:当前的迭代搜索完全基于随机采样,没有利用混合空间的几何结构信息。可以考虑使用贝叶斯优化或进化算法来更高效地探索高维混合空间。第四,评估指标的单一性:DeMix 仅使用 benchmark 排名作为优化目标,没有考虑训练稳定性、收敛速度等其他重要指标,实际预训练中这些因素同样关键。
未来方向
基于 DeMix 的成果,未来研究可以从多个方向延伸。首先,探索 DeMix 在更大规模模型(如 70B、400B)上的适用性是一个重要方向,因为大规模预训练的数据混合优化需求最为迫切,但也面临计算资源的巨大挑战。其次,可以将 DeMix 与 curriculum learning 结合,在训练过程中动态调整数据混合比例——前期侧重通用数据多样性,后期增加数学和代码数据比例,论文中提到的三阶段训练框架已经初步体现了这一思想。第三,DeMix 的组件模型训练可以探索更高效的方式,如使用 LoRA 微调替代全参数训练,或使用知识蒸馏来压缩组件模型。第四,将 DeMix 扩展到多模态预训练的数据混合优化是一个自然的延伸,当前视觉语言模型(如 Qwen-VL)的图文比例优化同样面临类似挑战。最后,DeMix 发布的 DeMix Corpora 及其验证过的混合比例,为后续研究提供了标准化的基准,可以用于系统性地比较不同数据混合方法的优劣。
复现评估
DeMix 在可复现性方面做得相对较好。论文在 GitHub 上开源了代码和 DeMix Corpora(https://github.com/Lucius-lsr/DeMix),提供了详细的实现细节:组件模型使用 Qwen3-1.7B 作为骨干,全局 batch size 512,序列长度 8192,初始学习率 3e-4 余弦衰减(最小保持 20%),通用数据混合比例 β=0.5。回归预测器使用 LightGBM,学习率 0.02,迭代次数 300。评估使用 OpenCompass 框架,涵盖 5 个通用 benchmark、2 个代码 benchmark 和 2 个数学 benchmark。然而,复现存在一些门槛:组件模型训练需要 7×30B=211B token 的预算,约需数百 GPU 小时(H800),对于中小团队仍有一定成本;DeMix Corpora 的完整下载和预处理也需要较大的存储和计算资源。此外,论文未提供完整的训练日志和 checkpoint,可能影响精确复现。总体而言,对于拥有中等计算资源的团队,DeMix 是可以复现的,但需要仔细对齐所有超参数。
论文图表