Contrastive Code Representation Learning

1 INTRODUCTION . Programmers increasingly rely on machine-aided programming tools to aid software development ( Kim et al. , 2012 ) . However , the wide diversity of programs encountered in practice limits the generalization of hand-written rules . Catching semantic bugs such as naming errors requires deeper language understanding , motivating learning-based programming tools . Recent work uses machine learning for bug detection ( Pradel & Sen , 2018 ) and optimization ( Mendis et al. , 2019 ) . Consider predicting the type of the variable declaration “ var median = ... ; ” . Static analysis fails as the type is underspecified , but the variable name indicates the statement is a float . Programming language datasets suffer from scarce annotations due to the time and expertise required to label . State-of-the-art approaches generally rely on either ( 1 ) synthetic supervised datasets or ( 2 ) self-supervised pre-training . Synthetic auto-generated labels have been used for method naming ( Alon et al. , 2019a ; b ) and bug detection ( Ferenc et al. , 2018 ; Benton et al. , 2019 ; Pradel & Sen , 2018 ) . However , synthetic code datasets suffer from duplication issues ( Allamanis , 2019 ) and biases ( Shin et al. , 2019 ) which degrade generalization . Moreover , auto-generated data does not cover the diverse program behaviors encountered in the wild . In contrast , self-supervised learning can leverage large open-source repositories such as GitHub with limited or no annotations . Inspired by the success of pre-training in natural language processing , recent work uses self-supervision to learn code representations . Authors have explored context-based token embeddings ( Ben-Nun et al. , 2018 ) and masked language modeling , where tokens are corrupted and reconstructed ( Feng et al. , 2020 ; Kanade et al. , 2020 ) However , reconstruction focuses on superficial language reasoning and does not explicitly address the underlying program functionality . The resulting models attend to program implementation specifics such as variable names . We hypothesize that programs with the same functionality should have the same underlying representation for downstream code understanding tasks , a principle illustrated in Fig . 1 . While it is time intensive to identify equivalent programs in a large corpus , it is cheap to leverage static compiler transformations to automatically generate many equivalent versions of a particular source program . In this work , we develop ContraCode , a self-supervised representation learning algorithm that uses source-to-source compiler transformation techniques ( e.g. , dead code elimination , obfuscation and constant folding ) to generate syntactically diverse but functionally equivalent programs . ContraCode uses these equivalent programs to construct a challenging discriminative pretext task that requires the model to identify equivalent programs out of a large dataset of distractors . In doing so , it has to embed the functionality , not the form , of the code . In essence , the domain knowledge from our code transformations induces the knowledge of the structure of programs onto learned representations . The contributions of our work include : 1. the novel use of compiler-inspired transformations as data augmentations for code , 2. the concept of program representation learning based on functional equivalence , and 3. a detailed analysis of architectures , code transforms and pre-train strategies , where ContraCode improves static type inference top-1 accuracy by 9 % , learned inference by 2 % – 13 % , summarization F1 score by up to 8 % and clone detection AUROC by 5 % – 10 % . 2 RELATED WORK . Self-supervised learning ( SSL ) is a general representation learning strategy where some dimensions or attributes of a datapoint are predicted from the remaining parts . These methods are unsupervised in the sense that they do not rely on labels , but SSL tasks often adapt losses and architectures designed for supervised learning . Self-supervised pre-training has yielded large improvements in both NLP ( Howard & Ruder , 2018 ; Devlin et al. , 2018 ; Radford et al. , 2018 ; 2019 ) and computer vision ( Mahajan et al. , 2018 ) by improving generalization ( Erhan et al. , 2010 ; Hao et al. , 2019 ) . Weak visual features , such as orientation ( Gidaris et al. , 2018 ) , color ( Zhang et al. , 2016 ) , and context ( Pathak et al. , 2016 ) , are meaningful signals for representations ( Mahajan et al. , 2018 ) . Contrastive learning unifies many past SSL approaches that compare pairs or collections of similar and dissimilar items ( Hadsell et al. , 2006 ) . Rather than training the network to predict labels or reconstruct data , contrastive methods minimize the distance between the representations of similar examples ( positives ) while maximizing the distance between dissimilar examples ( negatives ) . Examples include Siamese networks ( Bromley et al. , 1994 ) and triplet losses ( Schroff et al. , 2015 ) . Contrastive predictive coding ( Oord et al. , 2018 ; Hénaff et al. , 2019 ) learns to encode chunks of sequential data to predict of future chunks with the InfoNCE loss , a variational lower bound on mutual information between views of the data ( Tian et al. , 2019 ; Wu et al. , 2020 ) inspired by noiseconstrastive estimation ( Gutmann & Hyvärinen , 2010 ) . In instance discrimination tasks ( Wu et al. , 2018 ) , views and not pieces of an entire image are compared . SimCLR ( Chen et al. , 2020a ) and Momentum Contrast ( He et al. , 2019 ; Chen et al. , 2020b ) recently made progress by using many negatives for dense loss signal . Beyond images , InfoNCE has been applied to NLP ( Chuang et al. , 2020 ; Giorgi et al. , 2020 ) , but may require supervision ( Fang & Xie , 2020 ) . Code representation learning There has been substantial work on architectures and tasks for machine learning on code ( Allamanis et al. , 2018 ) . We adopt the summarization task of Alon et al . function x ( maxLine ) { const section = { text : `` , data } ; for ( ; i < maxLine ; i += 1 ) { section.text += ` $ { lines [ i ] } \n ` ; } if ( section ) { parsingCtx.sections.push ( section ) ; } } Original JavaScript method function x ( t ) { const n = { 'text ' : `` , 'data ' : data } ; for ( ; i < t ; i += 1 ) { n.text += lines [ i ] + '\n ' ; } n & & parsingCtx.sections.push ( n ) ; } Renamed variables , explicit object style , explicit concatenation , inline conditional function x ( t ) { const n= { 'text ' : '' , 'data ' : data } ; for ( ; i < t ; i+= 1 ) n.text+=lines [ i ] +'\n ' ; n & & parsingCtx.sections.push ( n ) } Mangled source with compressed whitespace Figure 2 : A JavaScript method from the unlabeled training set with two automatically generated semantically-equivalent programs . The original method is from the StackEdit Markdown editor . ( 2019a ) , and the variable type inference task of DeepTyper ( Hellendoorn et al. , 2018 ) . Other authors have explored summarization ( Movshovitz-Attias & Cohen , 2013 ; Allamanis et al. , 2016 ; Iyer et al. , 2016 ) and type inference ( Pradel et al. , 2019 ; Pandi et al. , 2020 ; Wei et al. , 2020 ; Allamanis et al. , 2020 ; Bielik & Vechev , 2020 ) with different languages and datasets . The tree or graph structure of code can be exploited to encode invariances in the representation . Inst2vec ( Ben-Nun et al. , 2018 ) locally embeds individual statements in LLVM IR by processing a contextual flow graph with a context prediction objective ( Mikolov et al. , 2013 ) . Tree-Based CNN embeds the Abstract Syntax Tree ( AST ) nodes of high-level source code . Code2seq ( Alon et al. , 2019a ) embeds AST paths with an attention-based encoder and LSTM decoder for supervised sequence-to-sequence tasks . Kanade et al . ( 2020 ) ; Feng et al . ( 2020 ) pre-train the Transformer ( Vaswani et al. , 2017 ) on code using the masked language modeling objective ( Devlin et al. , 2018 ) , an instance of the cloze task ( Taylor , 1953 ) where the model reconstructs corrupted tokens . Recurrent networks have also been pre-trained on code ( Hussain et al. , 2020 ) as language models ( Peters et al. , 2018 ; Karampatsis & Sutton , 2020 ) . Wang & Christodorescu ( 2019 ) ; Wang & Su ( 2019 ) assess the stability of program analyzers under semi-automated program transformations . Concurrent work by Rabin & Alipour ( 2020 ) found that code2vec and code2seq often change their classifications when statements are permuted , variables are renamed , or other-semantic preserving transformations are applied . 3 METHOD : CONTRASTIVE CODE REPRESENTATION LEARNING . Understanding program functionality and global structure is important for difficult tasks like summarizing code in natural language . For these problems , learned code representations should be similar for functionally equivalent programs and dissimilar for non-equivalent programs ( Figure 1 ) . The principle of contrastive learning offers a simple objective for learning such representations if data can be organized into pairs of positives and negatives . We use each pair to shape representation space , drawing positives together and pushing negatives apart . However , a major question remains : given an unlabeled corpus of programs , how do we identify or generate similar programs ? We address this question in Sec . 3.1 , then introduce our learning framework in Sec . 3.2 . 3.1 COMPILATION AS DATA AUGMENTATION . Modern programming languages afford great flexibility to software developers , allowing them to implement the same desired functionality in different ways . Crowdsourced datasets mined from developers , such as GitHub repositories , have many near-duplicates in terms of textual similarity ( Allamanis , 2019 ) , and are bound to contain even more functional equivalences for common tasks . Satisfiability solvers can identify these equivalent programs ( Joshi et al. , 2002 ; Bansal & Aiken , 2006 ) , but functional equivalence is also undecidable in general ( Rice , 1953 ) . Also , formal documentation of semantics is required . Programs can instead be compared approximately using test-cases ( Massalin , 1987 ) , but this is costly and requires executing untrusted code . Instead of searching for equivalences , we propose correct by construction data augmentation . Our insight is to apply source-to-source compiler transformations to unlabeled code to generate many variants with the same functionality . For example , dead-code elimination ( DCE ) is a common compiler optimization that removes operations that leave the output of a function unchanged . While Code compression Identifier modification 3 Reformatting ( R ) 3 Variable renaming ( VR ) 3 Beautification ( B ) 3 Identifier mangling ( IM ) 3 Compression ( C ) Regularization 3 Dead-code elimination ( DCE ) 3 Dead-code insertion ( DCI ) 3 Type upconversion ( T ) 3 Subword regularization ( SW ) 3 Constant folding ( CF ) 7 Line subsampling ( LS ) 3 = semantics-preserving transformation 7 = lossy transformation Table 1 : We augment programs with 11 automated source-tosource compiler transformations . 10 of the 11 transformations are correct-by-construction and do not modify operational semantics . More details are in Section A.3 . DCE preserves program functionality , Wang & Christodorescu ( 2019 ) find that up to 12.7 % of the predictions of current algorithm classification models change after DCE—supervised datasets were not enough to acquire the domain knowledge that DCE does not matter . A particular source code sequence , e.g . “ W * x + b ” is parsed unambiguously into a tree-structured representation “ ( + ( * W x ) b ) ” . This tree is then transformed by automated traversal algorithms . A rich body of prior programming language work explores parsing then tranforming Abstract Syntax Trees to optimize a program prior to machine code generation . If source code is output rather than machine code , this is called source-to-source transformation . Source-to-source transformations are common for optimization and obfuscation purposes in dynamic languages like JavaScript . If each transformation preserves code functionality , then any composition also preserves code functionality . We leverage the Babel and Terser compiler infrastructure tools for JavaScript ( McKenzie et al. , 2020 ; Santos et al. , 2020 ) to parse code into an Abstract Syntax Tree ( AST ) and then perform correctnesspreserving transformations on method bodies . Table 1 and Appendix A.3 list all transformations , but we broadly group program transformations into three categories . Code compression changes the syntactic structure of code and performs correct-by-construction transformations such as precomputing constant expressions at compile time . Identifier modification substitutes method and variable names with random tokens , thereby masking part of the semantic information in programs . Finally , transformations for Regularization improve model generalization by reducing the number of trivial positive pairs with high text overlap ; this group potentially modifies program semantics through the line subsampling pass .

Summary: This paper studies the self-supervised code functional representation learning and proposes a method called ContraCode. ContraCode utilizes some code functionality invariant transformations to generate positive pairs from the same code and negative pairs from different codes. After that, these codes pairs will be used to do the contrastive pre-training. Experiment results based on two tasks are reported.