PaperlayerView exact revision on arXiv

Source-linked AI summary

Multi-task Learning based Pre-trained Language Model for Code Completion

Fang Liu, Ge Li, Yunfei Zhao, Zhi Jin

arXiv:2012.14631v1cs.SE

TL;DR

Existing code-completion models lose context-dependent token meaning and often ignore identifier types, making identifier completion difficult. CugLM uses a Transformer with jointly trained code-understanding and code-generation objectives, then predicts token types to assist token completion. On Java and TypeScript, it outperforms prior state-of-the-art methods, especially for identifiers, although generalization to other languages remains unvalidated.

  • Problem

    Existing models use static token representations and generally ignore identifier type information, while identifier completion remains challenging.

  • Method

    CugLM pre-trains a Transformer with masked bidirectional, next-code-segment, and unidirectional objectives, then jointly predicts completion tokens and their types.

  • Results

    CugLM achieves better token-completion results than previous state-of-the-art models, especially for identifiers, on Java and TypeScript datasets.

  • Takeaways & Limitations

    The results support using pre-trained contextual representations and type information for source-code completion.

  • Takeaways & Limitations

    The findings were evaluated on Java and TypeScript GitHub datasets, so further studies are needed to validate generalization to other programming languages.

Abstract

from arXiv · show

Code completion is one of the most useful features in the Integrated Development Environments (IDEs), which can accelerate software development by suggesting the next probable token based on the contextual code in real-time. Recent studies have shown that statistical language modeling techniques can improve the performance of code completion tools through learning from large-scale software repositories. However, these models suffer from two major drawbacks: a) Existing research uses static embeddings, which map a word to the same vector regardless of its context. The differences in the meaning of a token in varying contexts are lost when each token is associated with a single representation; b) Existing language model based code completion models perform poor on completing identifiers, and the type information of the identifiers is ignored in most of these models. To address these challenges, in this paper, we develop a multi-task learning based pre-trained language model for code understanding and code generation with a Transformer-based neural architecture. We pre-train it with hybrid objective functions that incorporate both code understanding and code generation tasks. Then we fine-tune the pre-trained model on code completion. During the completion, our model does not directly predict the next token. Instead, we adopt multi-task learning to predict the token and its type jointly and utilize the predicted type to assist the token prediction. Experiments results on two real-world datasets demonstrate the effectiveness of our model when compared with state-of-the-art methods.

1 INTRODUCTION

The introduction identifies contextual representation and identifier completion as key weaknesses in existing code-completion models, then presents CugLM as a Transformer-based, multi-task solution evaluated on Java and TypeScript.

  • Motivation: Existing code-completion models do not adequately use contextual information when representing source-code tokens.Static embeddings assign the same vector to a token across contexts, losing context-dependent meaning.
  • Motivation: Identifier completion remains difficult because most models ignore identifier types and often target easier punctuation-like tokens.More than two-thirds of completions concern non-identifiers, while IDEs use type information to suggest suitable names.
  • Approach: CugLM pre-trains contextual token representations by masking identifiers and predicting them from bidirectional code context.The identifier-focused masked language-modeling objective is designed to support source-code modeling and completion.
  • Approach: CugLM jointly pre-trains masked bidirectional, next-code-segment, and unidirectional language-modeling objectives before code-completion fine-tuning.The unidirectional objective matches generation settings where only leftward context is available.
  • Approach: During completion, CugLM predicts a token and its type jointly, using the predicted type to assist token prediction.The model first predicts the token type and then uses that prediction to guide the token prediction.
  • Results: 80% and 81% accuracy on Java and TypeScript for all-token completion improved Pointer Mixture Network by 17% and 24%, respectively.For identifier completion, accuracy reached 48% and 39%, with relative improvements over Pointer Mixture Network of 29% and 34%.
  • Contributions: The paper reports the first Transformer-based pre-training attempt for code completion and the best performance against its evaluated baselines on two real-world datasets.The evaluation covers Java and TypeScript programs collected from GitHub.

