[Paper Review] LambdaNet: Probabilistic Type Inference using Graph Neural Networks

Jiayi Wei, Maruth Goyal+2 more|arXiv (Cornell University)|Apr 29, 2020

Software Engineering Research24 references46 citations

TL;DR

LambdaNet uses a type dependency graph and graph neural networks to infer TypeScript types, including user-defined ones, achieving strong library-type accuracy and enabling open-vocabulary predictions.

ABSTRACT

As gradual typing becomes increasingly popular in languages like Python and TypeScript, there is a growing need to infer type annotations automatically. While type annotations help with tasks like code completion and static error catching, these annotations cannot be fully determined by compilers and are tedious to annotate by hand. This paper proposes a probabilistic type inference scheme for TypeScript based on a graph neural network. Our approach first uses lightweight source code analysis to generate a program abstraction called a type dependency graph, which links type variables with logical constraints as well as name and usage information. Given this program abstraction, we then use a graph neural network to propagate information between related type variables and eventually make type predictions. Our neural architecture can predict both standard types, like number or string, as well as user-defined types that have not been encountered during training. Our experimental results show that our approach outperforms prior work in this space by $14\%$ (absolute) on library types, while having the ability to make type predictions that are out of scope for existing techniques.

Motivation & Objective

• Motivate automatic type annotation inference for gradually-typed languages like TypeScript to reduce manual annotation effort.
• Introduce a novel program abstraction, the type dependency graph, that encodes typing constraints and contextual hints.
• Develop a GNN-based architecture to propagate information over the graph and predict types.
• Enable predictions of user-defined types not seen during training via a pointer-network-like prediction layer.
• Evaluate on real-world TypeScript projects and compare to prior approaches, showing improvements in library-type accuracy and open-vocabulary capability.

Proposed method

• Construct a type dependency graph from static analysis of TypeScript programs, with nodes as type variables and hyperedges encoding logical constraints and contextual hints.
• Compute node embeddings via multi-edge-type graph neural network message passing, handling Fixed, NAry, and NPairs hyperedges.
• Use attention-based aggregation to combine messages and enable long-range inference.
• Employ a pointer-network-like prediction layer that scores candidate types against type-variable embeddings, enabling open vocabulary including user-defined types.
• Represent both library and user-defined types, linking library type embeddings to variable nodes or fixed library-type vectors.
• Train end-to-end using Adam, with hyperparameters tuned on a validation set; predict per variable to ensure consistency across occurrences.

Experimental results

Research questions

• RQ1 Can LambdaNet outperform prior work in inferring library types for TypeScript?
• RQ2 Can the model accurately predict user-defined types not seen during training?
• RQ3 How do the logical and contextual hyperedges contribute to prediction accuracy?
• RQ4 Does a long-range message passing in the GNN improve inference quality?
• RQ5 How does the open-vocabulary, pointer-network-based prediction perform compared with traditional fixed-vocabulary approaches?

Key findings

• LambdaNet achieves 75.6% top-1 accuracy for library types (K=6) vs. 61.5% for DeepTyper (library types only).
• Overall top-1 accuracy including user-defined types is about 64.2%, substantially higher than the TypeScript compiler baseline.
• LambdaNet with open vocabulary predicts user-defined types, outperforming baselines like TypeScript compiler and SimilarName in top-1 and top-5 metrics.
• Ablation shows both contextual and logical edges are crucial, and longer message passing (K up to 6) improves accuracy.
• Attention-based aggregation and NPairs attention significantly boost performance, especially for user-defined types.
• LambdaNet outperforms JSNice on a subset of top-level function annotations.