GraphCodeBERT: Pre-training Code Representations with Data Flow

1 INTRODUCTION . Pre-trained models such as ELMo ( Peters et al. , 2018 ) , GPT ( Radford et al. , 2018 ) and BERT ( Devlin et al. , 2018 ) have led to strong improvement on numerous natural language processing ( NLP ) tasks . These pre-trained models are first pre-trained on a large unsupervised text corpus , and then fine-tuned on downstream tasks . The success of pre-trained models in NLP also promotes the development of pre-trained models for programming language . Existing works ( Kanade et al. , 2019 ; Karampatsis & Sutton , 2020 ; Feng et al. , 2020 ; Svyatkovskiy et al. , 2020 ; Buratti et al. , 2020 ) regard a source code as a sequence of tokens and pre-train models on source code to support code-related tasks such as code search , code completion , code summarization , etc . However , previous works only utilize source code for pre-training , while ignoring the inherent structure of code . Such code structure provides useful semantic information of code , which would benefit the code understanding process . Taking the expression v = max value−min value as an example , v is computed from max value and min value . Programmers do not always follow the naming conventions so that it ’ s hard to understand the semantic of the variable v only from its name . The semantic structure of code provides a way to understand the semantic of the variable v by leveraging dependency relation between variables . ∗Work done while this author was an intern at Microsoft Research Asia . Contact : Daya Guo ( guody5 @ mail2.sysu.edu.cn ) 1All the codes and data are available at https : //github.com/microsoft/CodeBERT . In this work , we present GraphCodeBERT , a pre-trained model for programming language that considers the inherent structure of code . Instead of taking syntactic-level structure of code like abstract syntax tree ( AST ) , we leverage semantic-level information of code , i.e . data flow , for pretraining . Data flow is a graph , in which nodes represent variables and edges represent the relation of “ where-the-value-comes-from ” between variables . Compared with AST , data flow is less complex and does not bring an unnecessarily deep hierarchy , the property of which makes the model more efficient . In order to learn code representation from source code and code structure , we introduce two new structure-aware pre-training tasks . One is data flow edges prediction for learning representation from code structure , and the other is variable-alignment across source code and data flow for aligning representation between source code and code structure . GraphCodeBERT is based on Transformer neural architecture ( Vaswani et al. , 2017 ) and we extend it by introducing a graph-guided masked attention function to incorporate the code structure . We pre-train GraphCodeBERT on the CodeSearchNet dataset ( Husain et al. , 2019 ) , which includes 2.3M functions of six programming languages paired with natural language documents . We evaluate the model on four downstream tasks : natural language code search , clone detection , code translation , and code refinement . Experiments show that our model achieves state-of-the-art performance on the four tasks . Further analysis shows that code structure and newly introduced pre-training tasks can improve GraphCodeBERT and the model has consistent preference for attending data flow . In summary , the contributions of this paper are : ( 1 ) GraphCodeBERT is the first pre-trained model that leverages semantic structure of code to learn code representation . ( 2 ) We introduce two new structure-aware pre-training tasks for learning representation from source code and data flow . ( 3 ) GraphCodeBERT provides significant improvement on four downstream tasks , i.e . code search , clone detection , code translation , and code refinement . 2 RELATED WORKS . Pre-Trained Models for Programming Languages Inspired by the big success of pre-training in NLP ( Devlin et al. , 2018 ; Yang et al. , 2019 ; Liu et al. , 2019 ; Raffel et al. , 2019 ) , pre-trained models for programming languages also promotes the development of code intelligence ( Kanade et al. , 2019 ; Feng et al. , 2020 ; Karampatsis & Sutton , 2020 ; Svyatkovskiy et al. , 2020 ; Buratti et al. , 2020 ) . Kanade et al . ( 2019 ) pre-train a BERT model on a massive corpus of Python source codes by masked language modeling and next sentence prediction objectives . Feng et al . ( 2020 ) propose CodeBERT , a bimodal pre-trained model for programming and natural languages by masked language modeling and replaced token detection to support text-code tasks such as code search . Karampatsis & Sutton ( 2020 ) pre-train contextual embeddings on a JavaScript corpus using the ELMo framework for program repair task . Svyatkovskiy et al . ( 2020 ) propose GPT-C , which is a variant of the GPT-2 trained from scratch on source code data to support generative tasks like code completion . Buratti et al . ( 2020 ) present C-BERT , a transformer-based language model pre-trained on a collection of repositories written in C language , and achieve high accuracy in the abstract syntax tree ( AST ) tagging task . Different with previous works , GraphCodeBERT is the first pre-trained model that leverages code structure to learn code representation to improve code understanding . We further introduce a graphguided masked attention function to incorporate the code structure into Transformer and two new structure-aware pre-training tasks to learn representation from source code and code structure . Neural Networks with Code Structure In recent years , some neural networks leveraging code structure such as AST have been proposed and achieved strong performance in code-related tasks like code completion ( Li et al. , 2017 ; Alon et al. , 2019 ; Kim et al. , 2020 ) , code generation ( Rabinovich et al. , 2017 ; Yin & Neubig , 2017 ; Brockschmidt et al. , 2018 ) , code clone detection ( Wei & Li , 2017 ; Zhang et al. , 2019 ; Wang et al. , 2020 ) , code summarization ( Alon et al. , 2018 ; Hu et al. , 2018 ) and so on ( Nguyen & Nguyen , 2015 ; Allamanis et al. , 2018 ; Hellendoorn et al. , 2019 ) . Nguyen & Nguyen ( 2015 ) propose an AST-based language model to support the detection and suggestion of a syntactic template at the current editing location . Allamanis et al . ( 2018 ) use graphs to represent programs and graph neural network to reason over program structures . Hellendoorn et al . ( 2019 ) propose two different architectures using a gated graph neural network and Transformers for combining local and global information to leverage richly structured representations of source code . However , these works leverage code structure to learn models on specific tasks from scratch without using pre-trained models . In this work , we study how to leverage code structure for pre-training code representation . 3 DATA FLOW . In this section , we describe the basic concept and extraction of data flow . In next section , we will describe how to use data flow for pre-training . Data flow is a graph that represents dependency relation between variables , in which nodes represent variables and edges represent where the value of each variable comes from . Unlike AST , data flow is same under different abstract grammars for the same source code . Such code structure provides crucial code semantic information for code understanding . Taking v = max value−min value as an example , programmers do not always follow the naming conventions so that it is hard to understand the semantic of the variable . Data flow provides a way to understand the semantic of the variable v to some extent , i.e . the value of v comes from max value and min value in data flow . Besides , data flow supports the model to consider long-range dependencies induced by using the same variable or function in distant locations . Taking Figure 1 as an example , there are four variables with same name ( i.e . x3 , x7 , x9 and x11 ) but with different semantic . The graph in the figure shows dependency relation between these variables and supports x11 to pay more attention to x7 and x9 instead of x3 . Next , we describe how to extract data flow from a source code . Figure 1 shows the extraction of data flow through a source code . Given a source code C = { c1 , c2 , ... , cn } , we first parse the code into an abstract syntax tree ( AST ) by a standard compiler tool2 . The AST includes syntax information of the code and terminals ( leaves ) are used to identify the variable sequence , denoted as V = { v1 , v2 , ... , vk } . We take each variable as a node of the graph and an direct edge ε = 〈vi , vj〉 from vi to vj refers that the value of j-th variable comes from i-th variable . Taking x = expr as an example , edges from all variables in expr to x are added into the graph . We denote the set of directed edges as E = { ε1 , ε2 , ... , εl } and the graph G ( C ) = ( V , E ) is data flow used to represent dependency relation between variables of the source code C . 4 GRAPHCODEBERT . In this section , we describe GraphCodeBERT , a graph-based pre-trained model based on Transformer for programming language . We introduce model architecture , graph-guided masked attention and pre-training tasks including standard masked language model and newly introduced ones . More details about model pre-training setting are provided in the Appendix A . 2https : //github.com/tree-sitter/tree-sitter 4.1 MODEL ARCHITECTURE . Figure 2 shows the model architecture of GraphCodeBERT . We follow BERT ( Devlin et al. , 2018 ) and use the multi-layer bidirectional Transformer ( Vaswani et al. , 2017 ) as the model backbone . Instead of only using source code , we also utilize paired comments to pre-train the model to support more code-related tasks involving natural language such as natural language code search ( Feng et al. , 2020 ) . We further take data flow , which is a graph , as a part of the input to the model . Given a source code C = { c1 , c2 , ... , cn } with its comment W = { w1 , w2 , ... , wm } , we can obtain the corresponding data flow G ( C ) = ( V , E ) as discussed in the Section 3 , where V = { v1 , v2 , ... , vk } is a set of variables and E = { ε1 , ε2 , ... , εl } is a set of direct edges that represent where the value of each variable comes from . We concatenate the comment , source code and the set of variables as the sequence input X = { [ CLS ] , W , [ SEP ] , C , [ SEP ] , V } , where [ CLS ] is a special token in front of three segments and [ SEP ] is a special symbol to split two kinds of data types . GraphCodeBERT takes the sequence X as the input and then converts the sequence into input vectors H0 . For each token , its input vector is constructed by summing the corresponding token and position embeddings . We use a special position embedding for all variables to indicate that they are nodes of data flow . The model applies N transformer layers over the input vectors to produce contextual representations Hn = transformern ( Hn−1 ) , n ∈ [ 1 , N ] . Each transformer layer contains an architecturally identical transformer that applies a multi-headed self-attention operation ( Vaswani et al. , 2017 ) followed by a feed forward layer over the input Hn−1 in the n-th layer . Gn = LN ( MultiAttn ( Hn−1 ) +Hn−1 ) ( 1 ) Hn = LN ( FFN ( Gn ) +Gn ) ( 2 ) where MultiAttn is a multi-headed self-attention mechanism , FFN is a two layers feed forward network , and LN represents a layer normalization operation . For the n-th transformer layer , the output Ĝn of a multi-headed self-attention is computed via : Qi = H n−1WQi , Ki = H n−1WKi , Vi = H n−1WVi ( 3 ) headi = softmax ( QiK T i√ dk +M ) Vi ( 4 ) Ĝn = [ head1 ; ... ; headu ] W O n ( 5 ) where the previous layer ’ s output Hn−1 ∈ R|X|×dh is linearly projected to a triplet of queries , keys and values using model parameters WQi , W K i , W V i ∈ Rdh×dk , respectively . u is the number of heads , dk is the dimension of a head , and WOn ∈ Rdh×dh is the model parameters . M ∈ R|X|×|X| is a mask matrix , where Mij is 0 if i-th token is allowed to attend j-th token otherwise −∞ .

Summary: The authors present Graph Code BERT, the first language model that leverages data flow to learn code representation. They use three objective functions: Masked Language Modeling, Edge Prediction, and Node Alignment. They claim their structure-aware pre-training can help improving performance on code-related downstream tasks, including code search, clone detection, code translation, and code refinement.