Arctic-SnowCoder揭秘：小数据如何炼成高性能代码模型？

一、结论写在前面

论文标题：Arctic-SnowCoder: Demystifying High-Quality Data in Code Pretraining

论文链接：​​https://arxiv.org/pdf/2409.02326​​

高质量数据对于语言模型的有效预训练至关重要。然而，“高质量”的精确定义仍未得到充分探索。

聚焦于代码领域，论文引入了Arctic-SnowCoder-1.3B，这是一个数据高效的基础代码模型，通过三个阶段的逐步精炼数据进行预训练，共处理了555B token：(1) 使用500B个标准质量代码token进行通用预训练，经过基本过滤、去重和去污染预处理；(2) 使用50B个高质量token进行持续预训练，这些token从第一阶段中通过BERT风格的质量注释器选出，该注释器经过训练以区分优质代码与随机数据，使用从高质量代码文件中提取的正例，以及来自Magicoder和StarCoder2-Instruct的指令数据；(3) 使用5B个合成数据进行增强预训练，这些数据由Llama-3.1-70B使用第二阶段数据作为种子生成，采用Magicoder的预训练方法。

仅仅训练有限数据集，Arctic-SnowCoder-1.3B在BigCodeBench上取得了SOTA的结果，这是一个专注于实际和具有挑战性的编程任务的编码基准，在训练了小于等于 1T tokens的类似大小的模型中。特别是，它比Phi-1.5-1.3B [20]高出36%。尽管训练数据量为555B tokens，与其他在万亿tokens上训练的最先进的小型代码模型相比，Arctic-SnowCoder在多个基准测试中达到了或超过了这些模型的性能。

例如，Arctic-SnowCoder-1.3B在所有评估的基准测试中击败了StarCoderBase-3B，后者训练数据超过1T tokens。Arctic-SnowCoder-1.3B在HumanEval+ （28.0 vs. 27.4）上优于StarCoder2-3B，后者训练数据超过3T tokens，这是一个评估函数级代码生成的基准，同时在BigCodeBench（19.4 vs. 21.4）上保持竞争力。论文进行了全面的消融研究，以验证训练Arctic-SnowCoder背后的设计决策：    

•首先，论文的研究结果表明，在一般预训练中，将文件级数据按编程语言划分成仓库后，显著优于仅按仓库名称分组数据的方法。

•此外，论文确定了最佳的学习率计划，包括重新预热阶段后进行线性衰减，以及在持续预训练期间高质量数据的理想重复次数，论文发现是四次。

•更重要的是，论文对基于模型的质量注释器的比较，这些注释器在各种数据组合上进行训练，强调了预训练数据与下游任务的对齐对于实现卓越性能至关重要。

图1：Arctic-SnowCoder-1.3B的三阶段预训练，逐步使用更高质量的数据。

二、论文的简单介绍

2.1 论文的背景

预训练LLMs通常依赖于大量数据。在代码等专业领域，这种对数据量的重视尤为明显，研究人员通过抓取GitHub等平台获取大量代码预训练数据集。然而，最近的研究越来越表明，高质量的数据对于有效的预训练至关重要, 包括代码领域。

在通用领域，研究人员探索了多种技术来策划高质量的语言模型预训练数据。FineWeb-Edu使用基于Snowflake-arctic-embed-m嵌入构建的线性回归器来评估网页的教育价值并选择高质量内容，而DCLM方法则采用基于fastText的过滤器，该过滤器在高质量在线资源和指令数据的正例以及随机负例网页上进行训练，以识别高质量文本。与使用未过滤的大规模数据集相比，这些基于模型的质量过滤器显著提高了语言模型在下游任务中的性能。

同样，研究人员也认识到高质量代码数据对于预训练的重要性，Phi-1使用随机森林分类器在Code-Gen嵌入上选择教育性代码样本，DeepSeek-Coder-V2采用多阶段fastText管道来召回与网络相关的代码数据和GitHub上的高质量代码，实现了最先进的编码性能。    

在本文中，论文介绍了Arctic-SnowCoder-1.3B，这是一个高性能的小型代码模型，通过一种新颖的三步训练方法创建，专注于数据质量的逐步提升。得益于这种方法，Arctic-SnowCoder-1.3B在所有评估的基准测试中均优于StarCoderBase-3B [19]，并在复杂且实用的BigCodeBench基准测试 [46] 上超越Phi-1.5-1.3B [20] 36%，该基准对实际编程至关重要。

