引言

想象你是一名厨师，任务是制作一道完美的汤。你的储藏室里装满了食材: 胡萝卜、土豆、牛肉、香料以及其他几十种东西。你知道，汤的成分构成——胡萝卜与土豆的确切比例——将决定它是美味可口还是平淡无奇。

在烹饪界，你可以尝一勺试试味道。但在大语言模型 (LLM) 的世界里，“烹饪”单一食谱 (在特定的数据混合上训练模型) 需要耗费数十万美元和数周的计算时间。如果你想测试 50 种不同的数据组合，你需要进行 50 次独立的训练。这就是数据消融 (data ablation) 的瓶颈。

随着 LLM 的规模越来越大，决定在什么数据上训练它们与架构本身一样重要。然而，由于测试每一种可能的数据组合极其昂贵，研究人员往往依赖直觉或有限的实验，这可能导致错失性能提升的机会。

在最近的一篇论文《Scalable Data Ablation Approximations for Language Models through Modular Training and Merging》中，来自艾伦人工智能研究所 (AI2) 、卡内基梅隆大学和华盛顿大学的研究人员提出了一种出色的变通方案。如果不去烹饪 50 种不同的汤，而是将每种食材单独烹饪，然后通过数学方法“混合”这些烹饪好的版本，来精确预测最终汤的味道，那会怎样？

这篇博客文章将探讨他们的模块化训练与合并 (Modular Training and Merging) 方法，这是一种允许研究人员以极低的计算成本模拟全面数据消融研究的技术。

图 1: 传统方法要求对每种数据混合都训练一个新模型。提出的模块化策略重用在数据分区上训练的模型，通过合并参数来近似完整训练的性能。

问题: 组合爆炸

要理解为什么这项研究是必要的，我们需要看看数据混合的数学原理。

假设你的训练语料库 \(\mathcal{C}\) 被分成了 \(n\) 个不同的分区 (例如: 维基百科、GitHub 代码、arXiv 论文、Reddit 帖子) 。你想知道这些分区的哪种组合能产生最好的模型。 如果你只是想包含或排除分区，那么可能的组合总数为 \(2^n\)。

  • 如果 \(n=10\)，你有 1,024 种组合。
  • 如果 \(n=30\)，你有超过 10 亿种组合。

即使你将自己限制在固定大小 \(k\) 的混合中，复杂度也是多项式级扩展的 (\(O(n^k)\))。在传统的“朴素”方法中，你必须为每一个你想评估的组合从头开始训练一个新模型。对于任何没有无限资源的人来说，这在计算上都是不可行的。

研究人员提出了一种线性扩展的方法，\(O(n)\)。你只需对每个分区训练一次，然后几乎可以瞬间近似出这些分区的任何组合的结果。

图 2: 计算复杂度对比。绿线显示了模块化近似方法的线性运行时间，而蓝线和橙线显示了传统训练的指数级和多项式级成本。

解决方案: 模块化训练与合并

核心创新依赖于一个称为参数平均 (parameter averaging) 的概念 (有时称为模型合并或 model soups) 。其假设简单而有力:

一个在数据 A 和数据 B 的混合物上训练的单一模型的困惑度 (Perplexity，即错误率) ，与两个分别在数据 A 和数据 B 上训练的独立模型的参数平均值的困惑度高度相关。

以下是作者提出的分步配方:

1. 数据分区

首先，训练语料库被分成“基本单元”或分区。这些可以按主题 (例如“生物学”、“计算机科学”) 或时间 (例如“2020年”、“2021年”) 来定义。

研究人员将此设置形式化如下: 表示 S2ORC 和 Wikipedia 数据分区并集的方程。

在他们的实验中，他们使用了两个主要数据集:

  1. S2ORC: Semantic Scholar Open Research Corpus (学术论文) 。
  2. M2D2 Wikipedia: 一个多领域的维基百科数据集。

他们确保数据经过了去污染处理 (根据评估集进行检查) ，以确保测试的公平性。

表 1: 使用的数据集统计信息，包括预训练和评估的 token 数量。

