GraphCodeBERT: Pre-training Code Representations with Data Flow

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.

Accepted by ICLR2021

References in corpus (6)

• Cross-lingual Language Model Pretraining
• Exploring Software Naturalness through Neural Language Models
• SCELMo: Source Code Embeddings from Language Models
• Detecting Code Clones with Graph Neural Networkand Flow-Augmented Abstract Syntax Tree
• Tree-structured Attention with Hierarchical Accumulation
• Coupling Retrieval and Meta-Learning for Context-Dependent Semantic Parsing

Cited by in corpus (21)

• On the Opportunities and Risks of Foundation Models
• CodeXGLUE: A Machine Learning Benchmark Dataset for Code Understanding and Generation
• CodeBLEU: a Method for Automatic Evaluation of Code Synthesis
• Natural Language Generation and Understanding of Big Code for AI-Assisted Programming: A Review
• Code Structure Guided Transformer for Source Code Summarization
• Deep Learning-based Software Engineering: Progress, Challenges, and Opportunities
• Unified Pre-training for Program Understanding and Generation
• AMMUS : A Survey of Transformer-based Pretrained Models in Natural Language Processing
• The Good, the Bad, and the Missing: Neural Code Generation for Machine Learning Tasks
• CoEdPilot: Recommending Code Edits with Learned Prior Edit Relevance, Project-wise Awareness, and Interactive Nature
• ProphetNet-X: Large-Scale Pre-training Models for English, Chinese, Multi-lingual, Dialog, and Code Generation
• How Does Naming Affect LLMs on Code Analysis Tasks?
• Is a Single Model Enough? MuCoS: A Multi-Model Ensemble Learning for Semantic Code Search
• An Empirical Study of the Imbalance Issue in Software Vulnerability Detection
• Identifying Non-Control Security-Critical Data through Program Dependence Learning
• GraphCode2Vec: Generic Code Embedding via Lexical and Program Dependence Analyses
• How could Neural Networks understand Programs?
• Distilling Transformers for Neural Cross-Domain Search
• UNICS: Multilingual Code Search via Unified Pseudocode and Contrastive Transfer Learning
• Assessing the Latent Automated Program Repair Capabilities of Large Language Models using Round-Trip Translation
• Using Document Similarity Methods to create Parallel Datasets for Code Translation