Source-linked AI summary

Arctic-SnowCoder: Demystifying High-Quality Data in Code Pretraining

Yuxiang Wei, Hojae Han, Rajhans Samdani

arXiv:2409.02326v1cs.CLcs.AI

TL;DR

Code pretraining has emphasized large crawled datasets, while recent work points to data quality as important. Arctic-SnowCoder applies three stages of progressively refined data and reports strong results against small code-model baselines, including on practical coding tasks.

  • Problem

    Code pretraining has relied on vast crawled datasets, while recent studies also identify data quality as important.

    Source: §1 · p. 2
  • Method

    Arctic-SnowCoder uses three progressive pretraining stages, including annotator-selected high-quality data and synthetic data seeded from earlier training data.

    Source: §1 · p. 2
  • Results

    Arctic-SnowCoder-1.3B beats StarCoderBase-3B across evaluated benchmarks and outperforms Phi-1.5-1.3B by 36% on BigCodeBench.

    Source: §1 · p. 2
  • Takeaways & Limitations

    The comparisons indicate that alignment between pretraining data and downstream tasks is crucial for superior performance.

    Source: §1 · p. 3

Abstract

from arXiv · show

Recent studies have been increasingly demonstrating that high-quality data is crucial for effective pretraining of language models. However, the precise definition of "high-quality" remains underexplored. Focusing on the code domain, we introduce Arctic-SnowCoder-1.3B, a data-efficient base code model pretrained on 555B tokens through three phases of progressively refined data: (1) general pretraining with 500B standard-quality code tokens, preprocessed through basic filtering, deduplication, and decontamination, (2) continued pretraining with 50B high-quality tokens, selected from phase one by a BERT-style quality annotator trained to distinguish good code from random data, using positive examples drawn from high-quality code files, along with instruction data from Magicoder and StarCoder2-Instruct, and (3) enhanced pretraining with 5B synthetic data created by Llama-3.1-70B using phase two data as seeds, adapting the Magicoder approach for pretraining. Despite being trained on a limited dataset, Arctic-SnowCoder achieves state-of-the-art performance on BigCodeBench, a coding benchmark focusing on practical and challenging programming tasks, compared to similarly sized models trained on no more than 1T tokens, outperforming Phi-1.5-1.3B by 36%. Across all evaluated benchmarks, Arctic-SnowCoder-1.3B beats StarCoderBase-3B pretrained on 1T tokens. Additionally, it matches the performance of leading small base code models trained on trillions of tokens. For example, Arctic-SnowCoder-1.3B surpasses StarCoder2-3B, pretrained on over 3.3T tokens, on HumanEval+, a benchmark that evaluates function-level code generation, and remains competitive on BigCodeBench. Our evaluation presents a comprehensive analysis justifying various design choices for Arctic-SnowCoder. Most importantly, we find that the key to high-quality data is its alignment with the distribution of downstream applications.

1 Introduction

Code pretraining has often emphasized large datasets, while recent work highlights the importance of data quality. Arctic-SnowCoder uses progressively refined data and reports strong benchmark performance, with data alignment and training design choices among its findings.

1 Introduction

  • Prior code-quality approaches include embedding-based selection and multi-stage filtering, while model-based filters have improved downstream performance over unfiltered data.The cited examples include Phi-1 and DeepSeek-Coder-V2.Source: §1 · p. 2
  • The three-stage method progressively refines training data, and Arctic-SnowCoder is trained on 555B tokens.The stages use general code data, high-quality code selected by an annotator, and synthetic code seeded from phase-two data.Source: §1 · pp. 2–3
  • Arctic-SnowCoder-1.3B beats StarCoderBase-3B across evaluated benchmarks and exceeds Phi-1.5-1.3B by 36% on BigCodeBench.It also surpasses StarCoder2-3B on HumanEval+ (28.0 vs. 27.4) while remaining competitive on BigCodeBench (19.4 vs. 21.4).Source: §1 · p. 2
  • Partitioning files by programming language before grouping them into repositories significantly outperforms grouping solely by repository name.Source: §1 · p. 2
  • The reported continued-pretraining choices are a re-warmup followed by linear decay and four repetitions of high-quality data.Source: §1 · p. 2
  • Comparisons of quality annotators indicate that alignment between pretraining data and downstream tasks is crucial for superior performance.The annotators were trained on different data combinations.Source: §1 · p. 3

2 Arctic-SnowCoder

Arctic-SnowCoder is trained on code from a shared raw corpus through general, continued, and synthetic-data pretraining. The stages use language- and repository-organized code, annotator-selected high-quality files, and synthetic code generated from selected seeds.

2 Arctic-SnowCoder

2.1 Raw data

  • The raw corpus combines cleaned The Stack v1 and GitHub crawls and contains code in 18 popular programming languages.The passage reports 400B unique tokens.Source: §2.1 · p. 3