2. 模块化训练

研究人员没有在所有数据上训练一个巨型模型，而是训练了许多小型的“专家”模型。每个模型只在一个基本单元分区上进行训练。

关键在于，所有这些模型共享一个“种子”优化轨迹。它们都从同一个预训练检查点 (在古腾堡书籍和维基百科的通用语料库上训练了 420 亿个 token) 初始化，然后再分支在各自的分区上进行训练。这种共享的历史确保了模型保持在损失地形 (loss landscape) 的同一个“盆地”中，这对下一步的成功至关重要。

3. 合并 (近似方法)

为了模拟一个在分区 A 和分区 B 的混合物上训练的模型，研究人员并不训练新模型。相反，他们提取模型 A 和模型 B 的权重 (参数) 并对其进行平均。

如果分区的大小不同，他们会根据使用的 token 数量或分区数量执行加权平均 。

它真的有效吗？实验结果

研究人员通过比较三个特定指标来验证这种方法:

  1. SEQ (顺序训练) : “金标准”。实际在组合数据混合物上训练的模型。这很昂贵，但代表了我们想要预测的“真相”。
  2. IND (单独) : 简单地取各个独立模型困惑度分数的平均值。 (即，模型 A 得分 X，模型 B 得分 Y，预测结果为 \((X+Y)/2\)) 。
  3. MERGED (合并) : 拟议的方法。将模型 A 和模型 B 的权重平均到一个新模型中，然后测量该新模型的困惑度。

实验 1: 简单的成对和三元组

在第一组实验中，作者混合了成对 (\(k=2\)) 和三元组 (\(k=3\)) 的数据分区。他们训练了 1.3 亿参数的模型 (按今天的标准来看很小，但足以作为概念验证) ，并将 SEQ 结果与他们的代理指标进行了比较。

结果如下图所示，显示出惊人的相关性。

图 3: 显示 SEQ (真实) 困惑度与代理指标之间相关性的散点图。下图显示了 Paloma 数据集上的 OOD 评估。绿色三角形 (MERGED) 显示出与实际训练性能的强相关性。

关键结论: 看图 3 的底行 (在分布外数据 Paloma 基准上的评估) 。绿色三角形( MERGED )与 SEQ 分数形成了一条紧密的对角线。这意味着，如果合并模型认为某种数据混合很好 (低困惑度) ，实际的顺序训练模型也会有效证实这一点。

相关性得分证实了这种视觉直觉:

表 2: 皮尔逊相关系数得分。MERGED 模型在 OOD 数据上与 SEQ 模型实现了高相关性 (0.95+) 。

在 Paloma 评估集上, MERGED 方法在 S2ORC 对上实现了 0.961 的皮尔逊相关系数，在 Wikipedia 对上达到了 0.993 。 这表明你可以使用廉价的合并方法可靠地对数据混合进行排名。

实验 2: 不均匀的数据大小

现实世界的数据很少是完美平衡的。你可能有 10GB 的生物学文本，却只有 100MB 的诗歌。如何处理在数据量差异巨大的模型上的合并？

作者探索了两种针对不均匀分区的合并策略:

  1. 宏观合并 (Macro-Merging) : 在大分区上训练一个模型，在小分区上训练一个模型，然后平均它们。
  2. 微观合并 (Micro-Merging) : 将大分区分解成许多大小与小分区相等的“基本单元”。在所有单元上训练模型。将所有生成的基模型均匀地平均在一起。

图 4: 不均匀分区大小的 SEQ 与代理分数对比。Macro-MERGED 分数 (红色倒三角形) 被证明是非常可靠的代理指标。

图 4 的结果显示, Macro-MERGED 分数 (红色倒三角形) 对真实模型性能具有极高的预测性 (\(r=0.984\))。

然而，当混合不同的高级来源 (例如，将学术论文与维基百科文章结合) 时，策略略有变化。

图 5: 混合源数据混合的结果。在这里，Micro-MERGED 模型 (紫色倒三角形) 提供了最佳的相关性 (0.989) 。

