论文笔记 - 基于抽象语法树的代码表示方法（ASTNN）

摘要

• 传统的信息检索将程序视为自然语言文本，会遗漏源代码的重要语义信息
• 基于抽象语法树（AST）的神经网络模型可以更好地表示源代码
• AST的规模较大，容易出现长期依赖的问题
• 文章中提出了一种新型的基于AST的神经网络进行源代码表示：
• • 将大型的AST分割为一系列小的语句树序列
• • 通过捕获语句的词汇和句法知识，将每一个语句树都编码为一个向量
• • 基于语句序列向量，利用双向RNN模型生成代码的向量表示
• • ASTNN在源代码分类和克隆代码检测上都得到了最优的结果

一、简介

• 代码分类、克隆检测、缺陷预测等软件工程方法都以源代码为研究对象
• 适当地表示源代码，可以有效地捕获源代码中存在的语法和语义信息
• 现有基于AST的神经网络模型有两个缺陷：
  ①梯度消失问题——使用自上而下的方式或滑动窗口技术对规模大、节点深的 AST 进行遍历和编码时可能会丢失一些长期上下文信息
  ②通常将 AST 视为完全二叉树——会破坏源代码的原始语法结构，导致生成的 AST 更加庞大、节点更加深入，会进一步削弱了神经模型捕获真实和复杂语义的能力

本文提出了一种新颖的方法：AST-based Neural Network

• 能够在不经过编译的情况下生成目标源码片段的代码表示
• 能够表示代码片段中贮藏的、语句(statement)级的词法和语法信息。
• • 这里的语句(statement)均采用普遍使用的程序语言规范
• • 并将 方法声明(MethodDeclaration) 也视为一种特殊的语句节点

下图是一段来自开源软件项目的源代码

• 第5行到第6行之间的代码段只包含了初始化变量sb的局部变量(Local Variable)语句
• 第7行到第15行之间的代码段包含了一个完整的Try语句
• 和Try相似的语句往往由两个部分构成：语句头以及包含具体功能语句的语句体。
• 本文将TRY这类语句的语句头和语句体中包含的若干语句分别视为独立的部分，并将其拆分成不同的单元。
• 通过这种方式，大型 AST 将最终被分解为一串小型的语句树序列。
• 之后，本文采用循环神经网络(RNN)将语句之间的顺序依赖关系编码成向量，通过这样的向量捕获源代码的固有性质，形成可以用作神经源代码表示的结构。

方法分为三步：

1. 首先，从代码片段构建 AST，然后再将 AST 拆分为小语句树(一个语句树由一条语句的AST节点组成，并以statement节点为根)。图 1 中，语句树用虚线表示，原始代码片段中相应的语句(或语句头)也用虚线标记
2. 在多路语句树上设计了一个递归编码器来捕获语句级的词汇和语法信息，然后将它们表示为语句向量
3. 最后，基于语句向量的顺序，本文使用门控循环单元(Gate Control Unit，GRU，一种递归神经网络)，结合 ② 获得的语句向量，最终获得整个代码片段的向量表示

astnn 的目标是学习更多关于源代码的语法和语义信息，而不是最先进的基于ast的神经模型。它具有通用性，可用于许多与程序理解相关的任务，如源代码分类和代码克隆检测。

二、背景

1. AST（抽象语法树）

• AST 是一种用于表示源代码的抽象语法结构的树。其组成部分对应于源代码的结构或符号。
• • 一方面，与普通源代码相比，ast是抽象的，不包括标点和分隔符等所有细节。
• • 另一方面，AST 可用于描述源代码的词汇信息和语法结构，如图1(b)中的方法名称readText和控制流结构WhileStatement。

2. 基于树的神经网络
   基于树的神经网络(TNNs)——接受ast作为输入。给定一棵树，TNNs 通过自底向上递归计算节点嵌入来学习它的向量表示。最具代表性的基于树的ast模型有递归神经网络(RvNN)、基于树的CNN (TBCNN)和基于树的长短期记忆(Tree-LSTM)。

3. 现有工作的限制
   这三种基于树的方法有两个主要的局限性。