2.2 General pretraining

  • General pretraining uses the 400B-token raw corpus, partitioning files by language before grouping them by repository and concatenating them in random order.Training runs for 500B tokens; the resulting model is Arctic-SnowCoder-alpha.Source: §2.2 · p. 3

2.3 Continued pretraining with high-quality data

  • Continued pretraining uses 50B tokens by repeating 12.5B top-percentile code-file tokens four times, selected by a BERT-based quality annotator.The annotator is trained with positive code-file and instruction examples alongside random code documents; file scores average the top, middle, and bottom sections.Source: §2.3 · p. 4

2.4 Enhanced pretraining with synthetic data

  • Enhanced pretraining uses Llama-3.1-70B-Instruct to generate problem-solving code files from top-percentile continued-pretraining documents.The stage increases the Python mix to approximately 50% and uses an adaptation of OSS-Instruct for pretraining.Source: §2.4 · p. 4

3 Experiments

Experiments compare small code models and evaluate choices in repo-level data grouping and continued-pretraining design. The evaluation uses code-generation benchmarks and reports pass@1 with greedy decoding.

3.1 Experimental setup

  • The experiments use HumanEval+, MBPP+, EvoEval, and BigCodeBench for baseline comparisons, reporting pass@1 with greedy decoding.Ablations use the base HumanEval and MBPP benchmarks and omit BigCodeBench for faster evaluation.Source: §3.1 · pp. 4–5

3.2 Baseline comparison and effectiveness of three-stage pretraining

  • The baseline comparison covers models below 3B parameters, divided by whether their training compute exceeds 1T tokens, and includes three Arctic-SnowCoder phase checkpoints.Alpha follows general pretraining, beta follows continued pretraining with high-quality data, and the final checkpoint follows enhanced pretraining with synthetic data.Source: Table 2 · p. 5

3.3 Repo-level data in general pretraining

  • Partitioning files by programming language before repository grouping significantly outperforms grouping files solely by repository name.The adopted language-first approach makes each training document focus on one language.Source: Table 3 · p. 6, §3.3 · p. 6

3.4 Design choices in continued pretraining

  • ANN-HQINS achieves the best downstream performance among the tested quality annotators in 10B continued pretraining.The table compares annotators trained with different data recipes.Source: Table 4 · p. 6, §3.4 · p. 7
  • Higher annotator ROC-AUC scores almost consistently track better benchmark pass rates, linking quality-filter performance to downstream-task alignment.The analysis evaluates annotators on benchmark solutions against unseen random pretraining data.Source: Figure 2 · p. 7, §3.4 · p. 7
  • A re-warmup to the maximum pretraining learning rate followed by linear decay performs best in continued pretraining.The authors use this schedule consistently in their other continued-pretraining experiments.Source: Table 5 · p. 7, §3.4 · p. 7
  • Four repetitions of the top 12.5B high-quality tokens yield the best overall downstream performance in 50B continued pretraining.This repetition setting has the highest HumanEval and EvoEval scores; five repetitions have the highest MBPP score but not the best overall metrics.Source: §3.4 · pp. 7–8, Table 6 · p. 8

4 Related Work

Code pretraining commonly starts from large crawled corpora that require filtering and deduplication. Prior work uses model-based quality filters, while Arctic-SnowCoder applies quality annotation and synthetic-data generation to code pretraining.

  • Large code corpora such as The Stack v1 and v2 are crawled from public repositories, then processed with deduplication, PII redaction, benchmark decontamination, or model-based filtering.Arctic-SnowCoder's corpus contains approximately 400B unique tokens from filtered The Stack v1 and GitHub crawls.Source: §4 · p. 8
  • Prior model-based filters include Phi-1's random forest, FineWeb-Edu's annotation-trained regressor, DCLM's fastText filter, and multi-stage filtering pipelines.The cited work describes filters used for educational value, high-scoring text selection, and code or math content retrieval.Source: §4 · p. 8
  • Arctic-SnowCoder uses a high-quality code annotator to select pretraining data and generates synthetic files seeded from that data, adapting Magicoder OSS-Instruct for pretraining.Source: §4 · p. 9

5 Conclusion

Arctic-SnowCoder-1.3B uses 555B tokens across general, high-quality, and synthetic-data pretraining stages. The paper concludes that downstream-distribution alignment is key to high-quality code data and offers several practical training-design guidelines.

  • Arctic-SnowCoder-1.3B uses 555B tokens across 500B general, 50B high-quality, and 5B synthetic-data pretraining stages.Source: §5 · p. 9
  • The paper identifies alignment with downstream application distributions as the key to high-quality data in code pretraining.Source: §5 · p. 9
  • The paper offers practical guidelines for repository-level data grouping, learning-rate scheduling, and repetition of high-quality data.Source: §5 · p. 9
Loading 2409.02326v1…