PaperlayerView exact revision on arXiv

Source-linked AI summary

Code Prediction by Feeding Trees to Transformers

Seohyun Kim, Jinman Zhao, Yuchi Tian, Satish Chandra

arXiv:2003.13848v4cs.SEcs.LG

TL;DR

Code autocomplete still leaves substantial room to improve next-token ranking, especially beyond type-based and existing predictive methods. The paper uses Transformers informed by code structure, and TravTrans significantly outperforms prior models across reported comparisons.

  • Problem

    Existing code-prediction methods leave substantial room for improving next-token ranking, while type-based autocomplete uses naive ranking and struggles to identify type-compatible tokens in dynamic languages.

  • Method

    The paper explores Transformer models that represent partial AST structure through tree paths and traversal order, including PathTrans and TravTrans.

  • Results

    TravTrans improves reciprocal rank from 36.6% to 54.9% versus SeqRNN, from 43.9% to 58.0% versus Deep3, and from 43.6% to 58.0% versus Code2Seq.

  • Takeaways & Limitations

    Transformers outperform existing code-prediction models, with code structural information providing even better predictive power and similar relative benefits on an internal Facebook Python dataset.

  • Takeaways & Limitations

    Evaluations were conducted only on Python, so the reported trends have not been demonstrated to carry over to other programming languages.

Abstract

from arXiv · show

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.

I. INTRODUCTION

Code prediction supplies likely next tokens for IDE autocomplete, where ranking the correct completion highly reduces typing effort and supports code discovery. Existing type-based and sequence-based approaches leave substantial accuracy room, motivating machine-learning methods that use code naturalness and syntax.

  • Code prediction: Code prediction predicts the immediate next token from code surrounding the cursor, commonly presenting likely completions in an IDE.
  • Autocomplete motivation: Top-ranked completions matter because selecting the first suggestion requires less effort, while lower-ranked choices require scrolling.
  • Autocomplete motivation: Autocomplete also supports code discovery, but only when contextually appropriate suggestions appear among the top choices.
  • Prior approaches: Type-based tools rank type-compatible names naively, and dynamic languages make accurate type-compatible candidate lists difficult to construct.
  • Prior approaches: A typical RNN-based method provides less than 37% mean reciprocal rank on py150, leaving substantial room for improvement.The paper reports that its techniques improve mean reciprocal rank by 14% to 18% over previous methods.
  • Illustrative example: In the running example, atoi is predicted at rank 1 only by the new methods proposed in the paper.

D. Feeding Trees to Transformers

The paper adapts Transformers to code by representing partial AST structure through paths or tree traversal orders, then evaluates these models against prior systems on py150. Its best tree-aware model, TravTrans, substantially improves reciprocal-rank accuracy and focuses on relevant tree regions.

  • D. Feeding Trees to Transformers: Transformers are motivated by their ability to handle long-range dependencies better than RNNs, while AST information offers an orthogonal route to improving code prediction.
  • D. Feeding Trees to Transformers: The paper introduces PathTrans, which decomposes an AST into paths, and TravTrans, which represents it through a tree traversal order.A further TravTrans variant incorporates more tree structure.
  • Model comparison: SeqTrans feeds source tokens linearly to a Transformer, whereas the model overview compares increasing AST information with increasingly sophisticated neural architectures.
  • E. Key Results: TravTrans significantly outperforms prior models, improving reciprocal-rank accuracy from 43.9% to 58.0% over Deep3 and from 43.6% to 58.0% over Code2Seq.
  • E. Key Results: 18.3 percentage points separate SeqRNN and TravTrans on the reported comparison, from 36.6% to 54.9%; the relative Transformer benefits are similar on an internal Python dataset.
  • Interpretability: Saliency analysis finds that TravTrans tends to focus on relevant parts of the tree, beginning with the parent node.

F. Contributions

