PaperlayerView exact revision on arXiv

Source-linked AI summary

GraphCodeBERT: Pre-training Code Representations with Data Flow

Daya Guo, Shuo Ren, Shuai Lu, Zhangyin Feng, Duyu Tang, Shujie Liu, Long Zhou, Nan Duan, Alexey Svyatkovskiy, Shengyu Fu, Michele Tufano, Shao Kun Deng, Colin Clement, Dawn Drain, Neel Sundaresan, Jian Yin, Daxin Jiang, Ming Zhou

arXiv:2009.08366v4cs.SEcs.CL

TL;DR

Existing code-pretrained models largely treat code as token sequences, overlooking its inherent semantic structure. GraphCodeBERT incorporates data flow through structure-aware pre-training and achieves state-of-the-art performance across four downstream code tasks.

  • Problem

    Existing code-pretrained models mainly represent source code as token sequences, leaving its inherent semantic structure underused for code understanding.

  • Method

    GraphCodeBERT incorporates data flow into a Transformer with graph-guided attention and two structure-aware pre-training tasks.

  • Results

    GraphCodeBERT achieves state-of-the-art performance on code search, clone detection, code translation, and code refinement.

  • Takeaways & Limitations

    The results support using code structure and structure-aware pre-training to improve code representations across four downstream tasks.

Abstract

from arXiv · show

Pre-trained models for programming language have achieved dramatic empirical improvements on a variety of code-related tasks such as code search, code completion, code summarization, etc. However, existing pre-trained models regard a code snippet as a sequence of tokens, while ignoring the inherent structure of code, which provides crucial code semantics and would enhance the code understanding process. 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 use data flow in the pre-training stage, which is a semantic-level structure of code that encodes the relation of "where-the-value-comes-from" between variables. Such a semantic-level structure is neat and does not bring an unnecessarily deep hierarchy of AST, the property of which makes the model more efficient. We develop GraphCodeBERT based on Transformer. In addition to using the task of masked language modeling, we introduce two structure-aware pre-training tasks. One is to predict code structure edges, and the other is to align representations between source code and code structure. We implement the model in an efficient way with a graph-guided masked attention function to incorporate the code structure. We evaluate our model on four tasks, including code search, clone detection, code translation, and code refinement. Results show that code structure and newly introduced pre-training tasks can improve GraphCodeBERT and achieves state-of-the-art performance on the four downstream tasks. We further show that the model prefers structure-level attentions over token-level attentions in the task of code search.

1 INTRODUCTION

GraphCodeBERT is a Transformer-based pre-trained programming-language model that incorporates semantic data flow rather than syntactic AST structure. It introduces structure-aware pre-training tasks and achieves state-of-the-art performance across four downstream code tasks.

  • Model motivation and contribution: GraphCodeBERT leverages data flow, whose nodes represent variables and edges encode where variable values come from, instead of abstract syntax trees.Data flow is less complex and avoids an unnecessarily deep hierarchy, making the model more efficient.
  • Pre-training method: The model introduces data-flow edge prediction and variable alignment between source code and data flow as two structure-aware pre-training tasks.These tasks learn representations from code structure and align representations across source code and its data-flow graph.
  • Training and evaluation: GraphCodeBERT is pre-trained on CodeSearchNet (Husain et al., 2019), which contains 2.3M functions in six programming languages paired with natural-language documents.The pre-training corpus consists of functions paired with natural-language documentation.
  • Results: GraphCodeBERT achieves state-of-the-art performance on natural-language code search, clone detection, code translation, and code refinement.The paper further reports that code structure and the newly introduced pre-training tasks improve performance on these four downstream tasks.

2 RELATED WORKS

Prior work developed pre-trained models for programming languages and task-specific neural networks using code structure, but did not combine structural information with pre-training. GraphCodeBERT addresses this gap with structure-aware Transformer attention and pre-training objectives.

  • Pre-Trained Models for Programming Languages: Pre-trained models for programming languages emerged from NLP pre-training and supported code-intelligence development.
  • GraphCodeBERT is presented as the first pre-trained model to leverage code structure for code representation and understanding.
  • It incorporates structure through graph-guided masked attention in the Transformer and two structure-aware pre-training tasks linking source code with code structure.
  • Neural Networks with Code Structure: Neural networks using structures such as ASTs achieved strong results across code tasks, but learned task-specific models from scratch without pre-training.

