Advanced Graph-Based Deep Learning for Probabilistic Type Inference

The goal of type inference is to automatically deduce the type of an expression in a program. The deduced type can be either partially or fully, depends on the simplicity of the target programming language or the robustness of the type inference analysis.

Static type inference algorithms infer type information via solving the type constraints collected from the target program. A type constraint represents the aggregation and reduction of the type information between expressions in the program. This allows type inference to be formulated as a graph analysis problem.

Graph neural networks (GNNs) are machine neural network models that directly operate on graph-structured data. A graph can be defined as $G=(V,E)$, where $V$ is the set of nodes and $E$ is the set of edges. The neighborhood of a node $v$ is defined as $\mathcal{N}(v)=\{u\in V\leavevmode\nobreak\ |\leavevmode\nobreak\ (u,v)\in E\}$. A node $v$ can have attributes represented as a feature vector $\mathbf{x}_{v}$, and similarly, an edge $(u,v)$ can have attributes represented as a feature vector $\mathbf{e}_{uv}$. GNNs learn the representations of graph structures by propagating the states of nodes to their neighbors iteratively and transforming the nodes states into output representations. Let $\mathbf{h}_{v}$ be the state vector of node $v$ and $K$ be the number of propagation steps,²²2In this paper, we use the term propagation step to represent the computation in an iteration. A propagation step is also called a GNN  layer in the literature. a general form of a GNN can be defined as:³³3Although there exists a more general form that allows edge features to be updated during the propagation [6], the one we present here is sufficient to formulate all the GNN architectures we will discuss in this paper.

The state of each node $\mathbf{h}_{v}$ is initialized with its feature vector $\mathbf{x}_{v}$ in Equation 1. At step $k$, a node receives messages from its neighbors through corresponding edges. A message vector $\mathbf{m}_{uv}$ is generated from a neighbor’s hidden state $\mathbf{h}_{u}^{(k-1)}$ at step $(k-1)$ and the feature vector of the edge where the message is passed through (i.e., $\mathbf{e}_{uv}$), using the message function $f_{\mathrm{message}}^{(k)}$ in Equation 2. The messages are then aggregated to a single vector using the aggregation function $f_{\mathrm{aggregate}}^{(k)}$ in Equation 3. At the end of step $k$, the node’s state $\mathbf{h}_{v}^{(k)}$ is updated from its state at the previous step $\mathbf{h}_{v}^{(k-1)}$ and the aggregated message vector $\mathbf{a}_{v}^{(k)}$ using the update function $f_{\mathrm{update}}^{(k)}$ in Equation 4.

The functions $f_{\mathrm{message}}^{(k)}$, $f_{\mathrm{aggregate}}^{(k)}$, $f_{\mathrm{update}}^{(k)}$ are neural networks containing trainable parameters. Those functions can be the same and share parameters across different steps (i.e., $f_{\_}^{(1)}=f_{\_}^{(2)}=\ldots=f_{\_}^{(K)}$) or they can be different functions and use a separate set of parameters for each step. Following the taxonomy in [26], we use the term recurrent graph neural networks to refer to GNNs with shared functions and parameters for all steps, and convolutional graph neural networks to refer to GNNs with separate functions and parameters for different steps.

There has been a large body of work that leverages machine learning to assist static analysis. By their tasks, those approaches can be categorized into the following categories:

Program Property Prediction. In this class of tasks, the goal is to assign property labels to (part of) a program. The typical approach is to apply supervised learning to build a model that predicts the property labels from a program’s representation. The program representations that have been adopted by previous work include token sequences [11], abstract syntax trees [4, 21], and some more complex graphs [2, 23].

Learning to Parametric Program Analysis. To make a trade-off between precision and cost, parametric analyses have been applied in real-world solutions [28, 18]. Unlike manual parameterization and heuristic parameter search, which rely on the expertise of users or heuristic developers, statistical learning approaches can learn from automatically generated data and have shown a good capacity for finding the optimal configuration for parameters including the granularity of program abstraction [15, 19].