The paper advances code autocomplete by applying Transformers to next-token prediction and incorporating syntactic structure through AST representations. It evaluates these approaches against prior systems and studies how the best model makes predictions.

  • Contributions: The paper introduces Transformer-based approaches for next-token prediction that communicate code’s syntactic structure to the model.The proposed representations include AST paths, tree traversal order, and a variant capturing additional tree structure.
  • Contributions: 14% to 18% gain in accuracy is reported for the best model relative to prior state-of-the-art methods on the py150 dataset.The paper systematically compares its models with effective prior systems applicable to next-token prediction.
  • Contributions: The paper presents a preliminary attribution study indicating that TravTrans conditions predictions on pertinent context tokens.The authors describe this interpretability analysis for autocomplete as, to their knowledge, a first.
  • Contributions: TravTrans provides the paper’s overall conclusion: Transformer-based models over ASTs have the best prediction power for autocomplete.The paper frames this as progress in accurate next-token prediction for modern IDEs.
  • Contributions: The paper defines code prediction as learning a model to approximate the empirical distribution of the next code unit given a partial program context.Training minimizes KL-divergence, equivalently cross-entropy, over code prediction locations.

B. SeqRNN

This section describes prior code-prediction approaches based on RNNs, learned AST decision trees, and Code2Seq path embeddings. It also explains why adapting Code2Seq to next-token prediction is computationally expensive.

  • SeqRNN: SeqRNN feeds source-code tokens into an LSTM, whose hidden state carries information from preceding tokens for next-token prediction.The model uses the final recurrent representation before classification.
  • Deep3: Deep3 combines a learned decision tree with count-based probabilities at its leaves.A TGEN program traverses selected AST paths from a starting node and returns a probability distribution over possible values.
  • Deep3: TGEN programs are learned per corpus by genetic search that selects paths and grows the decision tree using entropy minimization.The experiments use a pretrained model and its Python dataset.
  • Code2Seq: Code2Seq represents code using AST leaf-to-leaf paths, embedding tokenized endpoints and paths before combining them with attention.Endpoint values use LSTMs, path sequences use bidirectional LSTMs, and the resulting embeddings feed a neural network.
  • Code2Seq: Code2Seq must construct new leaf-to-leaf paths at every prediction point, unlike code summarization, which constructs paths once per method.Target-ending paths replace the target value with a placeholder to prevent information leakage.
  • Code2Seq: Treating each prediction point as a separate data point makes Code2Seq computationally very expensive because both training examples and path construction increase.

E. A Primer on Transformers

Transformers process sequences with stacked self-attention blocks that let each prediction use relevant earlier tokens. Causal masking prevents attention to unseen future tokens, while multiple heads capture different relationships.

  • Transformer architecture: A Transformer maps an input context through an embedding layer, repeated attention blocks, and a classification layer.The architecture is summarized as Trans(ctx) = Clsfr(AttnBlk(...AttnBlk(Emb(ctx))...)).
  • Self-attention: For the example sequence [“map”, “(”, “string”, “.”], the input is first converted into token embeddings and then into query, key, and value embeddings.The query, key, and value projections use trainable fully connected networks W_q, W_k, and W_v.
  • Self-attention: Self-attention queries keys and summarizes values according to their relevance to the location being predicted.The attention mechanism is the main component of each attention block.
  • Self-attention: The attention score matrix has size n × n for an input sequence of length n, with rows normalized through softmax.The key-vector dimension is denoted by d_k.
  • Self-attention: Causal self-attention forms a lower-triangular matrix because predictions cannot use tokens that have not yet been seen.The weighted values produce the representation used for the next-token prediction.
  • Self-attention: Multi-headed attention lets the model attend to previous information in multiple ways simultaneously, including across attention blocks.The implementation omits positional encoding.

III. OUR WORK: TRANSFORMER-BASED MODELS

The paper asks whether Transformers can benefit from code syntax and introduces models that communicate AST structure in different ways. SeqTrans applies a Transformer directly to source-token sequences, while PathTrans uses root paths and leaf-token embeddings.

  • The central question is how to communicate syntactic structure to a Transformer while preserving its sequence-processing design.
  • The proposed models include SeqTrans for source tokens, PathTrans for root paths, and TravTrans for richer AST-based representations.
  • SeqTrans: SeqTrans applies a Transformer to partial programs represented as source-token sequences and predicts the next source-code token.
  • SeqTrans: SeqTrans serves as a baseline for later models that incorporate more AST information.
  • SeqTrans: SeqTrans is already a strong direct comparison to the SeqRNN baseline in the experiments.
  • PathTrans: PathTrans combines root-path representations from an LSTM with leaf-token embeddings before feeding them to the Transformer, which predicts leaf tokens.

2) PathTrans:

PathTrans exposes AST structure through root paths while retaining a sequential Transformer input. It summarizes each leaf’s path to the root, combines that representation with the leaf embedding, and predicts leaf tokens.

  • PathTrans: PathTrans represents each leaf using its path to the AST root, recording the internal nodes encountered along its ancestors.
  • PathTrans: An LSTM summarizes each root path, and the resulting representation is added to the corresponding leaf-node embedding before Transformer processing.
  • PathTrans: The summarized root paths are intended to capture local syntactic information that can help prediction.
  • PathTrans: The model predicts only leaf tokens because prediction points are defined as AST leaf nodes and the loss is computed over those nodes.
  • PathTrans: To provide a sequence input, AST nodes are ordered by pre-order traversal, also described as depth-first-search order.
  • PathTrans: The path sequences are capped at 13 tokens, making LSTMs adequate for embedding them in the implementation.

4) Capturing even more AST structure?:

TravTrans+ augments traversal-based Transformer attention with explicit AST path relations. The evaluation uses Python datasets and compares Transformer, RNN, Code2Seq, and Deep3 systems under specified implementation settings.

  • Capturing even more AST structure?: TravTrans+ addresses the loss of detailed structural relationships caused by presenting AST nodes only in a predetermined traversal order.
  • Capturing even more AST structure?: It augments Transformer self-attention with a matrix R encoding the unique path between node pairs as sequences of up and down moves.
  • Capturing even more AST structure?: The path notation uses i and j for the numbers of upward and downward moves needed to travel between two nodes.
  • Evaluation: The experiments use py150 and an internal Facebook Python dataset, with 50k and 5000 evaluation files respectively.
  • Evaluation: Transformer models use a GPT-2-small adaptation with six blocks, six heads per block, a context window of 1000, and embedding dimension 300.
  • Evaluation: The study compares Transformer variants with SeqRNN, Code2Seq, and Deep3, using separately described implementations and training procedures.

C. Evaluation Metric

The study evaluates leaf-token and selected internal-node prediction with mean reciprocal rank, reporting stronger Transformer-based results against recurrent and AST-based baselines. Comparisons between sequence and tree models remain subject to context and prediction-granularity differences.

  • Evaluation Metric: Mean reciprocal rank (MRR) evaluates next-token prediction, with scores reported as percentages and ranks defined for each prediction location.
  • Evaluation Metric: MRR gives partial credit for correct labels ranked among the top predictions, using only ranks up to 10 to approximate autocomplete presentation.
  • Evaluation: 36.6% to 50.1%: on py150, MRR for all leaf tokens increases from SeqRNN to SeqTrans.
  • Evaluation: 23.8% vs 36.5%: on the internal dataset, SeqTrans has higher all-leaf-token MRR than SeqRNN.
  • Evaluation: 14.1%: TravTrans improves Deep3’s MRR from 43.9% to 58.0% on py150 for all leaf tokens.
  • Evaluation: 14.4%: TravTrans improves Code2Seq’s MRR from 43.7% to 58.0% on py150 for all leaf tokens.
  • Evaluation: TravTrans outperforms Deep3 across internal-node prediction types on both py150 and the internal dataset.
  • Evaluation: 58.0% vs 58.8%: TravTrans+ slightly increases all-leaf-node MRR, but mixed leaf-node benefits keep TravTrans as the flagship model.

VI. MODEL INSPECTION

