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: This paper focuses on the problem of training a neural model to understand source code. The authors argue that both graph information (such as the parsed abstract syntax tree) and sequence information (such as the raw program tokens) are useful for understanding code, and describe a particular method of adding raw program tokens to a graph called SCS. They also describe a modification of a transformer (called a GN-Transformer) that uses this SCS graph representation, and present results on code summarization tasks.