TreeCaps: Tree-Structured Capsule Networks for Program Source Code Processing

1 INTRODUCTION . Understanding program code is a fundamental step for many software engineering tasks . Software developers often spend more than 50 % of their time in navigating through existing code bases and understanding the code before they can implement new features or fix bugs ( Xia et al. , 2018 ; Evans Data Corporation , 2019 ; Britton et al. , 2012 ) . If suitable models for programs are built , they can be useful for many tasks , such as classifying the functionality of programs ( Nix & Zhang , 2017 ; Dahl et al. , 2013 ; Pascanu et al. , 2015 ; Rastogi et al. , 2013 ) , predicting bugs ( Yang et al. , 2015 ; Li et al. , 2017 ; 2018 ) , and providing bases for program translation ( Nguyen et al. , 2017 ; Gu et al. , 2017 ) . Different from natural language texts , programming languages have clearly defined grammars and compilable source code must follow rigid syntactical structures and can be unambiguously parsed into syntax trees . There can be complex control flows and data flows among various code elements all over a program that affect the semantic and functionality of the program . Some inter-dependent code elements can appear in an arbitrary order in the program ( e.g. , a function A calls another function B while A and B are spatially far away from each other ) ; some code elements , such as local variable names , have no significant impact on code functionality . In the literature , tree-based convolutional neural networks ( TBCNNs ) have been proposed ( Peng et al. , 2015 ; Mou et al. , 2016 ) to show promising results in programming language processing . TBCNNs accept Abstract Syntax Trees ( ASTs ) of source code as the input , and capture explicit , structural parent-child-sibling relations among code elements . Gated graph neural networks ( GGNNs ) ( Li et al. , 2016 ) are also proposed as a way to learn graphs , and ASTs are extended to graphs with a variety of code dependencies added as edges among tree nodes to model code semantics ( Allamanis et al. , 2018b ) . While GGNNs capture more code semantics than TBCNNs , many additional edges among tree nodes have to be added through program analysis techniques , and many of the edges may be noise , contributing longer training time and lower performance . A recent model , known as ASTNN , based on a sequence of small ASTs for statements ( instead of one big AST for the whole program ) shows better performance than TBCNNs and GGNNs ( Zhang et al. , 2019 ) . In this paper , we propose a novel tree-based capsule network architecture , named TreeCaps , to capture both syntactical structure and dependency information in code , without the need of explicitly adding dependencies in the trees or splitting a big tree into smaller ones . Capsule Networks ( CapsNet ) ( Sabour et al. , 2017 ) itself is a promising concept that has demonstrated vast potential in outperforming CNNs in various domains including computer vision and natural language processing , because it has the main advantage that it can discover and preserve the relative spatial and hierarchical relations among objects within an input ( e.g. , an image and a piece of texts ) . TreeCaps adapts CapsNet to ASTs for programs , proposes novel primary variable and primary static capsule layers , proposes a novel variable-to-static routing algorithm to route a variable set of capsules to generate a static set of capsules ( with the intention of preserving code dependencies ) , and connects the tree capsule layers to a classification layer to classify the functionality of a program . In our empirical evaluation , we take various sets of programs written in different programming languages ( e.g. , Python , Java , C ) collected from GitHub and the literature , and train our TreeCaps models to classify programs with different functionalities . Results show that TreeCaps outperforms other approaches in program classification by significant margins , while a study done on variants of the proposed model reveals the effectiveness of the proposed variable-to-static routing algorithm , effect of the dimensionality of the classification capsules on the model performance and the effectiveness of the use of additional capsule layers . 2 RELATED WORK . Capsule networks ( Sabour et al. , 2017 ; Hinton et al. , 2018 ) use dynamic routing to model spatial and hierarchical relations among objects in an image . The techniques have been successfully applied to different tasks , such as computer vision , character recognition , and text classification ( Jayasundara et al. , 2019 ; Rajasegaran et al. , 2019 ; Zhao et al. , 2018 ) . None of the studies has considered complex tree data as input . Capsule Graph Neural Network ( Zhang & Chen , 2019 ) has been recently proposed to classify biological and social network graphs , yet , has not been applied to trees for programming languages processing yet . On the other hand , tree- and graph-based neural networks have been studied for program language processing . TBCNNs ( Peng et al. , 2015 ; Mou et al. , 2016 ) have been used to model code syntactical structures . GGNNs ( Li et al. , 2016 ; Allamanis et al. , 2018b ) build dependency graphs from ASTs and use graph neural networks to encode the code dependencies . Variants of TBCNNs and GGNNs are also proposed to represent programs differently and aim to achieve better training accuracy and costs . For example , ASTNN ( Zhang et al. , 2019 ) splits an entire AST into a sequence of smaller ones and uses bidirectional gated recurrent units ( Bi-GRU ) to model the smaller ASTs that represent statements in programs . For another example , bilateral dependency tree-based CNNs ( DTBCNNs ) ( BUI et al. , 2019 ) are used to classify programs across different programming languages . TreeLSTM ( Wei & Li , 2017 ) has also been used to model tree structures in program code . Our work aims to model tree structures too , but designs special dynamic routing with the intention to capture code dependencies without the need of explicit program analysis techniques . More generally , there is huge interest in applying deep learning techniques for various software engineering tasks , such as program classification , bug prediction , code clone detection , program refactoring , translation and even code synthesis ( Allamanis et al. , 2018a ; Alon et al. , 2019 ; Hu et al. , 2018 ; Chen et al. , 2018 ; Pradel & Sen , 2018 ) . We are likely the first to adapt capsule networks for program source code processing to capture both syntactical structures and code dependencies , especially for the problem of program classification . In the future , it can be an exciting area to combine more kinds of semantic-aware code representations ( e.g. , symbolic traces ( Henkel et al. , 2018 ) ) and tailored program analysis techniques with deep learning to improve code learning tasks . 3 APPROACH OVERVIEW . The overview of our TreeCaps approach is summarized in Fig . 1 . The source code of the training sample program is parsed into an AST and vectorized with the aid of Word2Vec ( Mikolov et al. , 2013 ) or a similar technique that considers the AST node types , instead of concrete tokens , as the vocabulary words ( Peng et al. , 2015 ; BUI et al. , 2019 ) . The AST and the vectorized nodes are then fed in to our TreeCaps network , which consists of a Primary Variable Capsule ( PVC ) layer to accommodate the varying number of nodes in the AST . Subsequently , the capsules are routed to the Primary Static Capsule ( PSC ) layer using the proposed variable-to-static routing algorithm , followed by routing with the dynamic routing algorithm to the Code Capsule ( CC ) layer . Acting as the classification capsule layer , Code capsules capture and provide embeddings for the entire training sample , while denoting the probability of existence of the source code classes by the respective vector norms . Finally , a softmax layer is used on the vector norms to output the probabilities for the input code sample to belong to various functionality classes . Section 4 and 5 explain more about the AST vectorization and other major components in TreeCaps . 4 ABSTRACT SYNTAX TREE VECTORIZATION . Fig . 2 illustrates the vectorization of an AST . Every raw source code is parsed with an appropriate parser corresponding to the programming language to generate the AST.1 The AST represents the syntactic structure of the source code with a set of generalized vocabulary words ( i.e. , node type names ) . We use ASTs to train the embedding for node types by applying embedding techniques , such as Word2vec ( Mikolov et al. , 2013 ) , in the context of ASTs using techniques similar to Peng et al . ( 2015 ) , which learns a vectorized vocabulary of node types , xnode ∈ RV , where V is the embedding size . The learned vocabulary can subsequently be used to vectorize each individual node of the ASTs , generating the vectorized ASTs . 5 TREE-BASED CAPSULE NETWORKS . One of the main challenges in creating a tree-based capsule network is that the input of the network is tree-structured ( ASTs in our case ) . Tree-structured data are inherently different from generic image data , Ximg ∈ RH×W×C , where H , W , C are the fixed height , width and the number of channels respectively , or natural language data , Xnlp ∈ RL×E , where L , E are the fixed padded sentence length and the word embedding size respectively . Hence , the network architecture needs to be constructed to accommodate such tree-structured data , Xtree ∈ RT×V , where T , V are the variable tree size ( the number of nodes in the tree ) and the node embedding size respectively . 5.1 PRIMARY VARIABLE TREECAPS LAYER . A further challenge with trees is that the tree size varies from program to program , and the number of children varies from node to node . A naive solution to the problem can be to pad the sizes 1We use python AST parser for python programs , whereas we use srcML parser for C and Java programs . to reach a fixed , pre-defined length , following a similar approach in natural language domain to preserve a fixed sentence length . However , zero-padding is not appropriate in our case due to the degree of variability . For instance , the number of children per node can vary from zero to hundreds , causing challenges in deciding the fixed padding length and introducing sparsity . Mou et al . ( 2016 ) propose a more effective approach termed as continuous binary tree , where the convolution window is emulated as a binary tree , where the weight matrix corresponding to each node is represented as a weighted sum of three fixed matrices Wt , Wl , Wr ∈ RV ′×V and a bias term b ∈ RV ′ , where V ′ is the embedding size after the convolution , and the weighting coefficients are calculated by taking the positional value in to account . Hence , for a convolutional window of depth d in the original AST , and there are K + 1 nodes ( including the parent node ) which belong to that window with vectors [ x1 , ... , xK+1 ] , where xi ∈ RV , then the convolutional output y can be defined as follows , y = tanh ( K+1∑ i=1 [ ηtiW t + ηliW l + ηriW r ] xi + b ) ( 1 ) where ηti , η l i , η r i are weights defined corresponding to the depth and the position of the children nodes : ηti = di − 1 d− 1 ηri = ( 1− ηti ) pi − 1 k − 1 ηli = ( 1− ηti ) ( 1− ηri ) ( 2 ) where di is the depth of the node i in the convolutional windows , pi is the position of the node and k is the total number of node ’ s siblings . The output of tree-structured convolution resembles the input tree structure , Yconv ∈ RT×V ′ . In the Primary Variable Capsule layer , y obtained from Equation 1 corresponds to the output of one convolutional slice . We use ε such slices with different random initializations for W , b , similar to CNNs for image data . Subsequently , as illustrated by Fig 3 , we group the convolutional slices together to form Npvc = T×V ′×ε Dpvc sets of capsules with outputs ui ∈ RDpvc , i ∈ [ 1 , Npvc ] , where Dpvc is the dimensions of the capsules ( i.e. , Dpvc is the number of instantiation parameters of the capsules ) in the PVC layer . In order to vectorize each capsule output uj as ûj ( to represent the probability of existence of an entity by the vector length ) , we subsequently apply a non-linear squash function as follows , ûi = ||ui||2 ||ui||2 + 1 · ui ||ui||2 ( 3 ) where ||ûi||2 ≤ 1 . Hence , the output of the primary variable capsule layer is Xpvc ∈ RNpvc×Dpvc .

Summary: The paper proposes a neural architecture for summarizing trees inspired by capsule networks from computer vision. The authors re-use a tree convolution from previous work for the bottommost layer, and then propose adaptations to the dynamic routing from capsule networks so that it can be applied to variable-sized trees. The paper applies the proposed architecture to three different program classification datasets, which are in three different languages. The paper reports empirical gains compared to two architectures proposed by previous work.