[Paper Review] GraphCodeBERT: Pre-training Code Representations with Data Flow

Daya Guo, Shuo Ren+16 more|arXiv (Cornell University)|Sep 17, 2020

Software Engineering ResearchComputer Science41 references152 citations

TL;DR

GraphCodeBERT pre-trains code representations by incorporating data flow structure, introducing two structure-aware tasks, and achieving state-of-the-art on four code-related downstream tasks.

ABSTRACT

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.

Motivation & Objective

• Motivate improving code understanding by leveraging code structure beyond token sequences.
• Propose a semantic-structure aware pre-training framework for programming languages.
• Investigate data flow as a scalable alternative to ASTs for encoding code semantics.
• Demonstrate benefits of structure-aware pre-training on multiple code-related tasks.

Proposed method

• Represent code with a Transformer backbone augmented to incorporate data flow as a graph input.
• Introduce graph-guided masked attention to selectively attend to data-flow nodes and code tokens.
• Add two structure-aware pre-training tasks: data flow edge prediction and node alignment between code and data flow.
• Combine these with standard masked language modeling to learn joint code and structure representations.
• Pre-train GraphCodeBERT on CodeSearchNet and evaluate on four downstream tasks.

Experimental results

Research questions

• RQ1 Does incorporating data flow as a semantic code structure improve learned representations beyond token sequences?
• RQ2 Can two structure-aware pre-training tasks (edge prediction and node alignment) enhance code understanding across tasks?
• RQ3 How does GraphCodeBERT compare to token-only pre-trained models on code search, clone detection, translation, and refinement?
• RQ4 Do models attend more to data-flow nodes than to raw code tokens in practice?

Key findings

• GraphCodeBERT achieves state-of-the-art performance on four downstream tasks: code search, clone detection, code translation, and code refinement.
• In ablation studies, removing data flow or the two structure-aware tasks degrades performance, validating their usefulness.
• Attention analysis shows the model prefers structure-level (data-flow) attention over purely token-level attention for code search.
• Case studies indicate data flow helps disambiguate variable semantics and improves code understanding.
• Across tasks, leveraging code structure improves over baselines including CodeBERT and RoBERTa variants.