一句话:提出 InferCode,以 AST 上下文中自动识别的子树为预测目标,自监督预训练通用代码表示。
代码表示学习广泛用于代码分类、代码搜索、注释生成与缺陷预测等任务,但现有技术多依赖针对特定下游任务标注的数据集进行训练,所得表示难以迁移到其他任务;即使从无标注代码中学习表示,其在下游任务上的效果也不理想。为突破这一限制,本文提出 InferCode,将自然语言处理中的自监督学习思想迁移到代码的抽象语法树(AST)上:其核心创新在于把 AST 上下文中自动识别出的子树当作预测标签来训练代码表示,因而无需任何人工标注,也无需构建代价高昂的图结构,训练出的表示不再绑定于特定下游任务或代码单元。作者以基于树的卷积神经网络(TBCNN)为编码器,在大量 Java 代码上训练了一个 InferCode 实例。该预训练模型可直接用于代码聚类、代码克隆检测、跨语言代码搜索等无监督任务,也可在迁移学习框架下继续微调权重,用于代码分类、方法名预测等有监督任务。与 code2vec、code2seq、ASTNN 等先前方法相比,使用预训练的 InferCode 在多数任务上取得显著更高的性能,其中还包括涉及不同编程语言的跨语言任务。
Learning code representations has found many uses in software engineering, such as code classification, code search, comment generation, and bug prediction, etc. Although representations of code in tokens, syntax trees, dependency graphs, paths in trees, or the combinations of their variants have been proposed, existing learning techniques have a major limitation that these models are often trained on datasets labeled for specific downstream tasks, and as such the code representations may not be suitable for other tasks. Even though some techniques generate representations from unlabeled code, they are far from being satisfactory when applied to the downstream tasks. To overcome the limitation, this paper proposes InferCode, which adapts the self-supervised learning idea from natural language processing to the abstract syntax trees (ASTs) of code. The novelty lies in the training of code representations by predicting subtrees automatically identified from the contexts of ASTs. With InferCode, subtrees in ASTs are treated as the labels for training the code representations without any human labelling effort or the overhead of expensive graph construction, and the trained representations are no longer tied to any specific downstream tasks or code units. We have trained an instance of InferCode model using Tree-Based Convolutional Neural Network (TBCNN) as the encoder of a large set of Java code. This pre-trained model can then be applied to downstream unsupervised tasks such as code clustering, code clone detection, cross-language code search, or be reused under a transfer learning scheme to continue training the model weights for supervised tasks such as code classification and method name prediction. Compared to prior techniques applied to the same downstream tasks, such as code2vec, code2seq, ASTNN, using our pre-trained InferCode model higher performance is achieved with a significant margin for most of the tasks, including those involving different programming languages. The implementation of InferCode and the trained embeddings are available at the link: https://github.com/bdqnghi/infercode.