code2vec: Learning Distributed Representations of Code

0x00 前言

读完这篇paper，第一感觉是：写得太详细（啰嗦）。好多东西翻来覆去讲了好几遍，最后总长度30页，确定不是水字数的？

但是不得不承认，论文里提到的方法很好。

0x01 Introduction

这篇paper提供了一种学习code的 Distributed Representation 的方法，借鉴了 word2vec 和 Attention 的方法。

它使用代码的抽象语法树结构，对代码片段在树上的路径（包括Terminal在内）进行 embedding ，这样得到的 code 的 embedding 可以用来做一些 downstream 的任务，原文做的任务是给出方法的代码，预测方法名。

0x02 AST Paths

对于代码片段，定义它的抽象语法树是一个六元组，其中代表非终结符（Non-Terminal）集合，代表终结符（Terminal）集合，代表终结符所对应的值的集合，代表语法树的根节点，是一个从 非终结符 到它在语法树上的 子孙节点 集合的映射，是从终结符到它对应的值的映射。

一条 AST Path 就是语法树上从一个终结符到另一个终结符之间的路径，中间会经过两终结符的LCA，可以被表示为一个三元组 $ $，$ x_s $和$ x_t $代表这两个终结符对应的值，$ p$ 代表两个终结符之间的路径（即经过的节点和父子关系，包括这两个终结符）。

如语句

x = 7

就可以被表示为。

0x03 Embedding Model

AST Path Embedding

一条 AST Path 的 embedding 由它所对应的三元组（两个终结符的值 + 中间路径）所对应的 embedding 连接起来，再经过一个全连接层得到。对于终结符的值和路径的 embedding ，采用 word2vec 的方法，即用 one-hot 向量和 embedding 矩阵作乘法，即可得到对应的 embedding 向量。

具体来说，有和分别代表终结符值的 embedding 矩阵和路径的 embedding 矩阵，矩阵中的每一行代表一个元素的 embedding 向量，每一个 embedding 向量都是维的，这个是 empirical 的 hyperparameter 。用终结符的值/路径所对应的 one-hot 向量和 embedding 矩阵作乘法，实际就是选出 embedding 矩阵中对应的那一行，作为其 embedding 向量。

将这三个向量连接起来，得到，再与作乘法，即经过一个全连接层，在经过作为激活层，即得到作为一条 AST Path 的 embedding。

Code Snippet Embedding

在得到代码片段中所有 AST Path 的 embedding 之后，对每条 AST Path 进行 Attention 操作，得到 attention weight ，将所有 AST Path 的 embedding 按照这个 attention weight 进行加权求和，得到整个代码段的 embedding 。

具体来说，我们有一个全局的 attention 向量，将每个 AST Path 的 embedding 和这个 attention 向量作 softmax ，即可得到 attention weight 。那么什么是 softmax ？

softmax 是指，对于得到的 AST Path 的 embedding 们，通过下面的式子可以计算得到一个值

$\alphai=\frac{exp(c^T_i\cdot a)}{\sum{j=1}^n exp(c^T_j\cdot a)}$

，这个值满足，也就是可以看成一个概率分布。在这里，我们把它作为加权值，来计算代码段的 embedding 。

我们将每条 AST Path 对应的 embedding 和它对应的加权值乘起来，再相加，得到整个代码片段的 embedding，即

。

Label Embedding

对于每个代码段的标签（在这里就是这个方法对应的方法名），也采用 word2vec 的方法对其进行 embedding ，即有，将标签对应的 one-hot 向量与这个 embedding 矩阵相乘，得到对应的 embedding。

0x04 Prediction Model

在得到代码段和各个标签的 embedding 以后，我们需要预测这个代码段是每个标签的概率。我们这样计算代码段的标签是的概率：

$q(yi) = \frac{exp(v^T\cdot tags_vocab_i)}{\sum{y_j\in Y}exp(v^T\cdot tags_vocab_j)}$

，其中是标签集合。也就是用代码段的 embedding 和标签的 embedding 作 softmax ，之前提到过，这样的值可以作为一个概率分布，。

0x05 Training

使用交叉熵损失（Cross Entropy Loss）作为损失函数。，其中是真实的概率分布中，这一分量的概率，是模型预测中这一分量的概率。由于我们的标签都是 one-hot 向量，只有一个分量的概率为，其他分量均为，所以交叉熵损失在形式上和对率回归一样，即。

使用反向传播机制+梯度下降算法，即可完成模型训练。