Program Specification Identification. This scenario focuses on mining program specification from code corpora, including the API specifications [17, 8], the behaviors of a library [5] and code idioms [3]. These techniques can be applied for detecting API misuses, performing code completion, etc.

In this paper, we focus on type inference, which can be categorized as a program property prediction problem. We investigate the approach of using graph structure to present the input program and building machine learning models to predict program properties. A  GNN-based learning framework is designed to learn from graph-structured data extracted from code and predict types for program elements.

We formulate type inference as a graph node label prediction problem. The prediction is at the file level for better scalability and flexibility. Our GNN-based type inference framework is shown in Figure 1. For each source file, we construct a graph whose nodes correspond to program elements in the file. The graph edges include syntax edges extracted from the abstract syntax tree and over-approximated data flow edges derived from name-matching-based static analysis. A GNN model is trained to predict type for the graph nodes from a fixed type vocabulary. The target programming language we choose is JavaScript, and the unit of program element for prediction is the identifier, which covers variables, function parameters, and object properties in JavaScript. Following the methodology of previous work [11], we train the model on a dataset of TypeScript programs. The type labels are obtained from the TypeScript compiler, and the graphs are constructed from the code without using any information from existing type annotations.

A data-driven approach relies on high-quality data containing features that are relevant to the target problem. For the type inference problem, the relevant features include the source of type information (e.g., literals), the code patterns that have implications about types, and the paths of type propagation. Ideally, the input data to the learning model should at least cover those features. In this work, we defined a type flow graph (TFG) that carries those features mentioned above by encoding syntactic and approximate semantic information in a program.

A TFG is a directed graph whose nodes represent program elements in the code and edges encode type-relevant relationships between the nodes. It is constructed from a program’s abstract syntax tree (AST). We give the definitions of the nodes and edges and how they can be extracted in the following paragraphs.

Here we introduce our machine learning model design for static type inference. This learning model is presented as a customizable system composed of a core GNN model with different building blocks and several other components that converts node/edge features into vector representations.

As discussed in Section 2.2, our GNN-based type inference system allows multiple design choices across the different components. We studied how they could affect the accuracy of type prediction by evaluating various combinations of those design choices on our dataset.

For our approach to be applicable to real-world interactive programming environments, it is important for our techniques to be efficient in practice. To evaluate the efficiency of the models in our framework, We measured their sizes (i.e., the number of model parameters) and inference throughputs. The results are shown in Figure 8. Here we only compare the models that take in edge features.

The number of parameters in a model affects its applicability in several ways. A larger model costs more memory to be loaded and can sometimes bring more computation workload. The blue bars in Figure 8(a) shows the number of parameters of our GNN models. C-GNN contains the largest number of parameters due to its non-recurrent design that requires a separate set of parameters for each step. The size of R-GAT is slightly larger than R-GNN because of its attention mechanism that requires extra parameters. Among the R-GNN variants, the two with name segmentation are significantly smaller than the others thanks to their smaller name segment embedding size (32-dimensional), although the contextual layer increases the model size by bringing in more parameters.

Figure 8(b) compares the computation efficiency of the models by showing their inference throughputs. This metric is computed as the number of source files in the test set divided by the total runtime of performing inference on those files. The batch size used for inference is set to 64 graphs for all models. We took six measurements for each model and reported their mean and standard deviation. C-GNN and R-GNN are the two simplest models in terms of the architecture, requiring less computation during inference, and thus produced the highest throughputs. The attention mechanism in R-GAT adds additional computation overhead, causing it to produce a lower throughput than R-GNN. Similarly, name segmentation and the contextual layer both make the models containing them to be slower. It is worth noting that the addition of the contextual layer brings more computation overhead. This is because it is based on a recurrent neural network, whose workload is determined by the length of the input code token sequence.

We compared our approach with two state-of-the-art deep learning type inference approaches: DeepTyper [11] and LambdaNet [25].

