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)$ 。