尽管优化预训练数据混合对大型语言模型(LLM)的性能有显著影响，但确定最优数据配比仍然是一个亟待解决的挑战。为应对这一问题，英伟达提出了一种名为 CLIMB(CLustering-based Iterative data Mixture Bootstrapping)的自动化框架，该框架能够在预训练环境中系统地发现、评估并优化数据混合策略。CLIMB 通过在语义空间中嵌入并聚类大规模数据集，并结合小型代理模型与性能预测器，迭代搜索最优数据混合比例。
技术创新点
CLIMB 框架主要贡献体现在以下三个方面：
- 提出了一种基于语义嵌入的数据混合方法，能够自动识别、分组并混合高质量的数据聚类，实现高效的领域特定训练，同时消除了对人工预定义领域标签的依赖。
- 设计了创新的迭代搜索机制，能够在训练过程中动态优化数据混合比例，平衡多样性与领域相关性，同时有效解决了数据聚类与过滤过程中的扩展性挑战。
- 构建了一个包含 20 个语义聚类、经过质量过滤的 1.2 万亿词元语料库，为数据混合研究提供了新的实验基础，并进一步提炼出一个高质量的 4000 亿词元预训练数据集。
框架整体架构
CLIMB 框架的整体结构如下图所示：
数据预处理阶段
该阶段主要通过嵌入和聚类技术将原始数据进行预处理与分组。这些聚类构成了后续搜索空间的基础，其中混合策略被定义为一组用于组合不同聚类的权重向量。
混合自举阶段
在第 k 次迭代中，CLIMB 从可能的混合配置中采样 nk 个候选混合方案，对其中一部分进行代理模型训练，并更新性能预测器以估计所有候选混合的效果。预测器能够筛选出可能表现不佳的混合方案，确保只有最具潜力的配置能够进入后续迭代的完整代理训练评估。
最优混合权重确定
通过逐步优化搜索空间并剔除次优候选方案，CLIMB 最终收敛到经过优化的数据混合配比，有效平衡了模型的通用性能与领域特定能力，无需繁重的人工管理过程。
技术实现细节
数据预处理流程
数据预处理分为三个关键步骤：
文本嵌入处理
针对包含 n 个文档的大型原始数据集 D̂ = {D₁, D₂, ..., Dₙ}，采用嵌入模型 Me 将每个文档映射到语义空间，生成对应的嵌入向量集合 E = {E₁, E₂, ..., Eₙ}。
嵌入聚类分析
利用 k-means 等聚类算法对嵌入向量进行聚类，将数据分组为 Kinit 个初始聚类。为确保后续处理的粒度足够细致，通常将 Kinit 设置为较大值（如 1000）。
聚类优化与合并
基于质量指标进行聚类级别的剪枝，移除低质量聚类，保留 Kpruned 个高质量聚类。随后根据聚类质心之间的距离将相似聚类合并为 Kenhanced 个最终聚类，其中 Kenhanced < Kpruned < Kinit。经过此过程，原始数据集 D̂ 被精简为优化后的数据集 D。
迭代自举：混合权重搜索
双层优化问题建模
CLIMB 将混合权重搜索视为双层优化问题：给定一组数据聚类 D = {D₁, D₂, ..., Dₖ}和目标函数 ℓ(α, ω)，其中模型权重 ω 使用混合权重 α 进行训练，该函数输出在校准集上的性能 P。目标是识别最优混合权重 α* ∈ A，以最大化任务性能 ℓ(α, ω)。
性能预测器近似目标函数
为降低计算成本，CLIMB 引入预测器 fθ(α)，基于有限的(混合权重，性能)对来近似 ℓ(α, ω)，大幅降低训练开销。
其中 L 是预测器 fθ 的损失函数，F̃ 表示 ℓ 的所有可能近似集合，S := {S ⊆ A | |S| ≤ C}表示满足采样预算 C 的所有配置集合。
迭代自举优化过程
CLIMB 提出一种迭代方法来同步优化采样策略 S 和预测器 fθ。数学上，这等效于使用坐标下降法解决双层优化问题，在优化配置采样和预测器拟合子程序之间交替迭代，第 k 次迭代可表述为：
其中 TopN(P̃k)代表根据任务性能 P̃k 排名的前 N 个配置集合。
下图通过 t-SNE 可视化了 CLIMB 的迭代搜索过程，每个点代表搜索空间中的一个数据混合配置，不同迭代阶段(CLIMB-Iter1、CLIMB-Iter2、CLIMB-Iter3)展示了搜索空间随迭代优化的演变过程。
实现算法
CLIMB 的实现始于从配置空间 A 中随机采样初始配置并训练代理模型以获取性能指标，初始化采样集 S₁。随后，在迭代 k = 2,...,K 中，交替优化采样集 Sₖ 和预测器 fᵏθ。
配置采样子程序
在迭代 k+1 时，根据预测性能 P̃k 对权重空间 A 中的所有未采样配置进行排序。为平衡探索与利用，从排名前 N 的配置中随机采样 M 个新配置，与 Sₖ 组合形成 Sₖ₊₁。
预测器拟合子程序
通过最小化损失函数 L，使用 Sₖ₊₁ 中的采样配置训练预测器 fᵏ⁺¹θ。然后利用更新后的预测器评估配置生成预测性能 P̃ₖ₊₁。
通过在预定迭代次数内交替执行这两个子程序，CLIMB 能够逐步优化预测器并将搜索过程引导至更高质量的混合权重子空间，提升搜索结果的整体质量。最终，选择经最终预测器评估的最佳配置作为数据混合的最终权重。
实验评估与结果
与现有数据混合方法的比较
下表展示了 CLIMB 与其他数据混合基线方法的性能对比：
实验表明，CLIMB 在 350M 模型上达到了 54.83%的平均准确率，明显优于随机混合策略(52.17%)和此前最佳的 Regmix 方法(53.78%)。同样，在 1B 模型规模上，CLIMB 的平均准确率为 60.41%，同样超越了所有基线方法。
与 SOTA 语言模型的性能比较
下表展示了 CLIMB 与当前最先进语言模型在通用推理基准上的性能对比：
在同等规模(约 1B 参数)的模型比较中，CLIMB 在大多数通用推理基准测试中都显著优于其他模型，包括 Llama-3.2 和 AMD-OLMo。总体而言，CLIMB 获得了最高的整体平均分，比排名第二的 Llama-3.2 高出 2.0 个百分点，这一差距在统计上具有显著意义。
ClimbMix：高质量预训练数据集构建
研究团队将 CLIMB 应用于 Nemotron-CC 和 smollm-corpus 两个现有数据集，旨在构建一个性能更强的预训练数据集。首先将这两个数据集合并，然后应用 CLIMB 聚类方法进行语义重组和质量过滤，将数据分为 20 个语义聚类，形成了一个包含 1.2 万亿词元的高质量语料库，命名为 ClimbLab。
在此基础上，研究者利用 CLIMB 搜索算法从这些聚类中识别出最优的数据混合比例。基于此最优混合策略，进一步提取了一个名为 ClimbMix 的 4000 亿词元高质量数据集。为验证其效果，研究团队使用 ClimbMix 从头开始训练了一个 1B 参数规模的模型，并在相同词元预算条件下与其他数据集训练的模型进行比较。
如下图所示，在 ClimbMix 上训练的模型表现显著优于在现有公开数据集上训练的同等规模模型。这一结果充分证明了 CLIMB 框架在优化预训练数据混合方面的有效性。
CLIMB 算法识别出的最优数据混合权重分布如下图所示，该分布反映了不同语义类别内容在最终混合中的理想比例：
总结
CLIMB 框架是一种创新的基于聚类的迭代式数据混合自举方法，专为优化 LLM 预训练数据混合而设计。CLIMB 通过自动化数据混合的发现、评估与优化过程，以明确的性能指标为目标改进了大规模预训练效果。
通过结合无监督聚类、代理模型训练和性能预测技术，CLIMB 能够高效地探索庞大的数据组合空间，无需依赖预定义的领域标签或大量人工干预即可构建最优的领域感知混合策略。实验结果表明，使用 CLIMB 优化的数据混合训练的 350M 和 1B 参数规模模型，在 12 项推理任务上均达到了当前最先进的性能水平。
论文：https://avoid.overfit.cn/post/a326f45fca2f4e90828526ed568ee620
