【AI4Code】《InferCode: Self-Supervised Learning of Code Representations by Predicting Subtrees》ICSE‘21

《InferCode: Self-Supervised Learning of Code Representations by Predicting Subtrees》 ICSE 2021

致力于自监督的训练出代码的表征，不利用任何标签。

提出InferCode，利用自监督学习从代码的抽象语法树（AST）中学习代码的表示，通过预测AST中的子树构造训练标签，而子树结构是可以在构造AST过程中自动生成的。基于的假设是：相似的代码片段的AST包含相似的子树，比如两种写法的冒泡排序代码，包含大量相似代码结构：

InferCode的无监督的下游任务有代码聚类、克隆检测以及跨语言代码搜索；有监督的迁移学习有代码分类和代码推荐。

方法

概述

InferCode与Doc2vec类似，将AST视为文档，子树作为文档中的单词。给定AST集合 $\{T_1,T_2,\dots ,T_{n\}}$，$T_i$的所有子树集合 $\left\{\dots ,T_{ij},\dots \right\}$，将 $T_i$ 和 $T_{ij}$表示为D维向量，考虑子树 $T_{ij}$ 出现在 $T_i$ 上下文中，最大化对数损失：${\sum }_j\log P_r\left(T_{ij}\mid T_i\right)$。所以InferCode的步骤为：

1. 对数据集预处理，为每段代码生成AST，并且识别出其中的子树的集合。所有子树集合累计构造成子树语料库。
2. 用TBCNN编码AST生成代码向量，用于预测子树。
3. encoder训出来用于下游任务。

识别子树

遍历AST，将满足指定条件的节点作为子树的根节点，识别出子树。指定条件为（expr_stmt，decl_stmt，expr，condition）。此外也考虑一些关键节点，例如 if，for，while，这些单独的节点被作为一个子树。

因此子树的选择条件是子树结构比较小，作者称这样细粒度子树更容易在代码片段中出现。而粗粒度的子树，如if，while，for语句块，这些子树太大。作为一个单独的词，在代码库中出现不频繁，encoder很难直接学习到有意义的表示；在这些大的子树种虽然句法不同并不意味着代码的功能不同，encoder很难学习到他们的相同点。

文章给出了一冒泡排序代码识别出的子树的例子：

学习代码表示

首先将多叉树AST转为二叉树，然后使用Tree-based CNN（TBCNN）作为encoder。

首先每个节点的初始embedding 是 Type embedding和Token embedding经过一个线性层的输出。

然后通过TBCNN的卷积层可以学习节点的信息。

TBCNN (AAAI‘16)

TBCNN的思想类似于GCN，TBCNN设计了针对树结构数据的三个卷积核：${\mathbf{W}}^t,{\mathbf{W}}^l,{\mathbf{W}}^r\in {\mathbb{R}}^{D\times D}$ 分别表示“top”，“left”和“right”，因此对于AST卷积窗口 “depth d” 包含了 $K=2^d-1$ 个节点 $\left[{\mathbf{x}}_1,\dots ,{\mathbf{x}}_K\right]$，该窗口的最终输出为：
$$\mathbf{y}=\tanh \left(\sum _{i=1}^K\left[{\eta }_i^t{\mathbf{W}}^t+{\eta }_i^l{\mathbf{W}}^l+{\eta }_i^r{\mathbf{W}}^r\right]{\mathbf{x}}_i+\mathbf{b}\right)$$
即对于该窗口中的任一节点，卷积核的权重矩阵根据该节点在窗口中的深度和位置进行不同的线性组合，以实现树卷积的效果。

经过卷积后每个节点得到一个特征向量，InferCode使用attention来代替maxpooling，聚合所有节点的信息，得到代码片段的向量表示。是通过一个可学习的attention向量实现的：

最终 $\vec{v}$ 是这段代码的向量表示。

预测子树

注意我们在数据预处理已经抽取出了若干子树结构，这些子树结构会构成一个字典 $L$（源码里大小为10w），每个子树分配一个可学习的向量表示，然后用softmax结构表示该子树在代码片段中的概率：
$$\text{ for }{l}_i\in L:q\left(l_i\right)=\frac{\exp \left({\vec{v}}^T\cdot {\mathbf{W}}_i^{\text{subtrees }}\right)}{{\sum }_{l_j\in L}\exp \left({\vec{v}}^T\cdot {\mathbf{W}}_i^{\text{subtrees }}\right)}$$
所以这里实质是将子树作为label，通过判断该子树label是否存在该代码中，这样来检测从代码中学习到的向量表示是否包含了该子树的信息，或者说 学习出的代码表示是子树信息的组合。

所以InferCode的训练参数为 ${\mathbf{W}}^{\text{type }},{\mathbf{W}}^{\text{token }},{\mathbf{W}}^t,{\mathbf{W}}^l,{\mathbf{W}}^r\in {\mathbb{R}}^{D\times D},a\in {\mathbb{R}}^D,{\mathbf{W}}^{\text{subtrees }}\in {\mathbb{R}}^{|L|\times D}$，参数还是非常轻量的。

评估

三个无监督任务：代码聚类、克隆检测 和 跨语言代码搜索；两个有监督任务：代码分类、代码推荐。

代码聚类

两个数据集：OJ dataset，包含52000个C代码片段，属于104个类；Sortting Algorithm dataset，10类Java的排序代码，每个类别包含1000段代码。

评价聚类效果的指标 Adjusted Rand Index：用C表示实际的类别划分，K表示聚类结果。定义a 为在C中被划分为同一类，在K中被划分为同一簇的实例对数量。定义b为在C中被划分为不同类别，在K中被划分为不同簇的实例对数量。定义Rand Index（兰德系数）：
$$RI=\frac{a+b}{\left(\begin{array}{l}n\\ 2\end{array}\right)}$$
即从所有样本中任意选两个出来，聚类结果和GroudTruth一致。**RI越趋近于1越好。**Rand Index无法保证随机划分的聚类结果的RI值接近0。于是，提出了Adjusted Rand index（调节的兰德系数）：
$$ARI=\frac{RI-E[RI]}{\max (RI)-E[RI]}$$

克隆检测

构造了50000克隆代码对和50000非克隆代码对作为正样本和负样本。

跨语言代码搜索

收集了Java、C、C++、C#特定算法实现的代码样本3000左右从Rosetta Code，对于每种语言再随机5000个代码文件从Github中。

代码分类

监督任务，在InferCode基础上加一个分类器微调。

代码推荐

监督任务，在InferCode基础上加一个softmax层微调。

代理任务标签选择

除了可以选择用子树作为标签外，还可以选择代码token、方法名作为预训练任务的预测标签