T REE C APS :T REE -S TRUCTURED C APSULE N ETWORKS FOR P ROGRAM S OURCE C ODE P ROCESSING
Anonymous authors
Paper under double-blind review
A BSTRACT
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 ﬁx bugs in existing programs.
Being able to process programming language code automatically and provide summaries of code functionality accurately can signiﬁcantly 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 ﬂows and data ﬂows.
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.
I NTRODUCTION
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 ﬁx 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 deﬁned grammars and compilable source code must follow rigid syntactical structures and can be unambiguously parsed into syntax trees.
There can be complex control ﬂows and data ﬂows 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 signiﬁcant 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 classiﬁcation 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 classiﬁcation by signiﬁcant 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 classiﬁcation capsules on the model performance and the effectiveness of the use of additional capsule layers.
R ELATED W ORK
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 classiﬁcation (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 classiﬁcation, 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 ﬁrst to adapt capsule networks for program source code processing to capture both syntactical structures and code dependencies, especially for the problem of program classiﬁcation.
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.
A PPROACH O VERVIEW
The overview of our TreeCaps approach is summarized in Fig.
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 classiﬁcation 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.
A BSTRACT S YNTAX T REE V ECTORIZATION
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, x node ∈ R V , 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.
T REE - BASED C APSULE N ETWORKS
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, X img ∈ R H × W × C , where H, W, C are the ﬁxed height, width and the number of channels respectively, or natural language data, X nlp ∈ R L × E , where L, E are the ﬁxed padded sentence length and the word embedding size respectively.
Hence, the network architecture needs to be constructed to accommodate such tree-structured data, X tree ∈ R T × V , where T, V are the variable tree size (the number of nodes in the tree) and the node embedding size respectively.
5. P RIMARY V ARIABLE T REE C APS L AYER
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 to reach a ﬁxed, pre-deﬁned length, following a similar approach in natural language domain to preserve a ﬁxed 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 ﬁxed 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 ﬁxed matrices W t , W l , W r ∈ R V (cid:48)
× V and a bias term b ∈ R V (cid:48) , where V (cid:48) is the embedding size after the convolution, and the weighting coefﬁcients 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 [ x 1 , ..., x K +1 ] , where x i ∈ R V , then the convolutional output y can be deﬁned as follows,
[Equation]
where η ti , η li , η ri are weights deﬁned corresponding to the depth and the position of the children nodes:
[Equation]
where d i is the depth of the node i in the convolutional windows, p i 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,
Y conv ∈ R T × V (cid:48) .
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 N pvc =
T × V (cid:48) × ε D pvc sets of capsules with outputs u i ∈ R D pvc , i ∈ [1 , N pvc ] , where
D pvc is the dimensions of the capsules (i.e., D pvc is the number of instantiation parameters of the capsules) in the PVC layer.
In order to vectorize each capsule output u j 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,
[Equation]
where || ˆu i || 2 ≤ 1 .
Hence, the output of the primary variable capsule layer is X pvc ∈ R N pvc × D pvc .
5. P RIMARY S TATIC T REE C APS L AYER
The key issue with passing the outputs of the PVC layer, X pvc , to the Code Capsule layer is that the number of capsules, N pvc , 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 deﬁned 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.
5.2. V ARIABLE - TO -S TATIC C APSULE R OUTING
Therefore, we propose a novel variable-to-static capsule routing algorithm, summarized in Algo.
Algorithm
1: procedure R OUTING ( ˆu i , r, a, b ) 2: ˆU sorted ← sort ([ ˆu 1 , ..., ˆu N pvc ]) 3:
Initialize v j : ∀ i, j ≤ a, v j ← ˆU sorted [ i ] 4:
Initialize α ij : ∀ j ∈ [1 , a ] , ∀ i ∈ [1 , b ] , α ij ←5: for r iterations do 6: ∀ j ∈ [1 , a ] , ∀ i ∈ [1 , b ] , f ij ← ˆu i · v j 7: ∀ j ∈ [1 , a ] , ∀ i ∈ [1 , b ] , α ij ← α ij + f ij 8: ∀ i ∈ [1 , b ] , β i ← Softmax ( α i ) 9: ∀ j ∈ [1 , a ] , s j ← (cid:80) i β
ij ˆu i 10: ∀ j ∈ [1 , a ] , v j ← Squash ( s j ) 11: return v j
We initialize the outputs of the Primary Static Capsule layer with the outputs of the a capsules with the highest L 2 norms in the PVC layer.
Hence, the outputs of the PVC layer, [ ˆu 1 , ..., ˆu N pvc ] , are ﬁrst sorted by their L 2 norms, to obtain ˆU sorted , and then the ﬁrst a vectors of ˆU sorted are assigned as v j , j ≤ a .
The intuition is that, in practice, not every node of the AST contributes towards source code classiﬁcation.
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 ( L 2 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 ≤ N pvc .
We assign b = N pvc 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 coefﬁcients as α ij = 0 , equally to all the capsules in the primary variable capsule layer.
Subsequently, as illustrated by Fig 3, they are iteratively reﬁned based on the agreement between the current primary static capsule layer outputs v j and the primary dynamic capsule layer outputs
ˆu i .
The agreement in this case is measured by the dot product, f ij ← ˆu i · v j , and the routing coefﬁcients are adjusted with f ij 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 s j , which is then squashed to update v j .
[Figure]
Code Capsule layer is the ﬁnal layer of the TreeCaps network, which acts as the classiﬁcation capsule layer, as illustrated by Figure 4.
Since the outputs of the PSC layer X psc ∈ R N psc × D psc , where
N psc = a and D psc = D pvc , consist of a ﬁxed 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 v j by the transformation matrices W jm to produce the prediction vectors ˆv m | j = W jm v j .
The trainable transformation matrices learn the part-whole relationships between the primary static capsules and the code capsules, while effectively transforming v j ’s into the same dimensionality as z m , where z m ’s denote the outputs of the code capsule layer.
Similar to the variable-to-static capsule routing, we initialize the routing coefﬁcients γ jm equally, and iteratively reﬁne them based on the agreements between the prediction vectors ˆv m | j and the code capsule outputs z m , where z m = squash ( (cid:80) j γ jm ˆv m | j ) .
1: procedure R OUTING ( ˆv j , t, a, c ) 2:
Initialize ∀ j ∈ [1 , a ] , ∀ m ∈ [1 , c ] , δ jm ←3: for t iterations do 4: ∀ j ∈ [1 , a ] , γ j ← softmax ( δ j ) 5: ∀ m ∈ [1 , c ] , z m ← squash ( (cid:80) j γ jm ˆv m | j ) 6: ∀ j ∈ [1 , a ] , ∀ m ∈ [1 , c ] , δ jm ← δ jm + ˆv m | j · z m 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 ﬁnal capsule outputs X cc ∈ R N cc × D cc , where N cc = κ is the number of classes and D cc is the dimensionality of the code capsules.
Ultimately, we calculate the probability of existence of each class by obtaining L 2 norm of each CodeCaps output vector.
5. M ARGIN L OSS FOR T REE C APS T RAINING
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 deﬁned as follows,
[Equation]
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 .
1 as the lower bound for the correct class and the upper bound for the incorrect class respectively.
E MPIRICAL E VALUATION 6. D ATASETS AND I MPLEMENTATION
We used three datasets in three programming languages to ensure cross-language robustness.
The ﬁrst dataset ( A ) contains 6 classes of sorting algorithms, with 346 training programs on average per class, written in Python.
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, srcML 3 deﬁnes a uniﬁed 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 uniﬁed AST node types deﬁned by srcML.
We used Keras and Tensorﬂow 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 classiﬁcation accuracies, a weighted average ensembling technique (Krogh & Vedelsby
1995) was used.
[Table]
Table 1 compares our results to other approaches for program classiﬁcation.
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 deﬁned speciﬁcally 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 uniﬁed 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
[Table]
we achieved a perfect classiﬁcation result, outperforming TBCNN by a narrow margin of less than 1% .
However, the margin was more signiﬁcant 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 signiﬁcant margin of 17 .
The performance was further improved by 1 . 97% with the use of 3model 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 .
However, TreeCaps surpassed
TBCNN by a more signiﬁcant margin of 8 .
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 .
6. M ODEL A NALYSIS
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
[Table] [Table]
6.3. E FFECT OF THE VARIABLE - TO - STATIC ROUTING ALGORITHM
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, X pvc ∈ R N pvc
×
D pvc consists of a variable component, N pvc .
Using dynamic max pooling across all the N pvc capsules will result in one output capsule, X dmp ∈ R 1 × D pvc .
Since X dmp has no variable components across the training samples, 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 signiﬁcant margin of 8 .
68% , establishing the effectiveness of our proposed algorithm.
6.3. E FFECT OF THE NUMBER OF INSTANTIATION PARAMETERS
The instantiation parameters D cc of the Code Capsule layer acts as the ﬁnal embeddings used for classiﬁcation, 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 classiﬁcation accuracy.
On the contrary, if the dimensionality of the latent representation is too low, it may not be sufﬁcient to capture the variations in source code, leading to under-representation, reducing the classiﬁcation accuracy.
Hence, in an attempt to identify a suitable value for D cc for source code classiﬁcation, we investigate the effect of D cc in the accuracy.
As summarized in Table 2, we observed that the most suitable value was D cc = 8 for the dataset B.
6.3. E FFECT OF THE ADDITION OF A SECONDARY CAPSULE LAYER
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 classiﬁcation accuracy, the added computational complexity increases the inference time by 16% , from 10 .
3 ms to 11 .
9 ms per sample.
The usefulness of the addition of such a SC layer needs to be further investigated
. 6. D ISCUSSION OF L IMITATIONS & F UTURE W ORK
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 classiﬁcation accuracy.
Further, we intend to compare related pieces of code identiﬁed by TreeCaps to program dependencies identiﬁed 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.
C ONCLUSION
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 signiﬁcantly contribute to the high classiﬁcation accuracy of TreeCaps model for program classiﬁcation tasks for various programming languages.
To the best of our knowledge, our work is the ﬁrst 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.