2.2 论文的方法--Arctic-SnowCoder

这里论文将详细解释Arctic-SnowCoder-1.3B的训练方法，如图1所示。论文首先讨论原始训练数据的组成（见图1），然后概述通用预训练阶段。接下来，论文描述使用高质量数据的协同预训练过程，最后，论文详细阐述使用合成数据的增强预训练。模型架构基于Llama-2，具体细节见表1。

表1：Arctic-SnowCoder的模型架构细节。

2.2.1原始数据

用于训练Arctic-SnowCoder-1.3B的原始预训练数据仅包含代码，主要来源于用于训练Snowflake Arctic的编码数据。该数据结合了The Stack v1和GitHub爬取数据的清洗后版本。从这些数据中，论文选择了18种流行的编程语言进行训练，类似于StarCoder2-3B。这些语言包括Python、Java、C++、C 、JavaScript、PHP、C#、Go、TypeScript、SQL、Ruby、Rust、Jupyter Notebook、Scala、Kotlin、Shell、vart、Swift，总计400B唯一token。

2.2.2 通用预训练

在通用预训练阶段，模型使用Adam进行训练，训练500B token，序列长度为8192，批量大小为512。学习率在600次迭代后进行线性预热，随后遵循余弦衰减。论文设置最大学习率为5.3x 10^{-4}，最小学习率为5.3 x 10^{-5}，这一设置参考了DeepSeek-Coder。在此阶段，论文使用了全部400B原始数据，未进行额外的质量过滤。论文首先按编程语言对代码文件进行分区，按仓库进行分组，然后以随机顺序连接它们，类似于StarCoder2的方法。论文展示了首先按编程语言对代码文件进行分区的优势。论文将此阶段生成的模型命名为Arctic-SnowCoder-alpha。    

2.2.3 高质量数据继续预训练(Continued pretraining)

在通用预训练之后，论文使用从相同原始预训练语料库中提取的50B高质量token继续预训练Arctic-SnowCoder-alpha。这50B高质量token是通过将12.5B个由论文的代码质量注释器评分的前百分位代码文件token重复4次形成的。

受FineWeb-Edu和DCLM的启发，论文在基于BERT的先进嵌入模型Snowflake-arctic-embed-m之上训练了一个线性分类头。训练数据包括30万个正样本，采样自22万个高质量开源代码文件、8万个来自Magicoder和StarCoder2-Instruct的高质量指令数据，以及从预训练语料库中随机选择的300个代码文档。

关于代码质量的先前研究，如Phi-1，往往过分强调代码的“教育价值”，使模型偏向于像HumanEva这样的简单基准。论文展示了论文的注释方法能够带来更平衡的模型能力提升。

此外，鉴于这些代码文档通常超过1000个token，超过了BERT的512个token的上下文窗口大小，论文改进了FineWeb-Edu的流程，通过平均质量注释器产生的顶部、中部和底部部分的分数来计算每个文件的分数。在此阶段，论文从0到最大预训练学习率5.3x 10^{-4}进行1000次迭代的学习率预热，随后进行线性衰减至0。此阶段生成的模型称为Arctic-SnowCoder-beta。

2.2.4 使用合成数据增强预训练

在增强预训练阶段，论文利用Llama-3.1-70B-Instruct 生成比持续预训练阶段更高质量的数据，并将Python混合比例提高到约50%，同时保持其他语言的比例不变。Phi-1 [13]表明，类似教科书的合成预训练数据可以显著提升模型性能。

然而，过度依赖此类数据可能会导致模型分布偏斜，从而可能损害其在实际编码任务中的有效性。例如，论文后面展示的，Phi-1.5在HumanEvalt和MBPP+上表现出色，这些任务类似于教科书练习，但在BigCodeBench [46]中更复杂和实用的编码任务上表现较差。为了解决这个问题，论文改编了Magicoder的OSS-Instruct方法用于预训练。最初，OSS-Instruct旨在通过提示模型创建受开源代码片段启发的问答对来生成现实的指令调优数据。

