Aprendizado de Representação de Código em Escala

Code Representation Learning At Scale

February 2, 2024

Autores: Dejiao Zhang, Wasi Ahmad, Ming Tan, Hantian Ding, Ramesh Nallapati, Dan Roth, Xiaofei Ma, Bing Xiang

cs.AI

Resumo

Estudos recentes mostraram que modelos de linguagem de código em grande escala demonstram ganhos significativos de desempenho em tarefas subsequentes, como a geração de código. No entanto, a maioria dos trabalhos existentes sobre aprendizado de representação de código treina modelos em uma escala de centenas de milhões de parâmetros usando corpora de pré-treinamento muito limitados. Neste trabalho, impulsionamos o aprendizado de representação de código com uma vasta quantidade de dados de código por meio de um esquema de pré-treinamento em duas etapas. Primeiro, treinamos os codificadores usando uma mistura que aproveita tanto a aleatoriedade na modelagem de linguagem com mascaramento quanto o aspecto estrutural das linguagens de programação. Em seguida, aprimoramos as representações por meio de aprendizado contrastivo com negativos difíceis e positivos difíceis construídos de maneira não supervisionada. Estabelecemos um modelo de codificador pronto para uso que supera consistentemente os modelos existentes em uma ampla variedade de tarefas subsequentes por grandes margens. Para compreender os fatores que contribuem para o sucesso do aprendizado de representação de código, realizamos ablações detalhadas e compartilhamos nossas descobertas sobre (i) um esquema personalizado e eficaz de remoção de ruído em nível de token para código-fonte; (ii) a importância de negativos difíceis e positivos difíceis; (iii) como o aprendizado contrastivo bimodal proposto impulsiona o desempenho da busca semântica cross-lingual; e (iv) como os esquemas de pré-treinamento determinam a escala de desempenho das tarefas subsequentes com o tamanho do modelo.

English

Recent studies have shown that code language models at scale demonstrate significant performance gains on downstream tasks, i.e., code generation. However, most of the existing works on code representation learning train models at a hundred million parameter scale using very limited pretraining corpora. In this work, we fuel code representation learning with a vast amount of code data via a two-stage pretraining scheme. We first train the encoders via a mix that leverages both randomness in masking language modeling and the structure aspect of programming language. We then enhance the representations via contrastive learning with hard negative and hard positive constructed in an unsupervised manner. We establish an off-the-shelf encoder model that persistently outperforms the existing models on a wide variety of downstream tasks by large margins. To comprehend the factors contributing to successful code representation learning, we conduct detailed ablations and share our findings on (i) a customized and effective token-level denoising scheme for source code; (ii) the importance of hard negatives and hard positives; (iii) how the proposed bimodal contrastive learning boost the cross-lingual semantic search performance; and (iv) how the pretraining schemes decide the downstream task performance scales with the model size.