论文速递基于聚类的迭代数据混合引导：优化数据混合以预训练大语言模型（LLM）的框架

论文：https://arxiv.org/abs/2504.13161
数据集：
https://huggingface.co/datasets/nvidia/ClimbLab
https://huggingface.co/datasets/nvidia/ClimbMix

识别最佳预训练数据混合仍然是一个具有挑战性的问题，尽管它对预训练性能有很大好处。为了解决这个问题，本文[1]提出了基于聚类的迭代数据混合引导（CLIMB），这是一个在预训练环境中发现、评估和细化数据混合的自动化框架。CLIMB在语义空间中嵌入和聚类大规模数据集，然后使用较小的代理模型和预测器迭代搜索最佳混合。

主要贡献：

提出一种嵌入驱动的数据混合方法：以自动识别、分组、混合高质量集群，实现高效的特定领域训练，同时消除对手动预定义特定领域数据的依赖。
引入迭代搜索过程：在整个训练过程中动态细化数据混合以优化多样性和域相关性，同时解决聚类和数据过滤中的扩展挑战。
创建了新的数据集：创建了一个经过过滤的1.2万亿令牌语料库，其中包含20个集群，作为数据混合研究的新平台；以及一个新的高质量4000亿令牌数据，用于高效预训练。

框架概述

如下图所示：

i）数据预处理

CLIMB首先通过嵌入原始数据并将其聚类到组中来预处理原始数据。
这些聚类作为搜索空间的基础，其中混合被定义为一组权重以组合不同的聚类。

ii）混合引导

CLIMB在迭代k中对 $n_k$ 个混合进行采样，在其中的一个子集上训练代理模型，并更新预测器以估计性能。
预测器会删减可能表现不佳的混合，因此只有最有希望的混合在后续迭代中进行完整的代理训练。

iii）最佳混合重量

通过逐步细化搜索空间和消除次优候选，CLIMB收敛到优化的数据混合，并平衡一般和特定领域的性能，而无需详尽的手动管理。

详细框架

i）数据预处理

遵循三个步骤：

a）文本嵌入

给定一个包含n个文档的大型原始数据集 $\hat{D} = {D_1, D_2, \ldots, D_n}$，使用嵌入模型 $M_e$ 将文档映射到嵌入空间。
嵌入模型的输出是一组嵌入向量，$E = {E_1, E_2, \ldots, E_n}$。

b）嵌入集群

使用聚类算法（如k-means）对嵌入进行聚类，用于分组到 $K_{init}$ 个集群中。
要创建尽可能细粒度的集群以供后续处理，在此阶段首选将 $K_{init}$ 设置为相对较大的值，例如1000。

c）集群合并

进行集群级修剪以去除低质量集群，保留基于基于模型的分类器修剪的 $K_{pruned}$ 个高质量集群作为修剪度量。
根据质心之间的距离将簇合并为 $K_{enhanced}$ 个增强簇，其中 $K_{enhanced} < K_{pruned} < K_{init}$。
整个数据集从 $\hat{D}$ 减少到 $D$。

ii）迭代引导：混合权重搜索

a）混合物重量搜索作为双层最佳化问题

给定一组数据簇 $D = {D_1, D_2, \ldots, D_k}$ 和目标函数 $\ell(\alpha, \omega)$，其中模型权重 $\omega$ 是用混合权重 $\alpha$ 训练的，该目标函数在校准集上输出可实现的性能 $P$，目标是识别最佳混合权重 $\alpha^* \in A$，以最大化任务性能 $\ell(\alpha, \omega)$。

b）使用任务绩效预测器近似目标

建议使用预测器 $f_{\theta}(\alpha)$ 基于（混合权重、性能）对的子集来近似 $\ell(\alpha, \omega)$，从而显著降低训练成本。

其中 $\mathcal{L}$ 是预测器 $f_{\theta}$ 的损失函数，$\widetilde{\mathcal{F}}$ 表示 $\ell$ 的所有可能近似值的集合，$S := {S \subseteq A | |S| \leq C}$ 表示满足抽样预算 $C$ 的所有配置。

c）迭代无融资创业解决上述方程

提出一种迭代方法来进化采样策略 $S$ 和预测器 $f_{\theta}$。
数学上表述为使用坐标下降法解决双层最佳化问题，该方法在优化配置采样和预测器拟合子程序之间交替，其中迭代 $k$ 可以表述为：

其中 $Top_N(\widetilde{P}_k)$ 表示一组前 $N$ 个配置，根据任务性能 $\widetilde{P}_k$ 进行排名。