首先，在基于梯度的树拓扑训练中，通过反向传播结构来计算梯度。由于程序的复杂性，ast的结构通常又大又深，尤其是嵌套结构。因此，从叶节点到根节点的自下而上的计算可能会遇到梯度消失的问题，并且很难捕获长期依赖关系，这将错过一些由距离根节点很远的节点携带的语义，例如叶节点中的标识符。
其次，现有的基于树的方法将ast视为二叉树，将父节点的三个或更多子节点移动到新的子树中进行简化，这改变了源代码的原始语义，使长期依赖问题更加严重。

三、本文方法

本节将介绍基于AST的神经网络(ASTNN)。ASTNN的整体架构如图2所示。
首先，将源代码片段解析为AST，并设计一个先序遍历算法，将每个AST分割为语句树序列(st-树，它是由作为根的语句节点和相应的语句AST节点组成的树)，如图1所示。
所有st-树都由语句编码器编码为向量，表示为e1，····等。然后我们使用双向门控循环单元(Bi-GRU)来建模语句的自然性。Bi-GRU的隐藏状态通过池化采样到单个向量，这是代码片段的表示。

1. 拆分ast和构造st-树序列
   首先，源代码片段可以通过现有的语法分析工具转换为大型AST。对于每个AST，我们按照语句的粒度对其进行拆分，并通过先序遍历提取语句树的序列。

• 给定一棵AST树 T 和一组语句(statement)节点 S ，T 中的每一个语句节点 s∈S 对应源代码中的一条语句。
• 方法声明(MethodDeclaration) 也视为一种特殊的语句节点，所以S∪{MethodDeclaration}
• 对于嵌套语句（图1），我们定义了一组独立的节点 P = {block, body}，其中block用于拆分嵌套语句（如Try和While语句）的语句头和语句体，body用于方法声明。
• 语句节点 s∈S 的所有后继结点用 D(s) 表示。对于任意d∈D(s)，如果存在从 s 到 d 经过节点 p∈P的一条路径，则表示节点 d 被一条语句包含在语句s的body中。我们称节点 d 为 s 的一个子语句（substatement ）节点。
• 那么，以语句节点 s∈S 为根的语句树（statement tree）就是由节点 s 及其所有子节点（不包括 T 中的子语句节点）组成的树。
• 例如，图1(b)中以MethodDeclaration为根的第一个语句树被虚线包围，其中包括“static”、“public”和“readText”等头部分，并排除了body中两个局部变量、一个Try和一个Return语句的节点。
• 由于一个 st树的节点可以有三个或更多的子节点，我们也称它为多路st树，以区别于二叉树。这样，一个大的AST可以分解成一个不重叠的多路 st树序列。

AST 的拆分和 st树序列的构造是通过一个遍历器和一个构造器直接实现的。

• 遍历器以深度优先的先序遍历方式访问 AST 中的每个节点
• 构造器递归地创建 st树，并按顺序将其添加到 st树序列中。

这样的做法保证了一个新的 st树是按照源代码中的顺序追加的。这样，就得到 st树的序列作为ASTNN的原始输入。选择 st树，因为语句是携带源代码语义的基本单元。我们的实验结果表明，所提出的语句级粒度更好，因为它在st树的大小和句法信息的丰富度之间有很好的权衡

2. 在多路st-树进行语句编码

① 语句向量：
设计了一个基于RvNN的语句编码器，用于学习语句的向量表示。

• 我们通过对 AST 进行前序遍历，得到所有的符号作为训练语料。
• 使用word2vec学习符号的无监督向量，并将训练好的符号嵌入作为语句编码器的初始参数
• 由于叶子节点通常包含词法信息(如标识符)，因此符号嵌入可以很好地捕获词法信息。
  
  
  
  
• 以图1中MethodDeclaration的节点为根的第一个st-树为例，编码器遍历st-树，递归地将当前节点的符号作为新的输入，连同其子节点的隐藏状态进行计算。如图3所示。这里只展示了前两层。
• 在st树中，两个子节点readText(即方法名)和FormalParameter(即定义方法参数的语法结构)以及其他兄弟节点丰富了MethodDeclaration的含义。如果将st树转换为一棵二叉树，例如将readText的节点移动到FormalParameter节点的一个子节点或后代节点上，则可能会破坏原有的语义。相反，我们将原始的多路st树作为输入。
  
  
  
  
  
