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: This paper proposes GraphCodeBERT as a Transformer-based pretrained model for programming language that incorporates data flow information in the graph representation of variables in the code. The data flow graph encodes the structure of variables based on “where-the-value-comes-from” from the AST parse. The pretrained model is jointly trained on the code, the natural language comment of the code, and the data flow graph of the code. In addition to the Masked Language Modeling objective, two new pretraining objectives are proposed including predicting the edge of the data flow graph and predicting the alignment of variables between data flow graph and code. The graph-guided masked attention is used such that the attention can only occur if two variables have an edge in the data flow graph or there is an alignment between data flow graph and code. The experiments show that GraphCodeBERT can deliver improvements on Natural Language Code Search, Code Clone Detection, Code Translation, and Code Refinement.