在图 5 中，我们看到微观合并 (Micro-Merging) (将所有内容分解为小的、相等的单元并平均它们) 在混合不同领域时表现更优。这表明，维护一个在等大小数据块上训练的模型“池”是一种稳健的模拟策略。

扩展规模: 这适用于“真正”的模型吗？

对学术界 LLM 研究的一个常见批评是，在 1.3 亿参数模型上有效的方法可能会在数十亿参数的规模上失效。作者通过使用 11 亿 (1.1B) 参数的模型重复实验来解决这个问题。

他们发现了两个有希望的结果:

  1. 自我预测: 1.1B 参数的合并模型成功预测了 1.1B 参数顺序训练模型的性能。
  2. 跨规模预测: 更令人兴奋的是, 1.3 亿参数的合并模型可以预测 1.1B 参数顺序训练模型的排名。

图 6: 跨规模预测。小型 130M 模型的代理分数 (x 轴) 有效预测了大型 1.1B SEQ 模型的性能 (y 轴) 。

这种跨规模能力是一个巨大的效率胜利。这意味着你甚至不需要在完整的目标规模上执行模块化训练。你可以使用小型、廉价的模型运行数据消融搜索，找到最佳的数据配方，然后投入你的重型计算资源在该特定配方上训练大模型。

表 4: 相关性得分显示，130M 宏观合并模型 (0.926) 与 1.1B SEQ 模型性能高度相关。

表 4 中的数据支持了这一点，显示小型代理模型与大型顺序模型之间存在 0.926 的相关性。

特定领域的验证

研究人员还仔细研究了 S2ORC 语料库中的特定研究领域，以确保这一趋势在生物学、物理学和计算机科学等不同主题中都成立。

图 7: 各个研究领域的散点图。这种关系在各个领域都成立，增强了该方法的稳健性。

图 7 说明，无论你是评估数学还是生物学，代理 (合并) 与真相 (SEQ) 之间的关系都保持正向线性关系。

为什么这行得通？

论文触及了线性模式连通性 (Linear Mode Connectivity) 。 由于所有模块化模型都从同一个预训练“种子”开始，它们保持在同一个优化盆地中。当你平均它们的权重时，你不是跳到高维空间中的随机点；你是在寻找一个能有效捕捉两个父模型能力的点。这不同于集成 (ensembling，即平均输出) ；这是平均模型的实际“大脑”。

其他同期工作也注意到了这一现象，暗示了更广泛的数据混合“定律”。

图 10: 与 Ye 等人 (2024) 的同期工作比较，显示了从小模型到大模型预测带来的效率收益的兼容性。

启示与结论

“模块化训练与合并”方法本质上将数据消融从一个 NP-hard 问题转化为一个可管理的工程任务。

高效数据选择的配方:

  1. 分解你的海量训练语料库为等大小的块。
  2. 训练一个小型、高效的模型在每个块上 (并行进行) 。
  3. 缓存这些模型。
  4. 模拟数十亿种可能的数据混合，通过平均这些缓存模型的权重。
  5. 选择基于模拟表现最佳的混合。
  6. 训练你的最终、巨大的 LLM 在该特定混合上。

这种方法随新数据量的增加呈线性扩展。如果你下个月得到一个新的数据集，你不需要重新训练所有东西来看看它是否合适；你只需在新数据上训练一个模块，并将其“合并”到你现有的模拟中，看看它是否改进了配方。

图 11: 即使某些设置下的相关性较弱，该方法对于选择性能最佳的数据混合仍然有用。

虽然存在局限性——相关性不是完美的 (1.0)，且该方法依赖于共享的“种子”训练阶段——但效率的收益是不可否认的。对于计算预算有限的研究人员和公司来说，这项技术提供了穿越数据选择黑暗森林的地图，使他们能够找到最佳路径，而无需走遍每一条小径。

通过将模型视为可以堆叠和混合的模块化积木，我们更接近于数据组合的科学，而不是依赖直觉的炼金术。