3 DATA FLOW

Data flow represents variable dependencies through “where-the-value-comes-from” relations, providing semantic code structure that is consistent across abstract grammars. It is extracted by parsing source code into an AST, identifying variables as graph nodes, and connecting value dependencies with directed edges.

  • Data Flow Concept: Data flow models dependencies between variables, with nodes representing variables and edges indicating where each variable’s value comes from.Unlike ASTs, data flow remains the same across different abstract grammars for identical source code and provides semantic information for code understanding.
  • Data Flow Extraction: The extraction process first parses source code into an AST, whose terminal nodes identify the variable sequence used to construct the data-flow graph.For source code C = {c1, c2, ..., cn}, the variables are represented as V = {v1, v2, ..., vk}.
  • Data Flow Extraction: A directed edge ε = ⟨vi, vj⟩ indicates that the value of variable vj comes from variable vi.Figure 1 illustrates this extraction procedure and the resulting graph of variable value dependencies.

4 GRAPHCODEBERT

GraphCodeBERT extends a bidirectional Transformer with source code, comments, and data-flow variables as joint input. Graph-guided attention and three pre-training tasks incorporate data-flow dependencies and align code tokens with variables.

  • Model Architecture: GraphCodeBERT uses a multi-layer bidirectional Transformer that jointly processes comments, source code, and data-flow variables, with special position embeddings marking variables as graph nodes.The input concatenates comments, source code, and variables; data flow represents direct where-the-value-comes-from relations between variables.
  • Graph-Guided Masked Attention: Graph-guided masked attention permits variable nodes to attend along direct data-flow edges or themselves, while linking variables and code tokens when the data flow identifies their correspondence.Disallowed query-key pairs receive an infinitely negative attention score, making their softmax attention weight zero.
  • Pre-training Tasks: GraphCodeBERT pre-trains with masked language modeling, data-flow edge prediction, and node alignment across source code and data flow.The structure-aware tasks respectively predict masked variable edges and masked code-token–variable edges to learn dependencies and align representations.
  • Pre-training Tasks: Masked language modeling randomly samples 15% of source-code and comment tokens, replacing them with [MASK] 80% of the time, random tokens 10% of the time, or leaving them unchanged 10% of the time.The objective is to reconstruct the original sampled tokens.

GraphCodeBERT

GraphCodeBERT’s Node Alignment task masks graph edges involving randomly sampled nodes and trains the model to recover them. It also frames alignment over candidate nodes using balanced positive and negative samples.

  • Node Alignment: It masks edges between code tokens and sampled nodes, then predicts the masked edges.
  • Node Alignment: The Node Alignment task randomly samples 20% of graph nodes for masking.
  • Node Alignment: For node alignment, the candidate set is paired with equal numbers of positive and negative samples.

5 EXPERIMENTS

GraphCodeBERT is evaluated on code search, clone detection, code translation, and code refinement, where it consistently achieves state-of-the-art or superior performance over competing methods. Additional analyses show that its structure-aware objectives and data-flow representation contribute to performance and receive substantial model attention.

  • Downstream Tasks: GraphCodeBERT achieves state-of-the-art or superior performance across code search, clone detection, code translation, and code refinement benchmarks.It significantly outperforms competing methods on code search and clone detection, while pretrained initialization improves translation and refinement beyond conventional baselines.
  • Code Search: 2% MRR improvement over baselines gives GraphCodeBERT state-of-the-art code-search performance, with statistically significant gains (p < 0.01).The code-search evaluation uses filtered queries and the full code corpus rather than only 1,000 candidates, approximating a more realistic retrieval setting.
  • Clone Detection: GraphCodeBERT significantly outperforms other clone-detection methods (p < 0.01), demonstrating the benefit of incorporating code structure.The evaluation uses BigCloneBench and compares against AST- and flow-based approaches.
  • Code Translation: Pretrained models outperform PBSMT and Transformer baselines for code translation, with GraphCodeBERT achieving state-of-the-art performance.The encoder is initialized from pretrained models, while the decoder and source-to-target attention are randomly initialized before training.
  • Code Refinement: GraphCodeBERT outperforms other pretrained models on both code-refinement datasets, while pretrained models further surpass Transformer models.The refinement models fine-tune the whole encoder-decoder system after pretrained encoder initialization.
  • Ablation Study: Removing Node Alignment or Edge Prediction lowers code-search performance from 71.3% to 70.3%–70.7%, while data-flow nodes receive disproportionate attention.Although nodes comprise only 5%–20% of tokens, node attention exceeds their proportion by approximately 10%–32% across programming languages; data flow also outperforms AST-based alternatives across sequence lengths.