2 BACKGROUND

The background traces code modeling from statistical and recurrent language models to attention-based Transformers, while motivating pre-training and multi-task learning as ways to improve representations and generalization.

  • Statistical Language Models: Statistical language models assign probabilities to token sequences and exploit predictable patterns in programming languages.They estimate how natural or likely a sequence is, but long contexts make probabilities difficult to estimate directly.
  • Statistical Language Models: N-gram models address long-context estimation with a Markov assumption that limits token dependence to the previous n−1 tokens.These models have been widely applied to code completion and capture repetitive source-code regularities.
  • Neural Language Models: LSTM and GRU language models use recurrent connections and gates to capture longer dependencies than N-gram models.Their gating mechanisms loosen the fixed context-size limitation of N-gram models.
  • Attention-based Models: Transformers replace recurrent layers with multiheaded self-attention, reducing sequential computation and capturing longer-range dependencies.Transformer-XL further introduces recurrence to model very long-term dependencies.
  • Multi-task Learning: Multi-task learning transfers knowledge across related tasks by sharing hidden layers and leveraging task-specific training signals.The shared representations can capture common features and reduce over-fitting risk.
  • Pre-trained Language Models: Language-model pre-training learns broad representations, improves downstream initialization, and can regularize models trained on small datasets.These benefits motivate adapting pre-trained models to downstream software-engineering tasks.

3 CugLM

CugLM uses a Transformer with task-specific attention masks and shared multi-task pre-training, then fine-tunes through type-aware token completion. Its objectives combine bidirectional code understanding with unidirectional generation.

  • 3.1 Model Architecture: CugLM encodes program token sequences with an L-layer Transformer to produce contextual representations at multiple levels.The model also produces an aggregated [CLS] representation for classification tasks.
  • 3.1 Model Architecture: Different attention masks support bidirectional modeling for understanding and left-to-right modeling for generation.In unidirectional modeling, each token attends only to leftward context and itself.
  • 3.2 Input Representation: The input representation sums token, segment, and position embeddings, with special [CLS] and [SEP] tokens supporting sequence aggregation and segment separation.Learned positional embeddings support sequence lengths up to 128 tokens.
  • 3.3 Pre-training Procedure: Pre-training jointly optimizes masked bidirectional LM, Next Code segment Prediction, and Unidirectional LM with shared parameters.The MLM objective masks identifiers with type information, while ULM supports language generation tasks such as code completion.
  • 3.4 Fine-tuning Procedure: Fine-tuning predicts a masked token’s type first, then concatenates its predicted type embedding with the mask representation to predict the token.The fine-tuning objectives include UMLM and ULM, and their cross-entropy losses are summed.
  • 3.4 Fine-tuning Procedure: Jointly learning the fine-tuning objectives is intended to improve predictions for identifiers and other tokens.The experiments use Java and TypeScript programs from public GitHub repositories, with separate pre-training and fine-tuning project splits.

4 EXPERIMENTS AND ANALYSIS

Experiments evaluate CugLM on Java and TypeScript code-completion datasets against established baselines and through ablations of its training objectives. The model outperforms LSTM-based systems, while pre-training, fine-tuning tasks, and type prediction each contribute to performance.

  • Data preparation: CugLM is pre-trained and fine-tuned on Java and TypeScript programs collected from duplicate-filtered, publicly available GitHub repositories.Projects are split into 60% for pre-training and 40% for fine-tuning.
  • Experimental setup: The evaluation uses top-1 accuracy, measuring whether the correct suggestion appears first in the model’s ordered list.The vocabulary contains the 50,000 most frequent training tokens.
  • Comparison with baselines: 84.06% token accuracy and 55.19% identifier accuracy are achieved on the Java large test set, with relative improvements over Pointer Mixture Network of 23.07% and 43.69%.On TypeScript, the corresponding accuracies are 82.14% and 41.85%, with relative improvements of 19.47% and 23.96%.
  • Comparison with baselines: CugLM substantially outperforms Transformer-XL trained from scratch and recurrent baselines, indicating that pre-training and fine-tuning are crucial to its success.Transformer-XL also outperforms the recurrent baselines on both datasets.
  • Research questions and results: Removing any pre-training task reduces performance, with MLM contributing most to the improvements, especially for identifier completion.Removing NCP performs worse than the full model but better than removing ULM.
  • Research questions and results: Removing UMLM or ULM during fine-tuning hurts performance on both datasets, particularly identifier prediction, while predicted identifier types reach 68.89% accuracy on Java and 79.31% on TypeScript.The fine-tuning tasks improve contextual representations and use identifier type information.

