[论文解读] InferCode: Self-Supervised Learning of Code Representations by Predicting Subtrees

Nghi D. Q. Bui, Yijun Yu+1 more|arXiv (Cornell University)|Dec 13, 2020

Software Engineering Research参考文献 65被引用 9

一句话总结

InferCode 通过在代码的抽象语法树（AST）中训练基于树的卷积神经网络（TBCNN）来预测子树，提出了一种用于代码表示的自监督学习方法，从而实现可迁移的、与任务无关的嵌入表示。该模型在大规模无标签 Java 代码上进行预训练，实现了在代码聚类、克隆检测和跨语言搜索等多种下游任务上的最先进性能，显著优于先前的方法（如 code2vec 和 ASTNN）。

ABSTRACT

Building deep learning models on source code has found many successful software engineering applications, such as code search, code comment generation, bug detection, code migration, and so on. Current learning techniques, however, have a major drawback that these models are mostly trained on datasets labeled for particular downstream tasks, and code representations may not be suitable for other tasks. While some techniques produce representations from unlabeled code, they are far from satisfactory when applied to downstream tasks. Although certain techniques generate representations from unlabeled code when applied to downstream tasks they are far from satisfactory. This paper proposes InferCode to overcome the limitation by adapting the self-supervised learning mechanism to build source code model. The key novelty lies in training code representations by predicting automatically identified subtrees from the context of the ASTs. Subtrees in ASTs are treated with InferCode as the labels for training code representations without any human labeling 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 trained an InferCode model instance using the Tree-based CNN as the encoder of a large set of Java code and applied it to downstream unsupervised tasks such as code clustering, code clone detection, cross-language code search or 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 previous code learning techniques applied to the same downstream tasks, such as Code2Vec, Code2Seq, ASTNN, higher performance results are achieved using our pre-trained InferCode model with a significant margin for most tasks including those involving different programming languages.

研究动机与目标

• 为解决现有任务特定代码表示方法依赖人工标注且在软件工程任务间缺乏泛化能力的局限性。
• 开发一种自监督方法，从无标签源代码中学习代码表示，无需人工标注或复杂的图结构构建。
• 为可被解析为 AST 的任意代码单元提供灵活且可重用的代码嵌入表示，支持多样化的下游任务。
• 相比现有的监督或弱监督方法，提升代码表示在不同编程语言和任务间的可迁移性与性能表现。
• 探究在 AST 中进行子树预测是否可作为学习鲁棒、通用代码嵌入的有效掩码任务。

提出的方法

• 模型使用基于树的卷积神经网络（TBCNN）作为编码器，将抽象语法树（ASTs）映射为密集的代码向量。
• 掩码任务为从 AST 上下文中预测随机采样的子树，利用代码向量作为上下文表示来重建子树。
• 子树通过自动从 AST 中识别，无需人工标注，作为预训练阶段的自生成标签。
• 模型在大规模无标签 Java 代码数据集上进行预训练，学习来自 AST 的结构与语义模式。
• 预训练完成后，编码器可微调用于监督任务，或直接用于无监督任务（如代码聚类与搜索）。
• 该方法避免了昂贵的图结构构建，且由于 AST 的通用可解析性，具有跨编程语言的泛化潜力。

实验结果

研究问题

• RQ1 在 AST 中进行子树预测是否可作为学习通用代码表示的有效自监督掩码任务？
• RQ2 所提出的自监督方法是否能产生比任务特定或弱监督方法更具泛化能力的代码嵌入？
• RQ3 InferCode 在代码聚类、克隆检测和跨语言搜索任务上的性能与 state-of-the-art 模型（如 code2vec、code2seq 和 ASTNN）相比如何？
• RQ4 预训练的 InferCode 模型是否可有效微调用于监督任务（如代码分类和方法名预测）？
• RQ5 该模型在不同编程语言上的性能是否稳健，特别是在跨语言代码搜索任务中？

主要发现

• 在代码聚类任务中，InferCode 的 ARI 达到 0.70，显著优于 code2vec（0.58），相对提升 17%。
• 在克隆检测任务（BigCloneBench）中，InferCode 的精确率高达 0.90，显著优于 code2vec（0.81）和 code2seq（0.45）。
• 在跨语言代码搜索任务（Rosetta Stone）中，InferCode 的 MRR 达到 0.57，优于 code2vec（0.41）和 code2seq（0.32）。
• 将预训练的 InferCode 模型微调用于监督任务，性能优于从零开始训练，证明了其强大的可迁移性。
• 模型在不同编程语言上的性能表现稳健，表明其具备超越单一语言的泛化能力。
• 与方法名预测相比，子树预测作为掩码任务能生成更优的表示，因为后者会迫使模型学习到名称相似但实现不同的虚假相关性。