DeepTyper treats a source file as a sequence of tokens and runs a two-layer bidirectional RNN on it to predict types for identifier tokens. To improve the consistency of type prediction for the same variable, the model also includes a consistency layer between the two RNN layers that takes the mean of hidden states for identifiers with the same name and adds it back to those hidden states. We implemented DeepTyper in our evaluation framework and used the same hyperparameters as suggested in their paper and open-source implementation (300-dimensional embeddings and 650-dimensional hidden layers).⁶⁶6https://github.com/DeepTyper/DeepTyper

LambdaNet is a GNN-based neural type inference model. It extracts a project-level type dependency graph (TDG) through static analysis and uses a convolutional GNN to perform type inference for the graph nodes, each of which is associated with a unique variable or an expression. LambdaNet’s TDG includes a manually defined set of edge types and labels. In contrast, our TFG automatically derives hundreds of edge labels from AST node types and parent-child relationships, which provide richer information to the GNN model and is also more easily adaptive to other languages. We also implemented LambdaNet in our evaluation framework. Unlike the original LambdaNet, our implementation gives a prediction for each occurrence of a variable instead of having one prediction per declaration, and uses a prediction space consisting of types in a fixed vocabulary. Those changes allow us to compare the efficacy of LambdaNet’s TDG and our TFG without the potential interference brought by different prediction schemes. We didn’t change the neural network architecture and used the hyperparameters suggested in their paper and open-source implementation (32-dimensional embeddings and hidden layers, six steps of graph propagation).⁷⁷7https://github.com/MrVPlusOne/LambdaNet

The reason why we chose to reimplement the two state-of-the-art approaches in our own framework, instead of using their released code, is to obtain an apples-to-apples comparison. We studied both DeepTyper’s and LambdaNet’s code repositories and had the following observations:

• •

  Their training code is tightly coupled with their data processing code, making it hard to reuse their code with a different data preprocessing approach such as that in our approach.

• •

  Their implementations are based on different deep learning frameworks which could result in efficiency variations for the same model, bringing challenges to our efficiency comparison.

• •

  LambdaNet’s authors implemented a variant of DeepTyper in their framework for comparison. However, our preliminary experiments showed that some changes introduced by the variant had a negative impact on accuracy on our dataset.

Based on those observations, we decided to use our own implementations of the two approaches in the experiments.

Part 4 of Table 2 shows the accuracy DeepTyper and LambdaNet achieved in our experiments. DeepTyper delivered a similar level of accuracy compared with R-GNN, our baseline GNN model. However, it is outperformed by the three variants of R-GNN that integrates name segmentation and/or the contextualized layer in terms of top-1 accuracy metrics. LambdaNet failed to reach the same level of accuracy as R-GNN. In terms of top-1 accuracy for predicting all 100 types, our best model, R-GNN${}_{\mbox{\footnotesize{\it NS-CTX}}}$, outperformed DeepTyper by 3.14% (absolute) and LambdaNet by 8.31%. Note that since we ported both prior approaches into our experimental framework, which contains different data processing procedures, the accuracy numbers we obtained could not be directly compared with the ones reported in their papers. This is especially the case for LambdaNet due to all the changes we made to it in our implementation.

The orange bars in Figure 8 shows the sizes and inference throughputs of DeepTyper and LambdaNet. As can be seen from Figure 8(a), DeepTyper contains significantly more parameters than other models. This is due to its large hidden state size. In contrast, LambdaNet employs a relatively small hidden state size, resulting in a more compact model. The sizes of our GNN models fall in between and are closer to LambdaNet’s. Shown in Figure 8(b), DeepTyper yielded the lowest throughput among all the models being compared. This is because DeepTyper runs multiple layers of RNN through the code token sequence, which can be very long. Our models with the contextual layer (R-GNN${}_{\mbox{\footnotesize{\it CTX}}}$ and R-GNN${}_{\mbox{\footnotesize{\it NS-CTX}}}$) suffered from the same situation, but still gave higher throughputs than DeepTyper, because that the contextual layer used only a single layer of RNN. The other four GNN models without the contextual layer achieved higher throughputs than LambdaNet. The reason comes from that LambdaNet generally contain more edges than our TFGs, thus requiring more computation resource during message passing and aggregation.