[Paper Review] Code Prediction by Feeding Trees to Transformers

Seohyun Kim, Jinman Zhao+2 more|arXiv (Cornell University)|Mar 30, 2020

Software Engineering Research66 references32 citations

TL;DR

This paper enhances next-token code prediction using Transformer models that leverage AST structure, introducing SeqTrans, PathTrans, TravTrans, and TravTrans+ to achieve significant accuracy gains over prior methods on Python code datasets. TravTrans delivers the strongest results, with up to ~58% mean reciprocal rank (MRR) on py150, surpassing prior state-of-the-art by 14–18 percentage points.

ABSTRACT

We advance the state-of-the-art in the accuracy of code prediction (next token prediction) used in autocomplete systems. First, we report that using the recently proposed Transformer architecture even out-of-the-box outperforms previous neural and non-neural systems for code prediction. We then show that by making the Transformer architecture aware of the syntactic structure of code, we further increase the margin by which a Transformer-based system outperforms previous systems. With this, it outperforms the accuracy of an RNN-based system (similar to Hellendoorn et al. 2018) by 18.3%, the Deep3 system (Raychev et al 2016) by 14.1%, and an adaptation of Code2Seq (Alon et al., 2018) for code prediction by 14.4%. We present in the paper several ways of communicating the code structure to the Transformer, which is fundamentally built for processing sequence data. We provide a comprehensive experimental evaluation of our proposal, along with alternative design choices, on a standard Python dataset, as well as on a Facebook internal Python corpus. Our code and data preparation pipeline will be available in open source.

Motivation & Objective

• Motivate and improve next-token code prediction for autocomplete in IDEs.
• Assess how integrating AST structure into Transformer models impacts predictive accuracy.
• Compare Transformer-based, tree-aware models against prior non-Transformer and AST-based approaches.
• Evaluate on standard and in-house Python datasets to demonstrate generalization.

Proposed method

• Represent partial programs as sequences of tokens (SeqTrans) or as AST-based structures (PathTrans, TravTrans) to feed into a Transformer.
• Introduce PathTrans: root-path representations from leaf nodes to root embedded and combined with leaf embeddings for Transformer input.
• Introduce TravTrans: feed AST DFS-order tokens to Transformer to incorporate structural information.
• Introduce TravTrans+: augment TravTrans with explicit tree-relational paths in attention via a tree-relational matrix R in AttnTreeRel.
• Compare against strong baselines (SeqRNN, Deep3, Code2Seq) on py150 and Facebook internal Python datasets.
• Use mean reciprocal rank (MRR) as the evaluation metric for next-token prediction on leaf tokens.

Experimental results

Research questions

• RQ1 Does a Transformer outperform an RNN-based model for next-token Python code prediction?
• RQ2 Do Transformer models that leverage AST structure outperform previous AST-based approaches for code prediction?
• RQ3 How do SeqTrans, PathTrans, TravTrans, and TravTrans+ compare in accuracy across datasets?
• RQ4 Is the tree-relational enhancement (TravTrans+) beneficial beyond TravTrans?
• RQ5 What insights can saliency maps provide about how TravTrans attends to tree structure during prediction?

Key findings

• TravTrans achieves the best overall performance, with 58.0% MRR on py150 and substantially higher accuracy than prior methods.
• SeqTrans (baseline Transformer) reaches 54.9–50.1% MRR depending on dataset and task, outperforming SeqRNN (36.6–23.8% MRR).
• PathTrans and TravTrans improve over non-Transformer AST baselines (Deep3, Code2Seq) with TravTrans achieving a 14.1–14.4 percentage-point gain in MRR over Deep3 and Code2Seq on py150.
• On the internal Facebook Python dataset, Transformer-based models show similar relative advantages over baselines.
• Saliency analysis indicates TravTrans attends to the most relevant tree regions, especially focusing on parent nodes when predicting tokens.