相比之下，论文通过使用Llama-3.1-70B-Instruct生成高质量和面向问题解决的代码文件来生成高质量的合成预训练数据，这些代码文件以持续预训练阶段中评分最高的代码文档为种子。后面论文展示了每个预训练阶段都显著优于前一个阶段，突显了逐步提高数据质量的有效性。

2.3 论文的效果

这里论文将Arctic-SnowCoder与最先进的小型语言模型进行比较，并展示了每个预训练阶段的性能提升，评估了两种形成通用预训练中仓库级别数据的策略，并对持续预训练中的设计选择进行了详细的消融分析。    

2.3.1 实验设置

论文考虑以下四个多样化的编程基准，以全面评估不同代码模型的代码生成能力：

•HumanEval+和MBPP+ 。HumanEval和MBPP是用于函数级代码生成的两个最广泛使用的基准。论文采用了EvalPlus增强的版本，提供了80倍/35倍的更多测试用例以进行严格评估。HumanEvalt和MBPP+分别包含164和378个编码问题。

•EvoEval 是一个程序合成基准测试套件，通过将现有基准测试演化为不同的目标领域而创建。论文采用了其五个默认的转换类别，即困难、创造性、微妙、组合和工具使用，总计500个任务。

•BigCodeBench 通过实际且具有挑战性的编程任务评估语言模型。它包含1140个编程任务，每个任务都是通过人类与语言模型的协作创建的，任务质量由人类专家保证。        

2.3.2 基线比较与三阶段预训练的有效性

表 2 ：将Arctic-SnowCoder与最先进的小型语言模型 ( \mathrm{< 3 B} ) 进行比较，按训练计算量 > 1T 标记进行划分。Arctic-SnowCoder-alpha和Arctic-SnowCoder-beta分别是通用预训练和继续使用高质量数据预训练后的检查点。Arctic-SnowCoder是使用合成数据增强预训练后的最终检查点。

表2展示了多种小型语言模型（参数少于3B）在多个编码基准上的综合比较，按其训练计算是否超过1T tokens进行分类。值得注意的是，Arctic-SnowCoder表现出色，特别是在其有限的训练数据下。Arctic-SnowCoder-1.3B在BigCodeBench上达到了与训练数据不超过1T token的同类模型相比的最新性能，显著优于StarCoderBase-3B、SmolLM-1.7B和Phi-1.5-1.3B。特别是，尽管Phi-1.5-1.3B在“教科书式”基准测试如HumanEval+、MBPP+和EvoEval上具有优势，但Arctic-SnowCoder-1.3B在更复杂和实用的Big-CodeBench上以36%的优势超越了Phi-1.5-1.3B。    

此外，Arctic-SnowCoder-1.3B在所有评估基准上均击败了StarCoderBase-3B，后者是StarCoder2-3B的前身，训练数据为1T tokens。尽管仅训练了555B tokens，Arctic-SnowCoder-1.3B在HumanEvalt上与经过更广泛训练的模型如StarCoder2-3B、StableCode-3B、CodeGemma-2B-v1.0和Qwen1.5-1.8B相媲美甚至超越。在EvoEval和BigCodeBench上，Arctic-SnowCoder仍然具有竞争力。

此外，该表还突显了Arctic-SnowCoder在其训练阶段的持续改进：Arctic-SnowCoder-alpha、Arctic-SnowCoder-beta和最终的Arctic-SnowCoder。每个阶段都建立在前一阶段的基础上，Arctic-SnowCoder在所有基准测试中均取得了最高分数。这种稳步提升强调了高质量和合成数据在最终阶段的关键作用。尽管从相同的数据开始，Arctic-SnowCoder的每次迭代都缩小了与最先进模型的差距，展示了整体训练方法的有效性。

2.3.3 通用预训练中的仓库级数据

在通用预训练阶段，论文采用了StarCoder2的方法，通过随机拼接文件内容将文件级数据随机分组到仓库中。在表3中，论文研究了两种方法：（1）仅按仓库名称对文件进行分组，这意味着每个训练文档可以是多语言代码文件的混合，如果仓库是用不同语言编写的；（2）在将文件分组到仓库之前，先按编程语言对文件进行分区，这意味着每个训练文档仅关注一种单一语言。

表3：两种预训练方法对仓库级别数据分组的比较。（1）“按仓库分组”将每个仓库视为一个可能混合多种语言的单一训练单元，以及（2）“按语言和仓库分组”在按仓库分组之前先按编程语言对数据进行分区。