6 CONCLUSION

GraphCodeBERT leverages data flow to learn code representations and, to the authors’ knowledge, is the first pre-trained model to consider code structure during pre-training. It introduces two structure-aware pre-training tasks and achieves state-of-the-art performance on four downstream code tasks.

  • 6 CONCLUSION: GraphCodeBERT leverages data flow to learn code representations.
  • 6 CONCLUSION: To the authors’ knowledge, GraphCodeBERT is the first pre-trained model to consider code structure for pre-training code representations.
  • 6 CONCLUSION: GraphCodeBERT introduces two structure-aware pre-training tasks and achieves state-of-the-art performance on code search, clone detection, code translation, and code refinement.

A PRE-TRAINING DETAILS

GraphCodeBERT uses a 12-layer Transformer pretrained on CodeSearchNet with the same dataset as CodeBERT for fair comparison. Pretraining alternates structure-aware objectives with masked language modeling and balances languages through multinomial sampling.

  • A PRE-TRAINING DETAILS: GraphCodeBERT uses 12 Transformer layers with 768-dimensional hidden states and 12 attention heads, pretrained on CodeSearchNet’s 2.3M functions across six programming languages.The model uses the same dataset as CodeBERT for fair comparison.
  • A PRE-TRAINING DETAILS: At each iteration, EdgePred and NodeAlign alternate alongside masked language modeling during pretraining.This combines the two structure-aware objectives with MLM rather than applying them in isolation.
  • A PRE-TRAINING DETAILS: Language-balanced batches are sampled from a multinomial distribution with α=0.7 to alleviate bias toward high-resource programming languages.Each batch is sampled from a single programming language according to probabilities based on language example counts.

B NATURAL LANGUAGE CODE SEARCH

On CodeSearchNet, GraphCodeBERT retrieves code by matching separately encoded natural-language queries and data-flow-aware source code, achieving state-of-the-art performance under the standard 1,000-candidate setting. The evaluation also uses a filtered dataset and a harder full-corpus retrieval setting for development and testing queries.

  • Data and preprocessing: The study evaluates natural-language code search on CodeSearchNet, using each function’s first documentation paragraph as its query following Husain et al. (2019).Queries containing unrelated content, special tokens, invalid language, or unsuitable lengths are filtered, as are examples whose code cannot be parsed into an abstract syntax tree.
  • Evaluation setting: Unlike Husain et al. (2019), the development and testing answers are retrieved from the whole candidate-code corpus rather than 1,000 candidates.The filtered-dataset statistics for this full-corpus setting are reported in Table 7.
  • Model and training: GraphCodeBERT separately encodes each query and data-flow-annotated source code, ranking candidates by the inner product of their [CLS] representations.Fine-tuning uses a learning rate of 2e-5, batch size 32, query and code maximum lengths of 128 and 256, and at most 64 nodes, with Adam and development-set early stopping.
  • Results: GraphCodeBERT achieves state-of-the-art performance when retrieving answers from 1,000 candidate codes under Husain et al. (2019)’s setting.The results are reported in Table 8.

C CODE CLONE DETECTION … F.1 NATURAL LANGUAGE CODE SEARCH

