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

1 INTRODUCTION

• Static embeddings in prior code-completion research [3] [25] [27] assign a token the same representation across contexts, losing distinctions such as overloaded method meanings.The introduction notes that surrounding tokens can provide information about a program entity’s role and meaning.Source: §1 · p. 2
• Existing language-model-based methods perform poorly on identifier completion, often target easier punctuation-like tokens, and commonly ignore identifier types that IDEs use to guide suggestions.The passage states that more than two-thirds of completions do not refer to identifiers.Source: §1 · p. 2
• CugLM jointly pre-trains contextual program representations with identifier-masked bidirectional modeling, next-code-segment prediction, and left-to-right language modeling.The left-to-right objective is included because generation tasks such as code completion use only leftward context.Source: §1 · pp. 2–3
• For code completion, CugLM predicts a token’s type first and uses the predicted type to assist token prediction.The model is fine-tuned on the code-completion task after pre-training.Source: §1 · p. 3
• On Java and TypeScript, CugLM achieves 80% and 81% accuracy for all-token completion, and 48% and 39% for identifier completion, respectively.For all-token completion, relative gains on Java and TypeScript are 17% and 24% over Pointer Mixture Network [27], and 19% and 24% over BPE NLM [25]. For identifier completion, the respective gains are 29% and 34%, and 11% and 9%.Source: §1 · p. 3
• The study pre-trains and fine-tunes CugLM on massive GitHub Java and TypeScript corpora and compares it with BPE NLM [25] and Pointer Mixture Network [27].The paper presents the approach as a Transformer-based pre-training effort for code completion that uses type information to improve identifier suggestions.Source: §1 · p. 3

2.1 Statistical Language Model

• N-gram models estimate token probabilities from the n−1 most recent tokens and have been applied to code completion, capturing repetitive source-code patterns.Programming languages have predictable statistical properties that can be modeled with language models.Source: §2.1 · p. 3
• LSTMs and GRUs capture longer dependencies than N-grams, while attention-based Transformers reduce sequential computation and Transformer-XL extends long-term context [6] [8] [48].Attention adds direct connections between distant tokens to help address gradient vanishing and explosion in gated recurrent networks [2] [48].Source: §2.1 · p. 3

2.2 Multi-task Learning

• Multi-task learning transfers knowledge across related tasks by sharing hidden layers, improving generalization and reducing over-fitting.Related-task training signals provide domain-specific information, and shared representations capture features common to the tasks.Source: §2.2 · p. 3

2.3 Pre-trained Language Models

• Pre-training can provide universal representations, improved initialization, and regularization; this work adapts a pre-trained code model to completion after addressing BERT’s generation limitation with UNILM [7] [10] [22] [39].Feature-based methods add pre-trained representations as features, whereas fine-tuning adapts pre-trained parameters directly; UNILM supports both understanding and generation tasks [10] [39].Source: §2.3 · p. 3
• CugLM shares parameters across MLM, NCP, and ULM pre-training objectives and uses different self-attention masks to control token access to context.The figure presents the overview of CugLM pre-training.Source: Figure 2 · p. 4

3.1 Model Architecture

• An L-layer Transformer produces contextual token representations, using bidirectional or left-to-right attention masks; its [CLS] representation aggregates the sequence for classification.Bidirectional masking lets tokens attend to each other, while unidirectional masking restricts attention to each token’s leftward context and itself.Source: §3.1 · p. 4

3.2 Input Representation

• Inputs pack two code segments and sum token, segment, and learned position embeddings, with positional embeddings supporting sequences up to 128 tokens.Distinct learned segment embeddings identify the first and second code segments.Source: §3.2 · pp. 4–5, Figure 3 · p. 5

3.3 Pre-training Procedure

• Pre-training jointly optimizes bidirectional masked-identifier prediction, next-code-segment classification, and left-to-right next-token prediction with shared parameters.For next-code-segment prediction, the second segment is the actual successor half the time and a random corpus segment half the time.Source: §3.3 · p. 5

3.4 Fine-tuning Procedure

• Fine-tuning uses left-context-only UMLM for typed identifiers and predicts type first, then conditions token prediction on that predicted type, alongside left-to-right ULM.The predicted type embedding is combined with the masked token’s hidden representation to predict the token.Source: §3.4 · pp. 5–6

4.1 Data preparation

• The corpus uses open-source GitHub projects split 60% for pre-training and 40% for fine-tuning, whose projects are divided 8:1:1 into train, validation, and test sets.For BPE NLM evaluation, the authors randomly sample 200 test files per language because its completion procedure is time-consuming.Source: §4.1 · p. 6

4.3 Research Questions and Results

