Arctic-SnowCoder: Desmistificando Dados de Alta Qualidade na Pré-treinamento de Código

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

September 3, 2024

Autores: Yuxiang Wei, Hojae Han, Rajhans Samdani

cs.AI

Resumo

Estudos recentes têm demonstrado cada vez mais que dados de alta qualidade são cruciais para o pré-treinamento eficaz de modelos de linguagem. No entanto, a definição precisa de "alta qualidade" ainda não foi explorada em profundidade. Focando no domínio do código, apresentamos o Arctic-SnowCoder-1.3B, um modelo de código base eficiente em dados pré-treinado em 555B de tokens por meio de três fases de dados progressivamente refinados: (1) pré-treinamento geral com 500B de tokens de código de qualidade padrão, pré-processados por meio de filtragem básica, deduplicação e descontaminação, (2) continuação do pré-treinamento com 50B de tokens de alta qualidade, selecionados da fase um por um anotador de qualidade no estilo BERT treinado para distinguir bom código de dados aleatórios, usando exemplos positivos retirados de arquivos de código de alta qualidade, juntamente com dados de instrução do Magicoder e StarCoder2-Instruct, e (3) aprimoramento do pré-treinamento com 5B de dados sintéticos criados pelo Llama-3.1-70B usando os dados da fase dois como sementes, adaptando a abordagem do Magicoder para pré-treinamento. Apesar de ser treinado em um conjunto de dados limitado, o Arctic-SnowCoder alcança desempenho de ponta no BigCodeBench, um benchmark de codificação focado em tarefas de programação práticas e desafiadoras, em comparação com modelos de tamanho semelhante treinados em não mais que 1T de tokens, superando o Phi-1.5-1.3B em 36%. Em todos os benchmarks avaliados, o Arctic-SnowCoder-1.3B supera o StarCoderBase-3B pré-treinado em 1T de tokens. Além disso, ele iguala o desempenho dos principais modelos de código base pequeno treinados em trilhões de tokens. Por exemplo, o Arctic-SnowCoder-1.3B supera o StarCoder2-3B, pré-treinado em mais de 3.3T de tokens, no HumanEval+, um benchmark que avalia a geração de código em nível de função, e permanece competitivo no BigCodeBench. Nossa avaliação apresenta uma análise abrangente justificando várias escolhas de design para o Arctic-SnowCoder. Mais importante ainda, descobrimos que a chave para dados de alta qualidade é sua alinhamento com a distribuição de aplicações futuras.

English

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.