• 给定一棵 st树，n表示非叶子节点，C表示其子节点的数量。
  一开始，使用预训练好的嵌入参数：$$W_e\in {\mathbb{R}}^{|V|\times d}$$
  这里的 V 是词汇量， d 是符号的嵌入维度，节点 n 的词法向量可以表示成（xn是符号n的 one-hot表示，vn是嵌入）：$$v_n=W_e^{\top }{x}_n（1）$$
  接下来，节点 n 的向量表示就是下面的公式：$$h=\sigma \left(W_n^{\top }{v}_n+\sum _{i\in [1,C]}{h}_i+b_n\right)（2）$$
  $$W_n\in {\mathbb{R}}^{d\times k}是编码维度为k的权重矩阵，b_n是偏置项h_i是每个子节点i的隐藏状态，h是更新后的隐藏状态，\sigma 是激活函数，通常为tanh或恒等函数。$$

本文采用恒等函数。类似地，我们可以递归地计算和优化 st-树 t 中所有节点的向量。

另外，为了确定节点向量的最重要特征，所有节点都被推入堆栈，然后通过最大池化（max-pooling） 进行采样。
我们通过下面这个公式得到 st-树 的最终表示形式和相应的表述，其中N为 st-树 的节点数。
$$e_t=\left[\max \left(h_{i1}\right),\cdots ,\max \left(h_{ik}\right)\right],i=1,\cdots ,N（3）$$

这些语句向量可以捕获语句的词汇级和语句级语法信息。

② 批处理

• 批处理——为了提高大数据集上的训练效率，需要设计批处理算法同时对多个代码片段进行编码。
• 通常在多路st树上进行批处理比较困难，因为同一批的父节点在同一位置上的子节点数量是不同的。
• 例如，给定下图中的两个父节点ns1和ns2，其中ns1有3个子节点，ns2有2个子节点，由于C（子节点的数量）不同，无法直接计算出同一批的两个父节点的公式2。

• ST-树动态批处理算法1：

直观地看，虽然父节点的子节点数量不同，但该算法可以动态地检测出所有可能的位置相同的子节点，并将其放入组中，然后利用矩阵运算批量地加快每组方程2的计算速度。

• 在算法1中，首先锁定同一批的L个样本，然后从根节点（第4行）开始对它们进行宽度优先遍历
• 对于批次中具有相同位置节点ns的样本，算法将首先计算等式 1（第10行）
• 然后根据节点位置对其所有子节点进行检测和分组（第12-16行）
• 如图4所示，我们将子节点按位置分为三个组，并在数组列表C和CI中记录这些组。
• 基于这些组，我们递归地在所有子节点上执行批处理（第17-21行）
• 在得到所有子节点的结果后，分批计算公式(第22行)
• 批处理所有的 ST-tree 节点向量压入堆栈S(第 24 行)
• 最后，通过公式3中的池化操作得到ST-树 样本的向量和对应的语句(第5行)
  

3. 表示语句序列

基于 ST-tree 向量序列，本文将利用 GRU（门控循环单元） 来跟踪语句的自然度（naturalness）
给定一个代码片段，假设我们从AST中提取出了ST树 T，令
$Q\in {\mathbb{R}}^{T\times k}=$ $\left[e_1,\cdots ,e_t,\cdots ,e_T\right],t\in [1,T]$

表示 ST-树 序列中已编码的 ST-树 向量，在时间 t 内，转换方程为：
r t = σ ( W r e t + U r h t − 1 + b r ) z t = σ ( W z e t + U z h t − 1 + b z ) h t ~ = tanh ⁡ ( W h e t + r t ⊙ ( U h h t − 1 ) + b h ) h t = ( 1 − z t ) ⊙ h t − 1 + z t ⊙ h ~ t （ 4 ）

