2024/09/27 10:20

用大模型优化大模型预训练数据，节省20倍计算量，实现显著性能提升！

AIxiv专栏是机器之心发布学术、技术内容的栏目。过去数年，机器之心AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com

生成式人工智能研究实验室（GAIR，主页：https://plms.ai/）是国内首个聚焦于生成式人工智能的高校研究组。汇聚了来自于 CMU、复旦、交大（ACM 班、IEEE 试点班等）等顶尖高校的年轻本硕博人才。实验室专注于三大核心领域：大模型基础研究、对齐系统和社会影响，致力于培养顶尖人工智能人才（具有原创、批判精神等）、开发尖端的生成式人工智能技术，赋能人类解决复杂问题，提升人类生活质量。

Program-Every-Example (ProX) 是上海交通大学生成式人工智能实验室 (GAIR Lab) 、上海人工智能实验室、新加坡 Sea AI Lab 联合研究的一个用于提升大规模预训练语料质量的数据优化框架。

1. 背景介绍

在大模型研发过程中，数据质量是决定模型成功与否的关键因素之一，而大规模预训练数据则是构建强大语言模型的基础。业界的实践通常是抓取互联网上的海量的网页数据，然后由人类专家设计专门的规则来进行数据清洗，这些规则虽然能够在处理数千亿甚至上万亿 token 级别的语料时保持较高的效率，但其局限性也逐渐显现。数据的清洗效果往往决于规则是否全面，然而这些规则并不能理解文档中的上下文语境，同时也很难为每一个文档去定制规则，故此面对多样化的样本需求，这些静态规则难以灵活应对，导致语料质量提升空间有限，制约了模型的进一步突破。

为此，上海交通大学生成式人工智能实验室 (GAIR Lab)、上海人工智能实验室以及新加坡 Sea AI Lab 联合提出了 Program-Every-Example (ProX) 框架。ProX 通过采用语言模型而非人类规则进行自动化的数据优化，灵活应对多样化的样本特征，显著提升了预训练语料的质量，为大规模预训练模型的性能突破提供了一种可能的数据优化思路

论文: https://huggingface.co/papers/2409.17115

团队还以开源的方式为社区提供了：

代码框架: https://github.com/GAIR-NLP/ProX
预训练实验的全实验细节和模型: https://huggingface.co/gair-prox
超过 100B 的高质量通用语料和 5B 的高质量数学语料: https://hf.co/collections/gair-prox/prox-dataset-66e81c9d560911b836bb3704

2. 用语言模型提升语言模型「自己」的数据质量

传统的预训练数据清洗和优化方法主要依赖人工设计的规则，虽然这些规则能够有效过滤低质量数据，但无法针对每个样本进行更细粒度的处理。这种方式难以处理复杂的多样化数据，最终可能导致模型性能的瓶颈。Program-Every-Example (ProX) 的出现旨在改变这一局面，它通过利用语言模型自动生成优化程序，能够对每个数据样本进行个性化处理，提升数据的整体质量。ProX 不仅可以识别低质量的文本，还能进行更精准的优化和清理，如字符串标准化、噪声行删除等操作，确保每个样本都能以最佳状态进入预训练过程。

^{图：ProX 总体框架，利用语言模型生成程序 + 执行，提升数据质量。}

3. 如何兼顾效率和质量？Program-Every-Example!

在利用语言模型提升数据质量的同时，如何提高数据处理的成本和效率也是研究中的一大挑战，特别是考虑到预训练所需要的语料往往突破了数万亿级别 (TB) token。

近期，业界和各个研究团队（如 meta 的 llama 团队，阿里千问团队）也在积极研究直接利用大语言模型来选择和过滤语料，这些方法主要应用于文档级别的数据过滤和选择，能够有效识别低质量的文档并加以排除，提升整体数据集的质量。然而，这类方法潜在的局限性在于其精细化处理能力不足，通常只停留在文档的选择层面，可能没有进一步对文档内的细节进行更深入的优化和清理，例如字符级别的规范化处理、噪声行的删除等操作，从而限制了数据优化的效果。

同时，直接的数据合成方法正在成为另一个研究热点。该方法通过大语言模型（LLM）基于种子数据或特定话题，主动生成新的高质量数据。这种合成技术尽管展示了生成多样化、丰富语料的潜力，但其挑战在于需要极大的计算资源支持，特别是在大规模生成时，计算代价非常高。此外，LLM 生成的数据还面临幻觉问题，即生成内容与事实不符，无法保证训练数据的真实性和一致性。这些问题给直接数据合成方法的实际应用带来了很大挑战，特别是在需要高精度和高可信度的领域。

于是在 ProX 中，研究者们提出了 Program-Every-Example 的设计。他们通过使用较小的语言模型（如 0.3B 的超小规模模型）来执行数据优化任务，并且将数据优化的过程转变为调用特定的 python 函数，对每个文档都生成特定的函数调用。ProX 框架会再通过执行这些程序，来对语料进行文档级别的过滤和噪声行、字符级别的处理。

这样简单的 Program 设计，大大节省了输出端的 token 数量，相比较于依赖于大语言模型数据合成方法（通常依赖于 7B 甚至 70B 规模的模型）显著减少了计算成本，同时又做到了简单过滤方法无法实现的精细化操作。

^{图：ProX 设计的函数接口，覆盖了文档级别过滤、噪声行去除、字符串清洗等操作。}

4. 如何构造数据优化模型？

如何构造这样的数据优化模型呢？是否要对不同 domain 都有不同的构造方法呢？在这里，ProX 采用了一种统一的构造思路，无需对特定领域进行额外设计就可以广泛适用。首先，针对数据优化的粒度，ProX 将优化分为两个阶段：文档级别和子文档 / 块级别 (将文档切分成更小的窗口)，分别对应文档级别的过滤和更细粒度的块内优化如行级别和字符级别的优化。

