[论文解读] Multi-task Learning based Pre-trained Language Model for Code Completion

Fang Liu, Ge Li+2 more|arXiv (Cornell University)|Dec 29, 2020

Software Engineering Research参考文献 46被引用 8

一句话总结

本文提出了一种基于多任务学习的预训练语言模型，采用Transformer架构，通过联合预测标记及其类型来改进代码补全，解决了静态嵌入表示和标识符补全效果差的局限性。该模型在两个真实世界数据集上实现了最先进性能，尤其在识别和建议正确标识符方面表现优异。

ABSTRACT

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.

研究动机与目标

• 解决静态词嵌入无法捕捉源代码中上下文相关语义的局限性。
• 通过在预测过程中引入类型信息，提升代码补全性能，特别是标识符的补全效果。
• 开发一个统一的预训练模型，支持代码理解与代码生成两类任务。
• 通过在预训练过程中采用混合目标的多任务学习方法，增强代码中上下文表征的学习能力。
• 展示联合标记与类型预测在IDE中实时代码补全中的有效性。

提出的方法

• 该模型采用基于Transformer的神经架构，学习代码标记的上下文表征。
• 在大规模代码仓库上使用三种目标函数进行预训练：掩码语言建模、下一项标记预测和类型预测。
• 采用多任务学习方法，联合预测下一个标记及其数据类型，其中类型预测用于引导标记生成。
• 在微调阶段，通过同时优化标记和类型预测的准确率，将模型适配于代码补全任务。
• 该模型利用双向上下文和注意力机制，捕捉代码序列中的长距离依赖关系。
• 采用混合预训练策略，结合代码理解与代码生成目标，以提升泛化能力。

实验结果

研究问题

• RQ1 与单任务模型相比，联合预测标记和类型是否能提升代码补全的准确性？
• RQ2 在预测过程中引入类型信息是否能提升标识符补全的性能？
• RQ3 多任务预训练语言模型是否能超越现有的基于静态嵌入的模型在代码补全中的表现？
• RQ4 所提出的模型在真实世界代码补全基准上的有效性如何？
• RQ5 该模型是否能在不同编程语言和代码模式间实现良好泛化？

主要发现

• 所提出的模型在两个真实世界代码补全数据集上实现了最先进性能，优于以往的SOTA方法。
• 该模型显著提升了标识符补全的准确率，解决了以往基于语言模型方法的关键局限。
• 联合预测标记与类型可生成更准确且更符合上下文的代码补全建议。
• 多任务预训练策略增强了表征学习能力，尤其对罕见或歧义标识符表现更优。
• 该模型在多样化代码模式和编程结构中表现出强鲁棒性与泛化能力。
• 消融实验证实，类型预测和多任务学习的引入对性能提升贡献最大。