The paper evaluates GraphCodeBERT across clone detection, code translation, code refinement, and natural-language code search using task-specific datasets and qualitative case studies. The examples show semantic clone classification, successful cross-language translation and bug fixing, and improved retrieval through data-flow understanding.

  • C CODE CLONE DETECTION: GraphCodeBERT evaluates clone detection on BigCloneBench, retaining 9,134 code fragments after filtering for tagged true and false clone pairs.The dataset contains over 6,000,000 true clone pairs and 260,000 false clone pairs across 10 functionalities.
  • C CODE CLONE DETECTION: A Java pair that downloads URL content and converts it to strings receives a semantic-similarity score of 0.983 and is classified as a true clone pair.The example illustrates that the model assigns a high score when two programs produce similar outputs for the same input.
  • D CODE TRANSLATION: For code translation, the model pairs Java and C# methods from shared open-source projects and successfully translates language-specific constructs between them.The example changes “boolean” to “bool” and replaces “String.valueOf(b)” with “b.ToString()”; Itext and JTS are excluded because of licensing.
  • E CODE REFINEMENT: The refinement dataset contains 46,680/5,835/5,835 small-subset training/development/test samples and 52,364/6,545/6,545 medium-subset samples.Variable and custom-method names are almost entirely normalized in the dataset.
  • E CODE REFINEMENT: GraphCodeBERT successfully fixes code-refinement bugs by changing “*” to “+” and removing an invalid “return” statement that conflicts with a “void” declaration.The evaluation uses buggy Java functions paired with their fixed versions from Tufano et al. (2019), across small and medium code-length subsets.
  • F.1 NATURAL LANGUAGE CODE SEARCH: In natural-language code search, GraphCodeBERT successfully retrieves correct source code for both illustrated queries, outperforming the interpretability of CodeBERT and RoBERTa (code) through data-flow-aware understanding.The first example highlights improved handling of the complicated expression “[(k, v) for k, v in self.items() if v is not self.EMPT”.

F.2 CODE CLONE DETECTION · F.3 CODE TRANSLATION AND CODE REFINEMENT

The paper uses case studies to examine GraphCodeBERT’s data-flow benefits in code clone detection, code translation, and code refinement. These comparisons contrast GraphCodeBERT with CodeBERT/RoBERTa or a Transformer without data flow using illustrative examples and model outputs.

  • F.2 CODE CLONE DETECTION: GraphCodeBERT is compared with CodeBERT and RoBERTa (code) on code clone detection through a case study illustrating how data flow supports semantic understanding.The example contrasts code returning HTML from a given URL with code returning the last line from a fixed URL, whose outputs make their semantics dissimilar.
  • F.2 CODE CLONE DETECTION: The code-search examples and retrieved results from different models are presented separately in Figure 9, providing contextual evidence for the paper’s broader case-study analysis.The supplied passage identifies the figure’s scope but does not report specific retrieval outcomes.
  • F.2 CODE CLONE DETECTION: Figure 10 presents a code-clone example with predictions from different models and highlights overlapping snippets between the source codes.The figure is used to visualize the case-study comparison rather than to provide aggregate benchmark values.
  • F.3 CODE TRANSLATION AND CODE REFINEMENT: GraphCodeBERT is compared with a Transformer without data flow on code generation tasks covering code translation and code refinement.The study lists three cases for each task in Tables 9 and 10.
  • F.3 CODE TRANSLATION AND CODE REFINEMENT: In the generation case study, the Transformer baseline makes mistakes including repeated tokens, logic errors, and syntax errors.The supplied passage introduces these failure types when contrasting [sys] with GraphCodeBERT, but the excerpt does not provide the complete case outcomes.
  • F.3 CODE TRANSLATION AND CODE REFINEMENT: Table 9 documents three Java-to-C# translation examples using source, reference, Transformer-without-data-flow, and GraphCodeBERT outputs.The notation is [src] for source input, [ref] for reference, [sys] for the baseline, and [ours] for GraphCodeBERT.
  • F.3 CODE TRANSLATION AND CODE REFINEMENT: Table 10 documents three code-refinement examples with the same source, reference, baseline, and GraphCodeBERT output notation.The table caption identifies the examples and comparison, without supplying aggregate numerical results.

G ERROR ANALYSIS

The error analysis identifies two main failure classes: difficulty understanding library APIs in code search and semantic or syntax errors in code translation. The authors suggest structure-aware remedies involving relevant API information and grammar-aware generation.

  • Potential remedies: The analysis proposes incorporating relevant API information for code search and a dedicated grammar-aware decoder or production-rule generation for code translation.These directions target API-meaning failures and generation errors caused by programming-language syntax and generation paradigms.
  • Code generation: In code translation, GraphCodeBERT produces both semantic errors, such as identifiers from nowhere, and syntax errors, such as a missing “}” before “return n”.The analysis presents two representative translation cases.
  • Code understanding: GraphCodeBERT’s code-search failures mainly involve retrieving source code that uses library functions whose API meanings lack relevant information.The examples include TensorFlow functions such as “tf.io.read file” and “tf.image.decode image,” and “GoogleCloudStorageHook.”
Loading 2009.08366v4…