由于直接让模型，尤其是较小的模型生成这样上述设计的函数调用是较为困难的，因此 ProX 合成了对应的微调数据，以便对基座模型进行微调以处理和优化真实的语料。在文档级别操作中，ProX 借鉴了近期纽约大学的 Self-Rewarding 以及 Hugginface 的 FineWeb-Edu 中所采用的评分机制，对数万个随机采样的种子文档进行了教育价值评分、文档格式评分，并构建出保留高分文档、去除低分文档的优化程序。在子文档 / 块级别级别的操作中，ProX 利用了 Llama-3 和少样本提示的方法，构造了对于多种噪声的函数调用，并对这些离散的函数调用进行语法检查和正确性分析，筛选然后合并为一个较为复杂的程序。

最终，ProX 对两个优化阶段，分别构建了适用于通用领域的约 50K 函数调用数据集，并以相同的方法针对数学领域生成了同等规模的数据集。ProX 随后通过 SFT（监督微调） 在一个小模型上进行训练，模型快速收敛，且成功应用于真实的大规模语料处理任务中，在后续的大量实验中展示了其高效的优化能力与广泛的适用性。

^{图：ProX 构造 SFT 数据和数据优化模型的流程示意图。}

5. 在广泛数据集上得以验证

那么 ProX 得到的数据效果到底如何？

超越人工规则，超越此前的基于语言模型的数据选择方法

初步实验发现，ProX 在小于 1B 规模的模型训练中超越了此前多个人工设计的规则过滤方法，并且在多数的标准化测试数据集中稳定达到更高的表现。

^{表：ProX 和人工规则的对比，表现出稳定且大幅度的增长。}

ProX 还和此前利用语言模型选择预训练数据的方法进行了比较。团队仅采用了文档级别的优化，以保持和数据选择方法较为公平的优化粒度，仍然在 8 个下游数据集上，大幅度提升了之前的利用语言模型进行数据选择方法 ( > 2.0%)。

^{表：Prox 和现有的利用语言模型进行数据选择方法对比，超越了所有基准方法。}

更大规模的实验：50B 语料训练媲美 3T 性能？

团队在 1.7B 模型上训练了 50B ProX 优化后的数据，发现其表现已经接近于甚至超越一些用超过数百 B 甚至 TB 级别语料训练的语言模型，例如 OLMo、Pythia、TinyLlama 等等。且训练后的模型和现有的利用超大模型直接数据合成语料上训练更充足的模型，表现也非常接近。这充分说明了 ProX 优化数据质量的有效性。

^{图：ProX 和原始数据、Tinyllama (3T token)、OLMo (2T token)、Pythia (300B token) 的训练代价及性能表现对比。}

在这些对比的模型中，Inst-LM 实际采用了更大规模的经过重构的下游数据集 + 微调的 Mistral-7B 模型，合成了更多的 100B tokens；Cosmo 则采用了更大规模的 Mistral-8x7B 模型合成了 25B tokens，并重复训练了 180B 总 tokens；Sheared-Llama 则是通过模型剪枝算法，从 Llama-2-7B 中剪枝并且继续训练 50B tokens 得到。相比于 ProX 仅用了 0.3B 的优化模型和 50B tokens 的训练代价，这些模型的数据开发和训练成本都要大得多。而 ProX 取得了和他们十分接近的数据集表现，也进一步反应了 ProX 的高效和节省计算量。

^{图：ProX 和 LLama-2 模型剪枝方法、其他合成语料训练的模型之间的性能比较。}

通过对 FLOPs 的定量分析，团队发现随着训练模型参数量的扩大，ProX 的数据优化的推理成本占训练成本的比例在不断下降，达到相同下游表现的总训练成本也在大幅下降。并且相比于利用大模型，直接进行大规模数据合成高质量语料所需要的代价要小得多。

^{图：达到相同性能，使用 / 不使用 ProX 所需的总计算量成本。}

特定领域预训练：无需额外设计，大幅提升数学性能

ProX 在数学语料上也有相似的提升。ProX 对现有最好的开源数学文本语料 Open-Web-Math 进行优化，得到了约 5B 高质量的数据。在对 7B 模型 (如，CodeLlama，Llama，Mistral）经过了 10B～15B 的预训练后，模型在 9 个数学相关榜单上最高提升 20% 的性能。注意到 ProX 仅用了约 1/20 的训练代价，就得到了效果堪比现有的使用开源数据训练 50B 甚至 200B 的专用数学大模型，再次证明了数据质量对于模型的重要性。

值得注意的是，ProX 并未对数学领域的数据优化模型、优化函数、数据收集进行额外的设计；他们直接使用了和通用领域相同的提示词、函数实现、构建策略，就得到了出色的效果；这也极大降低了定制化垂类大模型的开发成本。

^{图：ProX 在数学语料上的续训练以及在 9 个数学任务上的平均性能表现。}

6. 未来展望

ProX 可以怎么样帮助社区做下一步的研究？他们在论文的最后还构想中了如下几点：

设计出更灵活的接口，以覆盖多样化的数据优化需求，进一步提升数据质量。
将 ProX 发展到更多专用 domain：例如代码领域、多语言语料上，促进领域大模型的发展。
通过进一步减小模型尺寸、和其他加速手段，减小推理代价，使清洗 TB 级别文档不再遥远。

最后，随着近期 openai o1 模型的发布，test time scaling 成为研究者们和业界关心的话题。对此，ProX 团队也怀有相同的期待：在数据优化阶段投入更多的计算量来优化数据，或许同样能够成为模型突破瓶颈、迈向通用人工智能（AGI）的重要技术手段之一。

工程Program-Every-Example