Rewriting Pre-Training Data Boosts LLM Performance in Math and Code

大型语言模型（LLMs）在程序合成和数学推理方面的表现从根本上受限于其预训练语料库的质量。我们引入了两个开放许可的数据集，这些数据集在 Llama 3.3 社区许可下发布，通过系统性地重写公开数据，显著提升了 LLM 的性能。SwallowCode（约 161 亿 token）通过对 The-Stack-v2 中的 Python 代码片段进行四阶段处理来优化其质量：语法验证、基于 pylint 的风格过滤，以及一个两阶段的 LLM 重写过程，该过程确保代码风格一致，并将代码片段转化为独立且算法高效的示例。与以往依赖排除性过滤或有限转换的方法不同，我们的“转换并保留”方法能够升级低质量代码，最大化数据的使用价值。SwallowMath（约 23 亿 token）通过去除冗余内容、恢复上下文信息，并将解决方案重新格式化为简洁的分步解释，改进了 Finemath-4+ 数据集。在固定的 500 亿 token 训练预算内，使用 SwallowCode 对 Llama-3.1-8B 进行持续预训练，相较于 Stack-Edu，在 HumanEval 上的 pass@1 提升了 +17.0，在 HumanEval+ 上提升了 +17.7，超越了基线模型的代码生成能力。同样地，使用 SwallowMath 替代原有数据集后，在 GSM8K 上的准确率提高了 +12.4，在 MATH 上提高了 +7.6。消融研究表明，每一步处理流程都带来了增量改进，其中重写步骤贡献了最大的提升。所有数据集、提示模板和检查点均已公开，这有助于可复现的研究，并推动 LLM 在专业领域预训练技术的发展。

论文试图解决大型语言模型（LLMs）在代码生成和数学推理任务中的性能瓶颈问题，认为其根本原因在于预训练数据的质量不足。这是一个已知但尚未完全解决的问题。

论文提出了一种‘transform-and-retain’方法，通过系统性重写公共数据集来提升LLMs的性能。具体来说，SwallowCode和SwallowMath分别针对代码和数学领域进行了多阶段的数据优化，包括语法验证、风格过滤以及基于LLM的重写过程。这种方法不仅保留了原始数据的价值，还显著提升了数据质量，这是对传统排除式过滤方法的改进。

1. 提出了两个高质量的数据集：SwallowCode（约161亿tokens）和SwallowMath（约23亿tokens），均公开发布；2. 在固定50亿token的训练预算下，使用SwallowCode和SwallowMath进行持续预训练显著提升了LLama-3.1-8B模型的表现，例如HumanEval上的pass@1提高了+17.0，GSM8K上的准确率提高了+12.4；3. 每个数据处理阶段都经过了消融研究验证，证明了各步骤的重要性；4. 所有数据集、提示模板和检查点均已开源，便于复现和进一步研究。

相关研究包括：1. The-Stack-v2和Finemath-4+等现有代码和数学领域的数据集；2. CodeParrot和PolyCoder等专注于代码生成的LLMs；3. MATH和GSM8K等数学推理基准测试的研究；4. 近期关于数据清洗和增强对LLMs影响的工作，如《Improving Code Generation via Program Synthesis Data Augmentation》和《Enhancing Mathematical Reasoning with Curated Datasets》。