下图中使用t-SNE可视化CLIMB的迭代搜索过程。每个点代表搜索空间中的数据混合配置，不同的迭代（CLIMB-Iter1、CLIMB-Iter2、CLIMB-Iter3）说明了搜索空间是如何在迭代中细化的。

iii）实施

通过从 $A$ 中随机采样几个配置来初始化 $S_1$ 并训练代理模型以获得它们的性能。然后，对于迭代 $k = 2, \ldots, K$，交替地联合优化采样集 $S_k$ 和预测器 $f_{k\theta}$。

a）子程序1：配置采样

在迭代 $k + 1$ 时，对权重空间 $A$ 中的所有配置（不包括那些已经在 $S_k$ 中的）根据它们的预测性能 $\widetilde{P}_k$ 进行排序。
接下来，从基于 $\widetilde{P}_k$ 排名前 $N$ 的配置中随机采样 $M$ 个新配置，以平衡开发和探索。
这些新采样的配置，结合 $S_k$，形成 $S_{k + 1}$。

b）子程序2：（弱）预测器拟合

通过最小化损失 $\mathcal{L}$，使用 $S_{k + 1}$ 中的采样配置来训练预测器 $f_{k + 1\theta}$。然后使用学习到的预测器 $f_{k + 1\theta}$ 来评估配置并生成预测的性能 $\widetilde{P}_{k + 1}$。
通过在这两个过程之间交替进行预定义数量的迭代，可以逐步细化预测器并将采样过程引导到具有更高质量混合权重的子空间，从而提高搜索混合权重的平均质量。
最后，选择最终预测器预测的最佳配置作为最终数据混合权重。

结果

i）与数据混合基线的比较

下表显示了CLIMB与数据混合方法的比较。

CLIMB实现了54.83%的平均准确率，优于随机（52.17%）和表现最好的基线Regmix（53.78%）。
同样，对于1B模型，CLIMB实现了60.41%的平均准确率，高于所有基线。

ii）与最先进的语言模型（SOTA LMs）的比较

下表显示了与通用推理基准上最先进的语言模型的比较。

CLIMB在所有子500M和子1.2B模型中取得了最佳性能。例如，当比较类似规模的模型（围绕1B参数）时，CLIMB在大多数一般推理基准上始终优于其他基线，包括Llama-3.2和AMD-OLMo。
CLIMB获得了最高的总体平均分数，以明显的优势（2.0%）超过了次优方法（即Llama-3.2）。

ClimbMix：新的预训练数据

将CLIMB应用于两个现有数据集：Nemotron-CC和smollm-corpus，目标是构建一个强大的新预训练数据集。
首先结合Nemotron-CC和smollm-corpus，然后采用我们提出的CLIMB聚类方法将这个组合的数据集语义重组和过滤成20个不同的聚类，从而产生一个1.2万亿令牌的高质量语料库，名为ClimbLab。
随后，利用CLIMB-search从这些集群中识别出最佳数据混合。使用这种最佳混合，我们进一步提取了一个名为ClimbMix的4000亿令牌高质量数据集。
从零开始使用ClimbMix训练一个1B模型，并评估其相对于在相同令牌预算下在其他数据集上预训练的模型的性能。
结果（如下图所示）表明，在ClimbMix上训练的模型显著优于在现有数据集上训练的模型。

CLIMB识别的最佳数据混合权重如下图所示。

结论

引入CLIMB，一种新的基于聚类的迭代混合无融资创业框架，用于优化预训练大语言模型（LLMs）的数据混合。
CLIMB自动发现、评估和细化数据混合，改进具有明确目标的大规模预训练。
通过利用无监督聚类、代理模型训练和预测器，CLIMB有效地导航数据组合的广阔搜索空间，无需依赖预定义的域标签或大量手动管理即可构建最佳域感知混合。
通过使用CLIMB搜索的最佳数据混合训练350M和1B模型，我们在12个推理任务中实现了最先进的性能。

参考文献：

CLIMB: CLustering-based Iterative Data Mixture Bootstrapping for Language Model Pre-training

主要贡献：

框架概述

i）数据预处理

ii）混合引导

iii）最佳混合重量

详细框架

i）数据预处理

ii）迭代引导：混合权重搜索

iii）实施

结果

i）与数据混合基线的比较

ii）与最先进的语言模型（SOTA LMs）的比较

ClimbMix：新的预训练数据

结论

推荐阅读