GN-Transformer: Fusing AST and Source Code information in Graph Networks

1 INTRODUCTION . Code summarization is the task of generating a readable summary that describes the functionality of a snippet . Such task requires a high-level comprehension of a source code snippet thus it is an effective task to evaluate whether a Deep Learning Model is able to capture complex relations and structures inside code . Programming languages are context-free formal language , an unambiguous representation , Abstract Syntax Tree ( AST ) , could be derived from a source code snippet . A parse tree based representation of code is precise and without noise . An AST accurately describes the structure of a snippet and relationships between tokens which provides valuable supplementary information for code understanding . Using graph representations of source code has been the focus of multiple methods that perform code summarization . For example , Alon et al . ( 2019a ) encoded AST paths between tokens and aggregated them by an attention mechanism . Huo et al . ( 2020 ) used CNN and DeepWalk on Control Flow Graph ( CFG ) to learn statements representations . LeClair et al . ( 2020 ) applied Graph Convolutional Networks ( GCNs ) to learn AST representation . These methods proposed different ways to extract AST features , however the cross-modal interaction ( Veličković , 2019 ) is very limited since the AST and code features are independently extracted by separate models then simply concatenated or summed . In this paper we propose a novel architecture GN-Transformer shown in Figure 2 to fuse Graph information with an equivalent sequence representation . In summary : • We extend Graph Networks ( GN ) ( Battaglia et al. , 2018 ) to a novel GN-Transformer architecture that is a sequence of GN encoder blocks followed by a vanilla Transformer decoder . • We propose a novel method for early fusion of the AST representation and that of a code snippet sequence called Syntax-Code Graph ( SCG ) • We evaluate our approach on the task of code summarization and outperform the previous state of the art in two datasets and across three metrics . We evaluated our model on Java and Python datasets used by Ahmad et al . ( 2020 ) . We compared our results to those of Ahmad et al . ( 2020 ) . Two qualitative results are presented in Figure 1 . We make available our code , trained models and pre-processed datasets in our supplementary package , and we will open-source it after the review process concludes . 2 FUSING GRAPH AND SEQUENCE INFORMATION . Previous methods consider sequences and graphs as two modalities that are processed independently . For a sequence , recurrent architectures such as RNNs , LSTMs , GRUs are commonly used . CNNs have also been applied on sliding windows of sequences ( Kim , 2014 ) . Transformers ( Vaswani et al. , 2017 ) became a popular choice for sequences in recent years . For graph data , spectrum-based methods like GCNs ( Bruna et al. , 2014 ) capture graph structure through a spectrum . Non-spectrum methods like GraphSAGE ( Hamilton et al. , 2017 ) aggregates information from neighboring nodes using different aggregators , GAT ( Veličković et al. , 2018 ) introduced attention mechanism to aggregate neighboring information . Early fusion of multiple modalities is a challenging task . As a result , late fusion methods are used when considering multi-modal information in code summarization tasks . The cross-modal interactions are less efficient in late fusion as compared to early fusion Veličković ( 2019 ) . In Section 2.1 we discuss early fusion approaches of code sequence with an AST . In Section 2.2 we discuss representing sequence in a graph . 2.1 EARLY FUSION OF SEQUENCE AND GRAPH . For early fusion of a sequence and a graph , it is common to represent them under a single unified representation and input to a deep learning model . Random walks ( Perozzi et al. , 2014 ) and structure-based traversal ( SBT ) ( Hu et al. , 2018a ) demonstrated the advantages of flattening a graph structure into a sequence . SBT converts AST into a sequence that can be used by seq2seq architectures . However , structural information contained in the AST is lost when flattened into an unstructured sequence . The main motivation of previous methods to flatten a graph structure into a sequence instead of the opposite is due to the power of sequence models when compared to their graph counter-parts . Recent advances in general frameworks such as Graph Networks ( Battaglia et al. , 2018 ) , MessagePassing Neural Networks ( MPNN ) ( Gilmer et al. , 2017 ) proposed a unified graphical representation of data . Battaglia et al . ( 2018 ) proposed an extension and generalization of previous approaches that can learn a graphical representation of data under a configurable computation unit , GN block . We discuss our GN block based encoder in Section 4.2 . There are several benefits in graph representation . Firstly it can contain different information sources with arbitrary relational inductive bias ( Battaglia et al. , 2018 ) . Secondly , a graph representation can have explicit relational inductive biases among graph entities . Thirdly , flexible augmentation of graph structure of input through expert knowledge . Additionally , better combinatorial generalization ( Battaglia et al. , 2018 ) due to the reuse of multiple information sources simultaneously in a unified representation . Finally , measures for analyzing performance through graph structures could be naturally incorporated , such as average path length L and clustering coefficient C. We discuss it in better detail in Section 3.2 . 2.2 GRAPH REPRESENTATIONS FOR SEQUENCE . In Graph Networks , defining the graph structure of input data is the main way to introduce relational inductive biases ( Battaglia et al. , 2018 ) . Relational inductive biases impose constraints on the relationship among entities . In Graph Networks , the assumption that two entities have a relation or interaction is expressed by an edge between the corresponding entity nodes . The absence of an edge expresses the assumption of isolation which means no relation or direct influence of two nodes . Sequences are unstructured data in which relationships are implicit . Implicit relationships can be represented with each token as a graph node that is fully connected with all other nodes . Such representation allows each node to interact with every other node with no explicit assumptions on isolation . Thus it allows the model to infer the relationships between them implicitly . Transformers could be regarded as inferring on such a fully connected graph , see Figure 3 ( a ) . Each token in an input sequence corresponds to a node . The relationship between tokens is thus represented by attention weights , high attention values correspond to strong interaction while low attention means isolation . It is less efficient for a model to learn to infer relationships without any explicit assumptions of interactions and isolation . AST provides precise information about interaction and isolation among tokens in source code since it ’ s a representation without noise on how the tokens interact during the execution of a code snippet . Thus a natural way of utilizing information from an AST is fusing the graph structure with an input sequence . We can find an explicit mapping between tokens in a sequence and nodes in the AST through the scope information provided by a parser for a given programming language . We are then able to find relations thus build edges between a sequence and a graph . We ’ ll further discuss the graph structure of input data and how we fuse AST with it in Section 3 . The idea of fusing a sequence and a graph can be extended to broader cases apart from AST and code snippet when the mapping is not explicit . There are techniques that are applied in knowledge graphs to find a mapping between entities in knowledge graphs and words in sequence ( Peters et al. , 2019 ) . For our problem statement , we make the reasonable assumption that an explicit mapping between the sequence and a graph can be provided . 3 GRAPH STRUCTURE OF AST AND CODE . First , we propose a simple joint graph representation of source code and AST called standard SyntaxCode Graph ( SCG ) in Section 3.1 . In Section 3.2 we discuss the influence of the structure of SCG and the existence of a theoretically optimal graph structure . 3.1 SYNTAX-CODE GRAPH . Standard SCG consists of AST nodes directly from AST and token nodes created for each code snippet token or simply referred to the remainder of the paper as token node . The attribute of an AST node is the type on AST such as “ NameExpr ” . The attribute of a token node is the identifier name such as “ a ” , “ int ” , “ + ” . Standard SCG preserves the AST structure and introduces additional edges that connect token nodes with their direct parent node in AST . The direct parent node for a token node is found through the scope information for AST nodes . Taking the statement ‘ a [ i ] =4+2 ’ as an example , the AST is shown in Figure 3 ( b ) . The scope of an AST node is determined by the positional mark given by a compiler . The positional mark of the AST node ’ AdditiveExpr ’ is line 1 , col 6 ∼ line 1 , col 8 which corresponds to ’ 4+2 ’ , so the scope covers token nodes ’ 4 ’ , ’ + ’ , and ’ 2 ’ . However , a token may be covered by the scopes of multiple AST nodes , we only connect the token with its direct parent in AST . The direct parent of a token is the deepest AST node among all AST nodes which scope covered the token . Figure 3 ( c ) shows standard SCG . Standard SCG directly builds on top of the AST structure without introducing any additional relational assumptions . Thus it objectively reflects the program structure depicted by a compiler . 3.2 OPTIMAL SYNTAX-CODE GRAPH STRUCTURE . The graph structure of SCG plays a critical role in our method . You et al . ( 2020 ) represented neural network structure as relational graphs . Message passing in a relational graph is analogous to information propagation in Graph Networks . SCG can be directly defined as the relational graph of our model . Node features could be defined as node embedding in SCG . Likewise , the message function is node-wise FFN and the aggregation function is MHA . As a result , the graph structure of SCG will influence the performance of the model . The standard SCG structure discussed in Section 3.1 is efficient but may not be optimal . Qualitatively , we observed two problems of standard SCG . The first problem is the existence of long range dependencies . It is hard to pass information between nodes separated by long paths in SCG . In Graph Networks , each node is executing an information propagation and aggregation simultaneously with neighboring nodes . For example , One graph neural networks ( GNN ) layer could be regarded as executing one turn of information propagation and aggregation . In Figure 3 ( c ) , it would require three GNN layers for information to propagate from leaf node ‘ a ’ to root node ‘ AssignExpr ’ . Thus it is difficult to propagate information between nodes in a large AST . The second problem is the isolation among token nodes . Source code token nodes can only indirectly interact with each other through AST nodes . Taking ‘ 4+2 ’ in Figure 3 ( c ) as an example , there is no direct edge between token nodes ‘ 4 ’ and ‘ 2 ’ , they can only indirectly interact through the AST node ‘ AdditiveExpr ’ . However , the information that will be passed from ’ AdditiveExpr ’ to ’ 4 ’ and ’ 2 ’ is identical . Therefore , it makes it difficult for the token nodes to learn more complex relationships with other token nodes when they are under the same expression . We analyzed the test set of the Java dataset and present the results in Figure 4 ( a ) . Longer code usually contains more long range dependencies and complex expressions that correspond to the above two problems . Hand-engineering graph structure with expert knowledge may alleviate these problems thus improve performance . However , introducing redundant edges could be lead to performance degradation . We perform two experiments on opposite sides of the spectrum in our graph structure to examine the problems above . Variant 1 added shortcut edges between a code token and AST nodes with the scope covered by that token thus the distance between all AST nodes to token nodes within its scope was shortened to 1 . Variant 2 makes token nodes fully connected , thus there is no isolation at all . We reference to those variants as Variant 1 and Variant 2 in the remainder of this text . More details on the two variants are described in Appendix D. Both variants failed to improve performance due to loss of structural information and redundant connections respectively , leading us to conclude that a trade-off of both is required . We discuss it in further detail at Section 5.2 . The optimal graph structure was also analyzed through quantitative experiments . You et al . ( 2020 ) proposed to measure the effectiveness of the message passing in a relational graph by the Average path length ( L ) and the clustering coefficient ( C ) . The authors claim that the optimal structure is a balance between the value of C and L. Figure 4 ( b ) shows L and C for our variants , a Standard AST , and that of a fully-connected relational graph such as in a Transformer which could be formulated into a relational graph by the same way of ours with different relational graph structure . A vanilla Transformer encoder has a maximum C of 1 and a minimum L of 1 . Our standard SCG , at the other extreme , has a minimum C of 0 and a relatively high L due to its tree structure . The average L of standard SCG in Java and Python dataset is 6.28 and 6.35 respectively . Figure 4 ( b ) shows that we only explored a small area of the entire design space of the possible graph structure for various L and C. You et al . ( 2020 ) proposed that the optimal structure is usually located around a ” sweet spot ” in design space between extreme cases of the tree and fully-connected structures . Model performance will improve nearing this ” sweet spot ” . Thus we conclude that there is a large potential of improving model performance by improving the input structure either through hand-engineering or additional rules tailored for the specific problem domain .

Summary:  In this work, authors propose a new direction in summarizing code snippets by combining their AST and lexical code features in the form of a graph called SCG (which is shown not to be optimal). Their model, GN-Transformer, further extracts information from SCG to summarize the code snippet. This model is a combination of transformers and graph network and is benchmarked on two datasets of Java and Python source code. The results show that the model marginally outperforms the selected benchmarks.