论文可以观察到，第二种方法，即论文在一般预训练中最终采用的方法，明显优于第一种方法。

2.3.4继续预训练中的设计选择

在继续预训练中，论文从预训练语料库中提取高质量的token，并训练一个改进的基础模型。为了获得高质量的token，论文采用了基于模型的质量标注器。在本节中，论文实验了各种设计选择，包括标注器的训练数据、继续预训练中使用的学习率以及高质量token的最佳重复次数。    

基于模型的质量标注器 类似于FineWeb-Edu ，论文在Snowf lake-arctic-embed-m嵌入模型之上训练一个线性头来为每个代码文件评分。在表4中，论文实验了4种变体：

•ANN-EDU：论文提示Mixtral-8x7B-Instruct标注每个代码文件的教育价值（1到5）。使用400k标注数据训练一个线性回归头。对于以下变体，类似于DCLM ，论文随机采样负文档并仅更改正部分。使用线性分类头。

•ANN-INs：正样本是来自ANN-EDU的100k教育数据（3.5+）和来自Magicoder [41]和StarCoder2-Instruct [40]的100k高质量指令数据的混合。

•ANN-HQ: 正样本为220k开源、合成、高质量代码文件。

•ANN-HQINs: 正样本为220k ANN-HQ训练数据与80k来自Magicoder[41]和StarCoder2-Instruct[40]的指令数据的混合。

表4: 通过应用基于模型的质量标注器（使用不同配方训练）进行10B继续预训练的下游性能比较

训练标注器后，论文首先将每个标注器应用于整个预训练语料库，为每个文件获取一个分数。与仅扫描前2k字符的FineWeb-Edu不同，论文扫描代码文件的顶部、中部和底部部分，并平均这些分数。然后，论文根据这些分数按语言对代码文件进行排名，并选择前百分位的文档，直到达到大约10 B token。论文保持与预训练中使用的相同混合比例。表中显示，结合高质量文件和指令数据的ANN-HQINS实现了最佳的下游性能。

论文在图2中进行了额外的分析。对于每个标注者，论文创建了一个验证数据集，其中正样本来自代码解决方案基准，负样本来自训练期间未见过的随机预训练数据。论文使用ROC-AUC [6]（受试者工作特征曲线下面积）分数来评估标注者在基准数据排名中的表现。该图展示了每个基准的ROC-AUC分数与基准通过率之间的相关性。几乎一致的趋势是：较高的ROC-AUC分数导致更好的基准性能。良好的ROC-AUC分数表明标注者有效地塑造了下游任务的分布。因此，高质量的关键在于与下游应用分布的对齐。    

图2：注释者ROC-AUC评分与基准测试pass@1之间的相关性。

学习率调度 论文还在表5中研究了不同的学习率调度策略，包括（1）从最小预训练学习率线性退火至零，（2）使用最小预训练学习率的恒定调度，以及（3）重新预热至最大预训练学习率后线性衰减至零。根据经验，论文发现重新预热方法表现最佳，并在所有其他关于继续预训练的实验中一致使用。

高质量数据重复次数 最后，论文将预训练的token范围从10 B扩展到50 B。剩下的一个问题是如何确定高质量token的最佳重复次数。论文通过选择由ANN-HQINS排名的前百分位token进行实验，重复次数从1到5，如表6所示。在这种情况下，前百分位token是最高质量的token。例如，1 x 50B表示前50B token的一次重复，而4 X12.5B表示前12.5B token的四次重复，确保所选token的质量最佳。

根据表中的结果，重复高质量token四次（4 x 12.5B）在下游多个评估指标中任务中表现最佳，四次重复（4 x 12.5B）在HumanEval和EvoEval中得分最高。两次重复( 2 X25.0B )和三次重复( 3 X16.7B )也表现出强劲的性能，特别是在mbpp中。五次重复( 5 x10.0B )在MBPP中得分最高，但在总体指标上未超过四次重复。一次重复( 1 x50.0B )与多次重复相比，改进最小。

表6：使用ANN-HQINS在50B继续预训练中不同高质量数据重复次数的下游性能。    

本文转载自 ​​AI帝国​​，作者： 无影寺