Retrieval-Augmented Generation for Code Summarization via Hybrid GNN

1 INTRODUCTION . With software growing in size and complexity , developers tend to spend nearly 90 % ( Wan et al. , 2018 ) effort on software maintenance ( e.g. , version iteration and bug fix ) in the completed life cycle of software development . Source code summary , in the form of natural language , plays a critical role in the comprehension and maintenance process and greatly reduces the effort of reading and comprehending programs . However , manually writing code summaries is tedious and timeconsuming , and with the acceleration of software iteration , it has become a heavy burden for software developers . Hence , source code summarization which automates concise descriptions of programs is meaningful . Automatic source code summarization is a crucial yet far from the settled problem . The key challenges include : 1 ) the source code and the natural language summary are heterogeneous , which means they may not share common lexical tokens , synonyms , or language structures and 2 ) the source code is complex with complicated logic and variable grammatical structure , making it hard to learn the semantics . Conventionally , information retrieval ( IR ) techniques have been widely used in code summarization ( Eddy et al. , 2013 ; Haiduc et al. , 2010 ; Wong et al. , 2015 ; 2013 ) . Since code duplication ( Kamiya et al. , 2002 ; Li et al. , 2006 ) is common in “ big code ” ( Allamanis et al. , 2018 ) , early works summarize the new programs by retrieving the similar code snippet in the existing code database and use its summary directly . Essentially , the retrieval-based approaches transform the code summarization to the code similarity calculation task , which may achieve promising performance on similar programs , but are limited in generalization , i.e . they have poorer performance on programs that are very different from the code database . ∗Contact : shangqin001 @ e.ntu.edu.sg †Corresponding authors To improve the generalization performance , recent works focus on generation-based approaches . Some works explore Seq2Seq architectures ( Bahdanau et al. , 2014 ; Luong et al. , 2015 ) to generate summaries from the given source code . The Seq2Seq-based approaches ( Iyer et al. , 2016 ; Hu et al. , 2018a ; Alon et al. , 2018 ) usually treat the source code or abstract syntax tree parsed from the source code as a sequence and follow a paradigm of encoder-decoder with the attention mechanism for generating a summary . However , these works only rely on sequential models , which are struggling to capture the rich semantics of source code e.g. , control dependencies and data dependencies . In addition , generation-based approaches typically can not take advantage of similar examples from the retrieval database , as retrieval-based approaches do . To better learn the semantics of the source code , Allamanis et al . ( Allamanis et al. , 2017 ) lighted up this field by representing programs as graphs . Some follow-up works ( Fernandes et al. , 2018 ) attempted to encode more code structures ( e.g. , control flow , program dependencies ) into code graphs with graph neural networks ( GNNs ) , and achieved the promising performance than the sequencebased approaches . Existing works ( Allamanis et al. , 2017 ; Fernandes et al. , 2018 ) usually convert code into graph-structured input during preprocessing , and directly consume it via modern neural networks ( e.g. , GNNs ) for computing node and graph embeddings . However , most GNN-based encoders only allow message passing among nodes within a k-hop neighborhood ( where k is usually a small number such as 4 ) to avoid over-smoothing ( Zhao & Akoglu , 2019 ; Chen et al. , 2020a ) , thus capture only local neighborhood information and ignore global interactions among nodes . Even there are some works ( Li et al. , 2019 ) that try to address this challenging with deep GCNs ( i.e. , 56 layers ) ( Kipf & Welling , 2016 ) by the residual connection ( He et al. , 2016 ) , however , the computation cost can not endure in the program especially for a large and complex program . To address these challenges , we propose a framework for automatic code summarization , namely Hybrid GNN ( HGNN ) . Specifically , from the source code , we first construct a code property graph ( CPG ) based on the abstract syntax tree ( AST ) with different types of edges ( i.e. , Flow To , Reach ) . In order to combine the benefits of both retrieval-based and generation-based methods , we propose a retrieval-based augmentation mechanism to retrieve the source code that is most similar to the current program from the retrieval database ( excluding the current program itself ) , and add the retrieved code as well as the corresponding summary as auxiliary information for training the model . In order to go beyond local graph neighborhood information , and capture global interactions in the program , we further propose an attention-based dynamic graph by learning global attention scores ( i.e. , edge weights ) in the augmented static CPG . Then , a hybrid message passing ( HMP ) is performed on both static and dynamic graphs . We also release a new code summarization benchmark by crawling data from popular and diversified projects containing 95k+ functions in C programming language and make it public 1 . We highlight our main contributions as follows : • We propose a general-purpose framework for automatic code summarization , which combines the benefits of both retrieval-based and generation-based methods via a retrieval-based augmentation mechanism . • We innovate a Hybrid GNN by fusing the static graph ( based on code property graph ) and dynamic graph ( via structure-aware global attention mechanism ) to mitigate the limitation of the GNN on capturing global graph information . • We release a new challenging C benchmark for the task of source code summarization . • We conduct an extensive experiment to evaluate our framework . The proposed approach achieves the state-of-the-art performance and improves existing approaches by 1.42 , 2.44 and 1.29 in terms of BLEU-4 , ROUGE-L and METEOR metrics . 2 HYBRID GNN FRAMEWORK . In this section , we introduce the proposed framework Hybrid GNN ( HGNN ) , as shown in Figure 1 , which mainly includes four components : 1 ) Retrieval-augmented Static Graph Construction ( c.f. , Section 2.2 ) , which incorporates retrieved code-summary pairs to augment the original code for learning . 2 ) Attention-based Dynamic Graph Construction ( c.f. , Section 2.3 ) , which allows message passing among any pair of nodes via a structure-aware global attention mechanism . 3 ) HGNN , ( c.f. , 1https : //github.com/shangqing-liu/CCSD-benchmark-for-code-summarization Source Code Retrieval Database Section 2.4 ) , which incorporates information from both static graphs and dynamic graphs with Hybrid Message Passing . 4 ) Decoder ( c.f. , Section 2.5 ) , which utilizes an attention-based LSTM ( Hochreiter & Schmidhuber , 1997 ) model to generate a summary . 2.1 PROBLEM FORMULATION . In this work , we focus on generating natural language summaries for the given functions ( Wan et al. , 2018 ; Zhang et al. , 2020 ) . A simple example is illustrated in Listing 1 , which is crawled from Linux Kernel . Our goal is to generate the best summary “ set the time of day clock ” based on the given source code . Formally , we define a dataset as D = { ( c , s ) |c ∈ C , s ∈ S } , where c is the source code of a function in the function set C and s represents its targeted summary in the summary set S. The task of code summarization is , given a source code c , to generate the best summary consisting of a sequence of tokens ŝ = ( t1 , t2 , ... , tT ) that maximizes the conditional likelihood ŝ = argmaxsP ( s|c ) . Source Code : int pdc_tod_set ( unsigned long sec , unsigned long usec ) { int retval ; unsigned long flags ; spin_lock_irqsave ( & pdc_lock , flags ) ; retval = mem_pdc_call ( PDC_TOD , PDC_TOD_WRITE , sec , usec ) ; spin_unlock_irqrestore ( & pdc_lock , flags ) ; return retval ; } Ground-Truth : set the time of day clock Listing 1 : An example in our dataset crawled from Linux Kernel . 2.2 RETRIEVAL-AUGMENTED STATIC GRAPH . 2.2.1 GRAPH INITIALIZATION . The source code of a function can be represented as Code Property Graph ( CPG ) ( Yamaguchi et al. , 2014 ) , which is built on the abstract syntax tree ( AST ) with different type of edges ( i.e. , Flow To , Control , Define/Use , Reach ) . Formally , one raw function c could be represented by a multi-edged graph g ( V , E ) , where V is the set of AST nodes , ( v , u ) ∈ E denotes the edge between the node v and the node u . A node v consists of two parts : the node sequence and the node type . An illustrative example is shown in Figure 2 . For example , in the red node , a % 2 == 0 is the node sequence and Condition is the node type . An edge ( v , u ) has a type , named edge type , e.g. , AST type and Flow To type . For more details about the CPG , please refer to Appendix A. Initialization Representation . Given a CPG , we utilize a BiLSTM to encode its nodes . We represent each token of the node sequence and each edge type using the learned embedding matrix Eseqtoken and Eedgetype , respectively . Then nodes and edges of the CPG can be encoded as : h1 , ... , hl = BiLSTM ( E seqtoken v,1 , ... , E seqtoken v , l ) encode_node ( v ) = [ h→l ; h ← 1 ] encode_edge ( v , u ) = Eedgetypev , u if ( v , u ) ∈ E else 0 ( 1 ) Entry CFGEntryNode example ( ) FunctionDef int a = rand ( ) IdentifierDeclStatement a % 2 == 0 Condition int b = a ++ IdentifierDeclStatement EXIT CFGExitNode call ( b ) ExpressionStatement ParameterList void ReturnType int IdentifierDeclType a Identifier a = rand ( ) AssignmentExpr rand ( ) CallExpression a % 2 == 0 EqualityExpression a % 2 MultiplicativeExpression a Identifier 2 PrimaryExpression 0 PrimaryExpression int IdentifierDeclType b Identifier b = a ++ AssignmentExpr a ++ IncDecOp ++ IncDec a Identifier a Symbol call Callee b Argument b Identifierb Symbol b Identifier AST Flow To Control Define/Use Reach void example ( ) { int a = rand ( ) ; if ( a % 2 == 0 ) { int b = a++ ; call ( b ) ; } } A simple code example parse Figure 2 : An example of Code Property Graph ( CPG ) . where l is the number of tokens in the node sequence of v. For the sake of simplicity , in the following section , we use hv and ev , u to represent the embedding of the node v and the edge ( v , u ) , respectively , i.e. , encode_node ( v ) and encode_edge ( v , u ) . Given the source code c of a function as well as the CPG g ( V , E ) , Hc ∈ Rm×d denotes the initial node matrix of the CPG , where m is the total number of nodes in the CPG and d is the dimension of the node embedding .

Summary: This paper proposes a retrieval-augmented method for generating code summarization. The model encodes the input code based on its graph structure (Code Property Graph) with a hybrid GNN architecture. The model augments the initial graph representation of the input code based on the representation of the top-1 retrieval result. It also augments the final graph encoding with the BiLSTM encoding of the retrieved summary.