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

对于代码片段 \(C\) ，定义它的抽象语法树是一个六元组 \(<N, T, X, s, \delta, \phi>\) ，其中 \(N\) 代表非终结符（Non-Terminal）集合， \(T\) 代表终结符（Terminal）集合， \(X\) 代表终结符所对应的值的集合， \(s\in N\) 代表语法树的根节点， \(\delta: N\rightarrow (N\cup T)^\*\) 是一个从 非终结符 到它在语法树上的 子孙节点 集合的映射， \(\phi:T\rightarrow X\) 是从终结符到它对应的值的映射。

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

如语句

x = 7

就可以被表示为 \(<x, (NameExpr\uparrow AssignExpr\downarrow IntegerLiteralExpr), 7>\) 。

0x03 Embedding Model

AST Path Embedding

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

具体来说，有 \(valueVocab\in \mathbb{R}^{|X|\times d}\) 和 \(pathVocab\in\mathbb{R}^{|P|\times d}\) 分别代表终结符值的 embedding 矩阵和路径的 embedding 矩阵，矩阵中的每一行代表一个元素的 embedding 向量，每一个 embedding 向量都是 \(d\) 维的，这个 \(d\) 是 empirical 的 hyperparameter 。用终结符的值/路径 所对应的 one-hot 向量和 embedding 矩阵作乘法，实际就是选出 embedding 矩阵中对应的那一行，作为其 embedding 向量。

将这三个向量连接起来，得到 \(c'\in\mathbb{R}^{3d}\) ，再与 \(W\in\mathbb{R}^{d\times 3d}\) 作乘法，即经过一个全连接层，在经过 \(tanh\) 作为激活层，即得到 \(c=tanh(W\cdot c')\) 作为一条 AST Path 的 embedding。

Code Snippet Embedding

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

具体来说，我们有一个全局的 attention 向量 \(a\in\mathbb{R}^d\) ，将每个 AST Path 的 embedding 和这个 attention 向量作 softmax ，即可得到 attention weight 。那么什么是 softmax ？

softmax 是指，对于得到的 AST Path 的 embedding 们 \({c_1, c_2, ...,c_n}\) ，通过下面的式子可以计算得到一个值

\(\alpha_i=\frac{exp(c^T_i\cdot a)}{\sum_{j=1}^n exp(c^T_j\cdot a)}\)

，这个值满足 \(\sum_{i=1}^n\alpha_i = 1\) ，也就是可以看成一个概率分布。在这里，我们把它作为加权值，来计算代码段的 embedding 。

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

\(v=\sum_{i=1}^n\alpha_i\cdot c_i\)

。

Label Embedding

对于每个代码段的标签（在这里就是这个方法对应的方法名），也采用 word2vec 的方法对其进行 embedding ，即有 \(tagsVocab\in\mathbb{R}^{|Y|\times d}\) ，将标签对应的 one-hot 向量与这个 embedding 矩阵相乘，得到对应的 embedding。

0x04 Prediction Model

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

\(q(y_i) = \frac{exp(v^T\cdot tags\_vocab_i)}{\sum_{y_j\in Y}exp(v^T\cdot tags\_vocab_j)}\)

，其中 \(Y\) 是标签集合。也就是用代码段的 embedding 和标签的 embedding 作 softmax ，之前提到过，这样的值可以作为一个概率分布， \(\sum_{y_i\in Y}q(y_i) = 1\) 。

0x05 Training

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

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