FastMix：通过梯度下降的快速数据混合优化

# FastMix: 通过梯度下降实现快速数据混合优化
来源: https://arxiv.org/html/2606.14971
Haoru Tan¹,²  Sitong Wu³  Yanfeng Chen²,†  Jun Xia²  Ruobing Xie²  Bin Xia³  Xingwu Sun²  Xiaojuan Qi¹,†  
¹香港大学  
²混元大模型，腾讯  
³香港中文大学

###### 摘要

尽管大规模、多样化的数据集推动了大型模型的最新进展，但确定预训练和后训练的最优数据混合仍然是一个重要的未解难题。我们通过 **FastMix** 应对这一挑战，这是一个新颖的框架，仅训练**一个代理模型**即可自动发现数据混合。**FastMix** 不依赖于预定义的启发式规则或资源密集型的模拟，而是联合优化混合系数和模型参数，从而在效率和可扩展性上显著优于先前方法。**FastMix** 的核心是将混合选择重新表述为**双层优化**问题。在这种重新表述下，我们展示了优化混合比例在数学上等价于在统一源采样下为每个源分配可微的损失权重。这直接将混合系数嵌入到可微的迭代优化目标中，使得能够对混合和模型进行高效的、基于梯度的优化。为解决该优化问题，**FastMix** 实现了一种近似迭代优化过程，交替进行：(i) 根据当前混合比例对数据采样，更新模型参数（内循环）；(ii) 基于验证反馈更新混合比例（外循环）。在预训练和后训练中，**FastMix** 在显著降低搜索成本的同时优于基线方法。代码: [https://github.com/hrtan/fastmix](https://github.com/hrtan/fastmix)

参见说明图1: 各种数据混合策略的平均性能与时间成本（GPU小时）对比。  
(a) 预训练：我们提出的 **FastMix**（我们的）方法以最低的时间成本实现了最高的性能。注释强调，它比 CLIMB (Diao et al., 2025) 的时间效率高 **55×**，比 RegMix (Liu et al., 2024) 的时间效率高 **550×**，同时提供了显著的性能提升。  
(b) 后训练：在此设置中，**FastMix**（我们的）再次展示了最先进的性能和时间效率，以 **52×** 的时间成本降低优于 RegMix，并比 CLIMB 额外获得 **5.5** 个性能点。这说明了我们的方法在性能和时间成本之间实现了优越的权衡。

## 1 引言

大规模模型的性能 (Yang et al., 2024b; Dubey et al., 2024; Touvron et al., 2023; Hu et al., 2024) 关键依赖于训练所用的数据。尽管大规模、多样化的数据集推动了最新进展，但确定预训练 (Shukor et al., 2025) 和后训练 (Dong et al., 2023) 的最优数据混合仍然是一个重大挑战。

流行的诸如手动试错 (Yang et al., 2023; Tong et al., 2024) 或基于代理的方法 (Liu et al., 2024; Diao et al., 2025) 通常随着模型规模的增大而难以扩展。例如，基于代理的搜索方法如 RegMix (Liu et al., 2024) 和 CLIMB (Diao et al., 2025) 展示了强大的泛化能力和稳定性，但它们在搜索过程中需要训练大量代理模型。这导致了高昂的计算开销，使得随着模型和数据集的不断扩展，混合优化变得越来越不切实际。因此，核心问题是：我们如何高效地为大规模训练确定有效的数据混合？

我们通过 **FastMix** 应对这一挑战，这是一个新颖的框架，仅训练**一个代理模型**即可自动发现数据混合。**FastMix** 不依赖于预定义的启发式规则或资源密集型的模拟，而是联合优化混合系数和模型参数，从而在效率和可扩展性上显著优于先前方法。**FastMix** 的核心是将混合选择重新表述为一个带权重的**双层优化**问题（见公式 (2)）。具体地，我们展示了优化混合比例在数学上等价于在统一源采样下为每个源分配损失权重。这种重新参数化直接将混合系数嵌入到可微的迭代优化目标中，使得能够对混合和模型进行高效的、基于梯度的优化。为解决该优化问题 (Maclaurin et al., 2015; Franceschi et al., 2018)，**FastMix** 实现了一种近似迭代优化过程，交替进行：(i) 根据当前混合比例对数据采样，更新模型参数（内循环）；(ii) 通过基于梯度的优化器 (Kingma and Ba, 2014) 基于验证反馈更新混合比例（外循环）。

广泛的评估表明，**FastMix** 在预训练和后训练中均能跨模型规模和任务优化数据混合，以极低的计算成本优于基线方法（见图 1）。在预训练中，它仅用 **1.3 GPU 小时** 就在 14 个基准测试中取得了平均得分 48.2 和排名第一（其中 9 个最好），比 RegMix (Liu et al., 2024) 快 **550 倍**，比 CLIMB (Diao et al., 2025) 快 **55 倍**。在后训练（SFT）中，一个针对数学优化的混合能泛化到编程和 STEM-QA 任务，仅用 **2.2 GPU 小时** 就达到了 65.4 分（比次优高 5.5 分），而 CLIMB/RegMix 需要超过 115 GPU 小时。总体而言，**FastMix** 使得混合优化对下一代大型模型变得实用且可扩展。

## 2 相关工作

大型模型的快速进步 (Dubey et al., 2024; Touvron et al., 2023; Allal et al., 2024; Yang et al., 2023; 2024a) 在很大程度上依赖于来自不同来源数据的策略性混合，这些来源涵盖了语言 (Yang et al., 2023)、模态 (Gunasekar et al., 2023; Yang et al., 2024b) 和难度级别 (He et al., 2025)。这个**数据混合问题** (Ge et al., 2024) 不仅在预训练 (Shukor et al., 2025; Dubey et al., 2024; Yang et al., 2024b) 中提出了根本性挑战，在后训练 (Dong et al., 2023; Ming et al., 2025; Tong et al., 2024) 中也是如此。早期的实践主要依赖于手动启发式规则，这些规则缺乏标准化，并且常常无法在不同设置间泛化。最近，基于优化的方法 (Xie et al., 2024; Fan et al., 2023; Liu et al., 2024) 被引入以实现混合选择的自动化。

基于代理的方法 (Xie et al., 2024; Liu et al., 2024; Diao et al., 2025) 采用两阶段设计：先训练一个代理模型在候选混合下，然后用其性能推断最优采样比例。例如，DoReMi (Xie et al., 2024) 训练一个小型代理根据相对损失调整领域权重，然后将优化后的比例重用于训练更大的模型。RegMix (Liu et al., 2024) 通过训练数百个不同比例下的代理模型，对得到的混合-性能对拟合一个回归模型，并外推最优混合，从而扩展了这一思想。CLIMB (Diao et al., 2025) 通过迭代细化搜索区域提高了效率，减少了所需的代理模型数量。其他工作 (Ye et al., 2024; Shukor et al., 2025; Kang et al., 2024) 研究了跨尺度迁移：Shukor 等人 (2025) 提供了理论和实验证据，表明在小模型上找到的混合可以泛化到更大的模型，而 Ye 等人 (2024)；Kang 等人 (2024) 报告了混合比例与性能之间的函数关系。

相比之下，**动态方法** (Chen et al., 2024; Ming et al., 2025; Albalak et al., 2023) 通过实时调整混合来消除单独的搜索阶段。例如，IDEAL (Ming et al., 2025) 利用影响函数 (Koh and Liang, 2017) 估计每个领域对下游性能的贡献，并动态地重新平衡训练数据。

总体而言，像 RegMix 和 CLIMB 这样的基于代理的方法通常比动态方法获得更强、更稳定的性能，但计算成本很高。我们的方法 **FastMix** 在保持基于代理优化的可靠性的同时，将搜索时间从数百 GPU 小时减少到几乎一个 GPU 小时，实现了更高的效率和更强的泛化能力。

## 3 FastMix

### 3.1 通过重新参数化重新描述问题

#### 数据混合作为双层优化问题。

形式上，数据混合优化可以表述为一个双层优化问题。设 \(D = \{D_1, \dots, D_k\}\) 为数据源（或聚类）的集合，设 \(\alpha \in A \subset \mathbb{R}^k\) 为混合权重，其中可行集 \(A\) 是概率单纯形（\(\alpha_i \geq 0\) 且 \(\sum_{i=1}^k \alpha_i = 1\)）。给定混合 \(\alpha\) 和模型参数 \(w\)，训练目标为 \(\mathcal{L}_{\text{train}}(D, w \mid \alpha)\)。设 \(w^*(\alpha)\) 为在该 \(\alpha\) 下（近似）优化此训练目标得到的参数。目标是找到使验证损失最小化的混合权重 \(\alpha^*\)，即在 \(w^*(\alpha)\) 上评估的 \(\mathcal{L}_{\text{target}}(w) = \ell_{\text{val}}(V, w)\)：

\[
\min_{\alpha} \, \mathcal{L}_{\text{target}}\Big(w^*(\alpha)\Big) \quad \text{s.t.} \quad w^*(\alpha) = \arg\min_{w} \mathcal{L}_{\text{train}}\Big(D, w \mid \alpha\Big), \quad \sum_{i=1}^k \alpha_i = 1, \quad \alpha_i \geq 0.
\]
(1)

其中内循环旨在通过在给定混合权重 \(\alpha\) 的数据集上最小化训练损失来找到最优模型权重 \(w^*(\alpha)\)。外循环则旨在优化这些混合权重 \(\alpha\) 以最小化模型在目标任务上的最终损失。

虽然这种双层形式在概念上很有吸引力，但在实践中难以求解。关键在于处理混合权重 \(\alpha\)。与模型参数 \(w\)（可以高效地进行基于梯度的更新）不同，混合（采样）比例通常是不可微的，阻碍了端到端的反向传播。因此，从业者只能采用贪婪启发式或策略梯度（得分函数）更新来调整 \(\alpha\)。这些过程样本效率低，并且随着数据源数量的增加而扩展性差，使得混合搜索成为主要的计算瓶颈。

#### 可微形式。

通过一个简单的重新参数化，我们将原始的双层问题转化为一个数学上等价、完全可微的目标。关键思想是用均匀采样下的每个源的可微损失权重来替代按混合比例进行的随机采样，使得每个源的贡献通过其权重连续可控，从而得到以下形式：

\[
\min_{\alpha} \, \mathcal{L}_{\text{target}}\Big(w^*(\alpha)\Big) \quad \text{s.t.} \quad w^*(\alpha) = \arg\min_{w} \sum_{i=1}^k \alpha_i \mathcal{L}_{\text{train}}\Big(D_i, w\Big), \quad \sum_{i=1}^k \alpha_i = 1, \quad \alpha_i \geq 0,
\]
(2)

其中 \(\mathcal{L}_{\text{train}}(D_i, w)\) 表示模型在源 \(D_i\) 上的训练损失，该损失是在**均匀源采样**（每个源被选中的概率为 \(1/k\)）下计算的。内循环通过最小化来自 \(k\) 个不同数据领域的训练损失的加权和来找到最优模型权重 \(w^*(\alpha)\)。数据混合权重 \(\alpha_i\) 作为每个领域损失的权重。外循环则旨在优化这些比例 \(\alpha\) 以最小化模型在目标任务上的损失。这种重新参数化是关键：我们不再将混合比例视为不可微的采样概率，而是将其重新解释为缩放每个源损失的连续系数。因此，混合权重 \(\bm{\alpha} = (\alpha_1, \ldots, \alpha_k)\) 是完全可微的，适用于基于梯度的优化。标准的优化器（如 SGD 或 Adam）随后可以联合更新模型参数和数据权重，从而实现高效的端到端训练。

**等价性证明。** 设 \(D = \bigcup_{i=1}^k D_i\) 表示 \(k\) 个数据源（或聚类）的并集，设 \(\alpha = (\alpha_1, \dots, \alpha_k)\) 为混合权重，满足 \(\sum_i \alpha_i = 1, \alpha_i \geq 0\)。为了采样一个训练样本 \(x\)，首先从 \(i \sim \mathrm{Cat}(\alpha)\) 抽取一个源索引，然后从 \(x \sim D_i\) 采样。在该混合采样下的训练损失为：

\[
\mathcal{L}_{\text{train}}(D, w \mid \alpha) = \mathbb{E}_{i \sim \mathrm{Cat}(\alpha)} \mathbb{E}_{x \sim D_i} \big[ \ell(x, w) \big] = \sum_{i=1}^k \alpha_i \, \mathcal{L}_{\text{train}}(D_i, w),
\]
(3)

其中 \(\ell(x, w)\) 是每个样本的损失，而 \(\mathcal{L}_{\text{train}}(D_i, w) = \mathbb{E}_{x \sim D_i} [\ell(x, w)]\)...