$$\begin{aligned} r_t & =\sigma\left(W_r e_t+U_r h_{t-1}+b_r\right) \\ z_t & =\sigma\left(W_z e_t+U_z h_{t-1}+b_z\right) \\ \tilde{h_t} & =\tanh \left(W_h e_t+r_t \odot\left(U_h h_{t-1}\right)+b_h\right) \\ h_t & =\left(1-z_t\right) \odot h_{t-1}+z_t \odot \tilde{h}_t \end{aligned}$$ （4） rt​zt​ht​~​ht​​=σ(Wr​et​+Ur​ht−1​+br​)=σ(Wz​et​+Uz​ht−1​+bz​)=tanh(Wh​et​+rt​⊙(Uh​ht−1​)+bh​)=(1−zt​)⊙ht−1​+zt​⊙h~t​​（4）
其中：r是重置门，用于控制前一状态的影响
zt 是更新门，用于组合旧信息和新信息
ht是候选状态，用于与上一状态ht-1一起完成线性插值，以确定当前状态
ht，W,，Wz，Wh，Ur，Uz，Uh∈Rk×m是权重矩阵，而br，bz，bh 是偏差项。
通过迭代计算所有时间步骤的隐藏状态后，可以得到这些语句的自然顺序。


为了进一步增强递归层获取依赖信息的能力，本文采用双向GRU，将两个方向的隐藏状态连接起来，形成新的状态，如公式5所示：
$$\begin{array}{rl}\overrightarrow{h_t}& =\overrightarrow{GRU}\left(e_t\right),t\in [1,T]\\ \overleftarrow{h_t}& =\overleftarrow{GRU}\left(e_t\right),t\in [T,1]\\ h_t& =\left[\overrightarrow{h_t},\overleftarrow{h_t}\right],t\in [1,T]\end{array}（5）$$
与语句编码器类似，这些状态的最重要特征将通过最大池化或平均池化进行采样。考虑到直观上，不同语句的重要性是不相等的。例如，MethodInvocation 语句中的API调用可能包含更多的函数信息，因此本文默认使用最大池化来捕获最重要的语义。模型最终生成一个向量：
$$r\in {\mathbb{R}}^{2m}$$即目标代码片段的向量表示。

四、模型的应用

ASTNN是通用的，它可以被训练作为源代码片段的特定任务向量表示，为许多程序理解任务描述不同的源代码语义。

1. 代码分类

• 根据代码片段的功能对其进行分类
  已知代码片段向量 r 和类别数 M，我们通过这个公式获得logits（全连接层的输出，softmax的输入）
  $$\hat{x}=W_{o}r+b_o\text{, 这里}{W}_o\in {\mathbb{R}}^{2m\times M},b_o是偏执项$$
  使用交叉熵损失函数：
  $$J(\Theta ,\hat{x},y)=\sum \left(-\log \frac{\exp \left({\hat{x}}_y\right)}{{\sum }_j\exp \left({\hat{x}}_j\right)}\right)（6）$$
  其中 Θ 表示模型中所有权重矩阵的参数， y 为真实标签。

2. 克隆代码检测

• 克隆代码检测是：检测两个代码片段是否实现了相同的功能
  已知两个代码片段向量 r1 和 r2 ，用下面这个公式来衡量它们之间的语义相关性：
  $$r=\left|r_1-r_2\right|$$
  然后用输出 y 作为它们之间的相似性
  $\hat{y}=$ $\mathrm{sigmoid}(\hat{x})\in [0,1]$，这里 $\hat{x}=W_{o}r+b_o$
  损失函数定为二元交叉熵：$$J(\Theta ,\hat{y},y)=\sum (-(y\cdot \log (\hat{y})+(1-y)\cdot \log (1-\hat{y})))（7）$$

• 训练ASTNN模型，使损失最小化
• 本文采用AdaMax优化器，因为计算效率高
• 在所有的参数优化后，存储训练好的模型
• 对于新的代码片段，将其预处理为 st-树 序列，然后输入重新加载的模型进行预测
• 输出是不同标签的概率 p
  
  
• 代码分类是多类别任务，可以通过下面的方式得到预测值：$$\text{ prediction }=\underset{i}{\mathrm{argmax}}\left(p_i\right),i=1,\cdots ,M（8）$$
• 而克隆检测是二分类任务， p 是一个 [0,1] 的数，可以用下面的方式得到预测值：
  $$\text{ prediction }=\{\begin{array}{ll}\text{ True, }& p>\delta \\ \text{ False, }& p\le \delta \end{array}，\delta 是阈值（9）$$