Year: 2019

Abstract:
Program comprehension is a fundamental task in software development and maintenance processes. Software developers often need to understand a large amount of existing code before they can develop new features or fix bugs in existing programs. Being able to process programming language code automatically and provide summaries of code functionality accurately can significantly help developers to reduce time spent in code navigation and understanding, and thus increase productivity. Different from natural language articles, source code in programming languages often follows rigid syntactical structures and there can exist dependencies among code elements that are located far away from each other through complex control flows and data flows. Existing studies on tree-based convolutional neural networks (TBCNN) and gated graph neural networks (GGNN) are not able to capture essential semantic dependencies among code elements accurately. In this paper, we propose novel tree-based capsule networks (TreeCaps) and relevant techniques for processing program code in an automated way that encodes code syntactical structures and captures code dependencies more accurately. Based on evaluation on programs written in different programming languages, we show that our TreeCaps-based approach can outperform other approaches in classifying the functionalities of many programs.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 APPROACH OVERVIEW
text: 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.

heading: 4 ABSTRACT SYNTAX TREE VECTORIZATION
text: 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.

heading: 5 TREE-BASED CAPSULE NETWORKS
text: 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.

heading: 5.1 PRIMARY VARIABLE TREECAPS LAYER
text: 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 .

heading: 5.2 PRIMARY STATIC TREECAPS LAYER
text: The key issue with passing the outputs of the PVC layer, Xpvc, to the Code Capsule layer is that the number of capsules, Npvc, is variable from one training example to another. Prior to routing the lower level capsules to a set of higher level capsules, the lower dimensional capsule outputs need to be projected to the higher dimensionality, with the aid of the transformation matrix which learns the part-whole relationship between the lower and the higher level capsules (Sabour et al., 2017). However, a trainable transformation matrix cannot be defined in practice with variable dimensions. Thus, the dynamic routing in the literature (Sabour et al., 2017) cannot be applied between a variable set of capsules and a static set of capsules.

heading: 5.2.1 VARIABLE-TO-STATIC CAPSULE ROUTING
text: Therefore, we propose a novel variable-to-static capsule routing algorithm, summarized in Algo. 1. Algorithm 1 Variable-to-Static Capsule Routing
1: procedure ROUTING(ûi, r, a, b) 2: Ûsorted ← sort([û1, ..., ûNpvc ]) 3: Initialize vj : ∀i, j ≤ a,vj ← Ûsorted[i] 4: Initialize αij : ∀j ∈ [1, a],∀i ∈ [1, b], αij ← 0 5: for r iterations do 6: ∀j ∈ [1, a],∀i ∈ [1, b], fij ← ûi · vj 7: ∀j ∈ [1, a],∀i ∈ [1, b], αij ← αij + fij 8: ∀i ∈ [1, b],βi ← Softmax(αi) 9: ∀j ∈ [1, a], sj ← ∑ iβijûi
10: ∀j ∈ [1, a],vj ← Squash(sj) 11: return vj
We initialize the outputs of the Primary Static Capsule layer with the outputs of the a capsules with the highest L2 norms in the PVC layer. Hence, the outputs of the PVC layer, [û1, ..., ûNpvc ], are first sorted by their L2 norms, to obtain Ûsorted, and then the first a vectors of Ûsorted are assigned as vj , j ≤ a. The intuition is that, in practice, not every node of the AST contributes towards source code classification. Often, source code consists of non-essential entities, and only a portion of all entities determine the code class. Since the probability of existence of an entity is denoted by the length of the capsule output vector (L2 norm), we only consider the entities with the highest existence probabilities for initialization. It should be noted that the capsules with the a-highest norms are used only for initialization, the actual outputs of the primary static capsules are determined by iteratively running the variable-to-static routing algorithm.
A well-known property of source code is that dependency relationships may exist among entities that are not spatially co-located. Therefore, we route b nodes in the AST based on the similarity between them and the primary static capsule layer outputs, where a ≤ b ≤ Npvc. We assign b = Npvc in general to route with all the nodes in the AST. If computational complexity is critical, we can choose a smaller b and route with top-b nodes of the AST. a and b can be chosen empirically, where computational complexity also factors in when choosing b.
We initialize the routing coefficients as αij = 0, equally to all the capsules in the primary variable capsule layer. Subsequently, as illustrated by Fig 3, they are iteratively refined based on the agreement between the current primary static capsule layer outputs vj and the primary dynamic capsule layer outputs ûi. The agreement in this case is measured by the dot product, fij ← ûi · vj , and the routing coefficients are adjusted with fij accordingly. If a capsule γ in the primary dynamic layer has a strong agreement with a capsule δ in the primary static layer, then fγδ will be positively large, whereas if there is a strong disagreement, then fγδ will be negatively large. Subsequently, the sum of vectors ûi is weighted by the updated βij to calculate sj , which is then squashed to update vj .