5 DISCUSSION

The discussion examines token-type completion, model size, subword encoding, and threats to validity. CugLM captures syntactic and semantic information effectively, but results vary by token type and language, with generalizability still limited.

  • Token Completion Analysis: Punctuation dominates completions, and BPE NLM slightly outperforms CugLM on these easier tokens.Punctuation completions are less beneficial to developers than identifier completions.
  • Token Completion Analysis: CugLM substantially outperforms BPE NLM on numeral and operator completions, which are more related to program semantics.The discussion attributes this performance to representations that better capture syntactic and semantic information.
  • Model Complexity: CugLM has fewer trainable parameters than all baselines despite multi-task learning because tasks share one multi-layer Transformer network.The model design also targets training efficiency and over-fitting avoidance.
  • Subword Encoding: BPE vocabulary segmentation produces comparable Java performance but slightly decreases accuracy on TypeScript.The paper analyzes this difference in relation to masked identifier proportions and type annotations.
  • Threats to Validity: Further studies are needed to validate and generalize the findings beyond Java and TypeScript.The datasets consist of GitHub repositories, and the language choice reflects available identifier type information.

6 RELATED WORK

Related work progresses from heuristic and statistical code completion to recurrent and Transformer-based models, including approaches using ASTs, subwords, caches, and multimodal pre-training. The paper positions CugLM as a Transformer-based pre-trained model tailored specifically to code completion.

  • Statistical Code Completion: Early code completion relied mainly on heuristic rules and static type information, followed by statistical language models exploiting predictable source-code patterns and localness.N-gram models were widely used, while cache mechanisms were introduced to capture source-code localness.
  • Neural Code Completion: Deep recurrent models advanced source-code modeling, while pointer mixtures addressed out-of-vocabulary tokens and other systems targeted AST-level or structurally informed completion.The related work also includes LSTM-based method-call recommendation and Transformer models incorporating syntactic structure.
  • Pre-trained Language Models: UNILM showed that configurable self-attention masks can support both language understanding and generation, overcoming BERT’s difficulty with generation.This provides a pre-training precedent for combining understanding and generation capabilities.
  • Multimodal Pre-training: Multimodal pre-trained models learn alignments across modalities such as language-image, language-video, and language-code pairs.CodeBERT is cited as a bimodal model connecting natural language and programming language.
  • Positioning of CugLM: CugLM is presented as the first Transformer-based pre-trained language model tailored to code completion.Its design is inspired by prior pre-training work but targets source-code completion directly.

7 CONCLUSIONS AND FUTURE WORK

The paper proposes a multi-task, Transformer-based pre-trained language model for source-code modeling and reports better token-completion results than previous state-of-the-art models, especially for identifiers. Future work extends the model to other languages and tasks.

  • Conclusions: CugLM uses a Transformer-based architecture, three pre-training objectives, and fine-tuning for code completion.The paper frames this as a multi-task learning approach for code understanding and generation.
  • Conclusions: Experimental results show better completion than previous state-of-the-art models, especially for identifiers.The conclusion presents the model as the first application of pre-trained language modeling to code completion.
  • Future Work: Future work will apply the model to other programming languages and fine-tune it for additional tasks.
Loading 2012.14631v1…