The inspection study examines how TravTrans attributes next-token predictions to preceding code and uses gradient-based saliency to analyze token influence. It finds that predictions tend to emphasize structurally and contextually relevant locations, while often showing low sensitivity to minor input perturbations.

  • Interpretation scope: The study is preliminary rather than a principled complete inspection of TravTrans, but it is intended to illuminate what may support the model’s predictive power.The authors describe it as a first step toward interpretability.
  • Inspection method: Gradient-based saliency maps quantify each input token’s influence because stacked attention layers and multiple heads make direct attention visualization difficult to interpret.The method is intended as a more comprehensive account of token influence.
  • Influence findings: TravTrans generally assigns high importance to the parent AST node and the last token when predicting the next leaf-node value.The parent AST node is the internal node directly above the leaf.
  • Influence findings: TravTrans attributes more influence to internal nodes directly followed by relevant leaf nodes, including prior occurrences of values that inform predictions such as atoi, sys, and argv.The reported examples link predictions to locations preceding string, map, num_requests, and earlier values followed by 0 and 1.
  • Robustness observation: For many prediction locations, very small gradient magnitudes suggest that the model is less sensitive to minor perturbations of the input sequence.The paper presents this as robustness in an orthogonal observation.

VII. THREATS TO VALIDITY

The validity discussion bounds the evidence by highlighting unsupported OOV predictions, Python-only evaluation, restricted model comparisons, and dataset or evaluation-setting differences. It also situates the work among related Transformer and code-prediction studies while noting that production user evaluation is outside scope.

  • OOV handling: 21% of py150 and 35% of internal-test prediction locations contain OOV tokens, which the main evaluation models cannot correctly predict.The models do not handle OOV predictions, so such locations count as incorrect.
  • OOV handling: OOV handling is outside the paper’s core model-design scope, although adding it should improve performance and is not expected to reverse the reported findings.A quick py150 comparison found TravTrans ahead of PointerMixture even without OOV handling.
  • Python specificity: The evaluation uses Python only, so the paper does not demonstrate that its trends carry over to other programming languages.The cited Deep3 result provides related evidence for Python-to-JavaScript transfer, not a demonstration by this paper.
  • Model comparison: The comparison excludes techniques requiring compiler support beyond AST construction, def-use information, and graph neural networks, and focuses on Transformer architectures.These exclusions bound the scope of the model comparison.
  • Model comparison: PointerMixture outperforms the principal baselines on py150, while TravTrans outperforms PointerMixture there despite counting all OOV predictions as wrong.On the internal dataset, PointerMixture performs better for some prediction types because TravTrans was trained on py150 rather than the internal data.
  • Practical scope: The paper does not include a user evaluation of an autocomplete tool, and prior work cited in the discussion reports that synthetic-benchmark accuracy can drop on real-world data.Lightweight models for faster IDE computation are also identified as a separate practical concern.
  • Related work: The paper reports systematic Transformer evaluation for code prediction alongside practical interest in Transformer-based autocomplete and prior work extending Transformers beyond sequential structures.Related work also reports structural representations helping code correction and translation.

IX. CONCLUSION AND FUTURE WORK

The paper concludes that Transformers improve code prediction and that communicating code structure improves predictive power further. It also reports that TravTrans focuses on relevant code locations, while identifying OOV handling and multi-token completion as future directions.

  • Conclusion: Transformers outperform existing code-prediction models, and structural information improves their predictive power further.The conclusion summarizes the paper’s central performance finding.
  • Conclusion: TravTrans’s attribution study indicates that its best-performing model tends to focus on relevant code locations for prediction.The study is presented as preliminary model inspection.
  • Future work: Improving OOV handling is a future direction because test tokens may be absent from training data, especially arbitrary method and variable identifiers.The paper mentions copying mechanisms as a possible mitigation.
  • Future work: The paper also plans to explore predicting multiple tokens at once, including autocomplete of entire expressions.The current work focuses on next-token prediction.
Loading 2003.13848v4…