heading: 5.3 CODE CAPSULE LAYER
text: Code Capsule layer is the final layer of the TreeCaps network, which acts as the classification capsule layer, as illustrated by Figure 4. Since the outputs of the PSC layer Xpsc ∈ RNpsc×Dpsc , where Npsc = a and Dpsc = Dpvc, consist of a fixed set of capsules, it can be routed to the CC layer via the dynamic routing algorithm in the literature (Sabour et al., 2017) (summarized in Algo. 2). For
each capsule j in the PSC layer, and for each capsule m in the CC layer, we multiply the output of the primary static capsule vj by the transformation matrices Wjm to produce the prediction vectors v̂m|j = Wjmvj . The trainable transformation matrices learn the part-whole relationships between the primary static capsules and the code capsules, while effectively transforming vj’s into the same dimensionality as zm, where zm’s denote the outputs of the code capsule layer. Similar to the variable-to-static capsule routing, we initialize the routing coefficients γjm equally, and iteratively refine them based on the agreements between the prediction vectors v̂m|j and the code capsule outputs zm, where zm = squash( ∑ j γjmv̂m|j).
Algorithm 2 Dynamic Routing 1: procedure ROUTING(v̂j , t, a, c) 2: Initialize ∀j ∈ [1, a],∀m ∈ [1, c], δjm ← 0 3: for t iterations do 4: ∀j ∈ [1, a], γj ← softmax(δj) 5: ∀m ∈ [1, c], zm ← squash( ∑ j γjmv̂m|j)
6: ∀j ∈ [1, a],∀m ∈ [1, c], δjm ← δjm + v̂m|j · zm 7: return ẑm
The primary static capsule outputs are routed to the CC layer using the dynamic routing algorithm, as illustrated by Fig 4, to produce the final capsule outputs Xcc ∈ RNcc×Dcc , where Ncc = κ is the number of classes and Dcc is the dimensionality of the code capsules. Ultimately, we calculate the probability of existence of each class by obtaining L2 norm of each CodeCaps output vector.

heading: 5.4 MARGIN LOSS FOR TREECAPS TRAINING
text: We use the Margin Loss proposed by Sabour et al. (2017) as the loss function for TreeCaps. For every code capsule µ, the margin loss Lµ is defined as follows,
Lµ = Tµmax(0,m + − ‖vµ‖)2 + λ(1− Tµ)max(0, ‖vµ‖ −m−)2 (4)
where Tµ is 1 if the correct class is µ and zero otherwise. Following Sabour et al. (2017), λ is set to 0.5 to control the initial learning from shrinking the length of the output vectors of all the code capsules, and m+,m− are set to 0.9, 0.1 as the lower bound for the correct class and the upper bound for the incorrect class respectively.

heading: 6 EMPIRICAL EVALUATION

heading: 6.1 DATASETS AND IMPLEMENTATION
text: We used three datasets in three programming languages to ensure cross-language robustness. The first dataset (A) contains 6 classes of sorting algorithms, with 346 training programs on average per class, written in Python.2 The second dataset (B) is inherited from BUI et al. (2019), which contains 10 classes of sorting algorithms, with 64 training programs on average per class, written in Java. The third dataset (C) is inherited from Mou et al. (2016), which contains 104 classes of C programs, with 375 training programs on average per class. For the dataset A, we used the publicly available vectorizer (see Footnote 2) to generate embeddings for more than 90 AST node types in Python. For the datasets B & C, srcML3 defines a unified vocabulary for more than four hundred AST node types for C and Java (and several other languages, but not Python), and we adapted the same vectorizer to generate embeddings for the unified AST node types defined by srcML.
We used Keras and Tensorflow libraries to implement TreeCaps. To train the models, we used the RAdam optimizer (Liu et al., 2019) with an initial learning rate of 0.001 subjected to decay, on an Nvidia Tesla P100 GPU. To enhance the classification accuracies, a weighted average ensembling technique (Krogh & Vedelsby, 1995) was used.

heading: 6.2 PROGRAM CLASSIFICATION RESULTS
text: Table 1 compares our results to other approaches for program classification. It should be noted that, Mou et al. (2016) have used custom-trained initial embeddings for a small set of about 50 AST node types defined specifically for C language only (Peng et al., 2015) and reported a higher result in their paper, while our approach generates the initial embeddings for a much larger vocabulary of more than three hundred unified AST node types for both C and Java. For a fairer comparison based on the same set of AST node vocabulary, especially for the datasets B & C, we used our embeddings based on srcML node vocabulary as the initial embeddings across all models. We followed the techniques proposed in Allamanis et al. (2018a) and BUI et al. (2019) to re-generate the results for GGNN and the techniques proposed in Mou et al. (2016) to re-generate the results for TBCNN.
For the dataset A, we achieved a perfect classification result, outperforming TBCNN by a narrow margin of less than 1%. However, the margin was more significant for the datasets B and C. An average accuracy of 92.11% was achieved by our approach for the dataset B, outperforming other approaches at least by 7.11%. TreeCaps outperformed its convolutional counterpart (TBCNN) by a significant margin of 17.11%. The performance was further improved by 1.97% with the use of 3- model weighted average ensembling technique. For the dataset C, our approach was able to surpass the other approaches by 1.43%, achieving an accuracy of 87.95%. However, TreeCaps surpassed TBCNN by a more significant margin of 8.55%. Three-model weighted average ensembling on the dataset C provided a further improvement of 2.89% in comparison to the other approaches, achieving an accuracy of 89.41%.