• CugLM outperforms LSTM-based baselines on Java and TypeScript, especially for identifier completion.On Java, it achieves 84.06% token and 55.19% identifier accuracy, with relative improvements over Pointer Mixture Network of 23.07% and 43.69%, respectively; on TypeScript, it achieves 82.14% and 41.85%, with relative improvements of 19.47% and 23.96%.Source: Table 2 · p. 7, §4.3 · p. 8

4.2 Experimental Setup

• CugLM uses a six-layer, 516-dimensional Transformer and is evaluated by top-1 completion accuracy.It is pre-trained for 600,000 steps and fine-tuned for 300,000 steps; the token and type vocabularies each contain 50,000 entries, with out-of-vocabulary targets counted as incorrect.Source: §4.2 · p. 7

4.3 Research Questions and Results

• Removing any pre-training task hurts performance; MLM contributes most, ULM next, and NCP contributes less than ULM.MLM removal particularly harms identifier completion; ULM supports representations suited to both bidirectional pre-training and unidirectional fine-tuning, while NCP helps model relationships between code segments.Source: Table 3 · p. 8, §4.3 · p. 9
• Removing either fine-tuning task reduces performance, with UMLM especially important for identifier prediction.UMLM provides contextual representations and uses identifier type information, while ULM supports identifier prediction by helping capture code-segment semantics.Source: Table 3 · p. 8, §4.3 · p. 9
• CugLM exceeds BPE NLM on identifier and all-token completion, and substantially exceeds Transformer-XL trained from scratch.BPE NLM is the strongest baseline for identifier completion; Transformer-XL outperforms recurrent baselines, while the comparison supports the importance of CugLM’s pre-training and fine-tuning.Source: §4.3 · p. 9
• Type prediction reaches 68.89% accuracy on Java and 79.31% on TypeScript, and removing it worsens both identifier and all-token completion.The model predicts token and type jointly, using the predicted type to assist token prediction.Source: §4.3 · p. 9

5.1 The type of completions

• Punctuation dominates completions and is predicted accurately by both models, with BPE NLM [25] ahead of CugLM; CugLM leads on keywords, numerals, and operators.The authors attribute the keyword results to stronger syntactic and semantic representations from multi-task pre-training and fine-tuning; numeral and operator completions also favor CugLM substantially.Source: §5.1 · pp. 9–10, Table 4 · p. 10

5.2 Model complexity comparison

• CugLM has fewer trainable parameters than all baselines, with tasks sharing one Transformer network and parameter settings kept small for efficiency and to avoid over-fitting.The authors report that CugLM still outperforms the baselines by a large margin, attributing this to pre-training and fine-tuning.Source: §5.2 · p. 10, Table 5 · p. 10

5.3 Effect of applying BPE algorithm

• Applying BPE leaves Java performance comparable to the original model and slightly lowers TypeScript accuracy, so it does not improve CugLM.The authors explain that BPE can split masked identifiers into many masked subwords during pre-training and make type information semantically inconsistent with subword units during fine-tuning.Source: §5.3 · p. 10

5.4 Threats to Validity

• The authors identify dataset and language generalizability, empirical hyperparameter choices, baseline implementation and configuration, and accuracy as evaluation-measure limitations.The datasets contain GitHub Java and TypeScript programs; other languages require study. Li et al. [27] was adapted from AST-level to token-level input, while BPE NLM [25] was evaluated in its static setting; dynamic and maintenance scenarios remain future work. Accuracy measures correctly predicted next tokens and is widely used in prior code-completion work.Source: §5.4 · p. 10

Statistical Code Completion

• Statistical completion evolved from heuristic rules and static types to N-gram models, cache-based localness, and models handling unlimited vocabulary, nested scope, locality, and dynamism.Source: §6 · pp. 10–11
• Neural approaches include LSTM and pointer-mixture models [27], AST-level joint type-value prediction [29], syntax-aware Transformers [26], and open-vocabulary models [25].Source: §6 · p. 11

Pre-trained Language Models

• BERT contextual pre-training was extended to source code, where it outperformed LSTM and scratch-trained Transformer baselines on five program-classification tasks [24]; UNILM supports understanding and generation.Source: §6 · p. 11
• CodeBERT [12] pre-trains on paired natural-language and code data using masked language modeling and replaced-token detection, and evaluates on code search and documentation generation.Source: §6 · p. 11

• The authors present a transformer-based pre-trained model for code understanding and generation, tailored to code completion as its first application of pre-training for that task.Source: §6 · p. 11

• The proposed model outperforms previous state-of-the-art models in token completion, especially identifier completion.The authors describe it as the first application of a pre-trained language model to code completion.Source: §7 · p. 11
• The authors pre-train a Transformer-based code understanding and generation model on two massive datasets with three objectives, then fine-tune it for code completion.Source: §7 · p. 11
• Future work will apply the model to other programming languages and fine-tune it for additional tasks.Source: §7 · p. 11