heading: 6.3 MODEL ANALYSIS
text: We evaluate the effects of various aspects of the TreeCaps model, including the effect of the variableto-static routing algorithm, variations in the number of instantiation parameters in the CodeCaps layer, and the addition of a secondary capsule layer. We evaluate these variations on the dataset B.

heading: 6.3.1 EFFECT OF THE VARIABLE-TO-STATIC ROUTING ALGORITHM
text: We investigate the effect of the variable-to-static routing algorithm by replacing it with Dynamic Max Pooling (DMP). Since there is no alternative approach existing in the literature for routing a variable set of capsules to a static set of capsules, we compare the proposed routing algorithm with dynamic pooling. The output of the PVC layer, Xpvc ∈ RNpvc×Dpvc consists of a variable component, Npvc. Using dynamic max pooling across all theNpvc capsules will result in one output capsule, Xdmp ∈ R1×Dpvc . Since Xdmp has no variable components across the training samples,
2Collected from https://github.com/crestonbunch/tbcnn. 3https://www.srcml.org/
it can now be routed to the code capsules using the dynamic routing algorithm. However, it should be noted that DMP is not suitable for capsule networks, as it destroys the spatial and dependency relationships between the capsules. We use DMP here only for comparison purposes.
As summarized in Table 2, DMP yields a considerably lower accuracy of 83.43% than our routing algorithm by a significant margin of 8.68%, establishing the effectiveness of our proposed algorithm.

heading: 6.3.2 EFFECT OF THE NUMBER OF INSTANTIATION PARAMETERS
text: The instantiation parameters Dcc of the Code Capsule layer acts as the final embeddings used for classification, in other words, the dimensionality of the latent representation of source code. If the dimensionality of the latent representation is higher than required, it can introduce sparsity and/or correlations between the instantiation parameters, reducing the classification accuracy. On the contrary, if the dimensionality of the latent representation is too low, it may not be sufficient to capture the variations in source code, leading to under-representation, reducing the classification accuracy. Hence, in an attempt to identify a suitable value for Dcc for source code classification, we investigate the effect of Dcc in the accuracy. As summarized in Table 2, we observed that the most suitable value was Dcc = 8 for the dataset B.

heading: 6.3.3 EFFECT OF THE ADDITION OF A SECONDARY CAPSULE LAYER
text: We evaluate the addition of an extra capsule layer functionally similar to a primary static capsule layer, which we call the secondary capsule (SC) layer. With respect to Fig 1, we stack the SC layer in between the PSC layer and the CC layer. We use dynamic routing to route between the PSC and SC layers and between the SC and CC layers. Even though we observed a minor improvement of the classification accuracy, the added computational complexity increases the inference time by 16%, from 10.3ms to 11.9ms per sample. The usefulness of the addition of such a SC layer needs to be further investigated.

heading: 6.4 DISCUSSION OF LIMITATIONS & FUTURE WORK
text: Since TreeCaps is based on the capsule networks, it inherits the limitations of capsule networks such as the high computational complexity in comparison to CNNs, and performance reduction with the increasing number of classes. TreeCaps lacks a reconstruction network, similar to the reconstruction network for image data presented by Sabour et al. (2017), which is useful to investigate the interpretability of the capsule network, including the relationship between the learnt instantiation parameters and the physical attributes of data.
We intend to extend our work to further investigate the effects of different initial embeddings on the classification accuracy. Further, we intend to compare related pieces of code identified by TreeCaps to program dependencies identified by program analysis techniques, to evaluate the effectiveness of TreeCaps as an embedding generating technique, and to extend TreeCaps to other related tasks such as bug detection and localization.

heading: 7 CONCLUSION
text: In this paper, we proposed a novel tree-based capsule network (TreeCaps) to learn rich syntactical structures and semantic dependencies in program source code. The model proposed novel technical features that deal with variable sized trees across different programming languages, including primary variable and primary static capsule layers, and the variable to static routing algorithm. Our empirical evaluations show that these features significantly contribute to the high classification accuracy of TreeCaps model for program classification tasks for various programming languages. To the best of our knowledge, our work is the first to adapt capsule networks to trees and apply them to program source code learning. We believe that TreeCaps can capture more code semantics than previous code learning models and complement existing program analysis techniques well.

Review summaries:
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.