基于神经网络注意力机制和指针网络的代码补全

Posted 2021-04-26 数据挖掘与开源生态

摘要

智能代码补全已经成为加快现代软件开发的一项重要研究课题。为了方便编程语言的有效代码动态补全，我们从大型代码库中学习神经语言模型，并开发一个定制的注意力机制来完成代码补全。然而，标准的神经语言模型即使有注意力机制，也不能正确地预测词汇表外单词（OoV），这限制了代码补全性能。受程序源代码中局部重复词的广泛使用以及最近提出的指针复制机制的启发，本文提出了一种指针混合网络来更好地预测OoV词，进而实现代码补全。基于上下文，指针混合网络学习通过RNN组件生成词汇表内的单词，或者通过指针组件从局部上下文生成OoV词。在两个基准标记数据集上的实验，证明了我们的注意力机制和指针混合网络在代码补全任务上的有效性。

1 引言

集成开发环境（IDEs）已经成为现代软件工程师的基本范式，因为IDE提供了一组有用的服务来加速软件开发。智能代码补全是IDE中最有用的特性之一，它根据上下文中现有的代码推荐下一个可能的代码标记，比如方法调用或对象字段。传统的代码补全很大程度上依赖于编译时的类型信息来预测下一个标记[Tu等人，2014]，因此，该方法可以适用于静态类型语言，比如Java。然而，由于缺少类型注释，像javascript和Python这样的动态类型语言的代码补全更加困难，从而支持动态类型语言的代码补全也更少。

为了有效地实现动态类型语言的代码补全，最近学者们转向基于学习的语言模型的研究[Hindle et al.,2012;White et al.,2015;Bielik等,2016]。他们将编程语言视为自然语言，并通过学习大型代码库(如GitHub)来训练代码完成系统。特别是神经语言模型，如递归神经网络(RNNs)，可以捕捉序列分布和深层语义，因此变得非常流行。但这些标准的神经语言模型受到所谓的隐藏状态瓶颈的限制：所有关于当前序列的信息被压缩成一个固定大小的向量。这种限制使得RNNs难以处理长期依赖关系，这种关系在程序源代码中很常见，例如类标识符在使用前声明了很多行。

注意力机制(Bahdanau et al.,2014)为这一挑战提供了一种解决方案。通过注意力机制，神经语言模型学会了检索和利用之前的相关隐藏状态，从而提高了模型的记忆能力，为反向传播提供了更多的路径。为了处理代码补全中的长期依赖关系，我们开发了一个定制的注意力机制，该机制可以利用程序抽象语法树(AST，如图1所示)上的结构信息，抽象语法树将在后面进行描述。

但是，即使有了注意力机制，还有另一个关键的问题被称为单词未知问题。一般来说，神经语言模型的最后一个组成部分是一个softmax分类器，每个输出维数对应于预定义词汇表中唯一的单词。由于计算高维softmax的计算成本很高，一种常见的做法是使用语料中最常见的K个单词构建词汇表，并用一个特殊的单词替换其他OoV词，如UNK。直观地说，标准的基于softmax的神经语言模型不能正确地预测OoV词。在代码补全任务中，简单地推荐一个UNK标识对开发人员没有帮助。未知词问题限制了神经语言模型的性能，特别是当程序源代码等语料库中存在大量的唯一词时。

对于我们的代码补全任务，我们注意到在编写程序时，开发人员倾向于局部重复。例如，图1中的变量名my_salary相对于整个语料库可能很少见，被标记为UNK。但是在特定的代码块中，它重复了几次，并且频率相对较高。直观地说，在预测这些未知单词时，我们的模型可以学会在局部上下文中选择一个偶合点，然后复制该位置的单词作为我们的预测。实际上，最近提出的指针网络，通过使用注意力分数从输入序列中选择一个单词作为输出，可以做到这一点[Vinyals et al.,2015]。虽然指针网络可以更好地预测未知单词或稀有单词，但它们无法预测超出当前输入序列的单词，即缺乏全局视图。因此，它们在我们的代码补全任务中可能不能很好地工作。

为了提高代码补全的效率，本文提出了一种指针混合网络，它可以通过从全局词汇表生成一个单词或从局部上下文复制一个单词来预测下一个单词。对于前者，我们采用一个标准的带注意力机制的RNN，我们称之为全局RNN组件。对于后者，我们使用一个指针网络，我们称之为局部指针组件。实际上，这两个组件共享相同的RNN结构和注意力分数。我们的指针混合网络是两个组件的加权组合。在进行每次预测时，基于上下文信息学习一个切换器，它可以指导模型选择一个组件来生成下一个单词。通过这种方式，我们的模型学习何时何地从局部上下文中复制一个OoV词作为最终的预测。

本文的主要贡献如下：

• 我们提出了一种指针混合网络，在代码补全任务中来更好地预测OoV词，它学习从全局词汇或局部上下文生成下一个词。

• 我们开发了一个用于代码补全的注意力机制，它利用了抽象语法树（即AST）的结构信息（特别是父-子信息）。

• 我们在两个样本标记数据集(JavaScript和Python)上评估我们的模型。实验结果表明，在目前的水平上有很大的改进。

2 方法

2.1 项目展示

在我们的语料库中，每个程序都以抽象语法树(AST)的形式表示。任何编程语言都有明确的上下文无关语法，可以用来将源代码解析为AST，而且AST还能以一对一对应的方式转换回源代码。以AST形式处理程序是软件工程(SE)的一种典型实践[Mou等,2016;Li等,2017]。

图1显示了一个Python示例程序及其对应的AST，我们可以看到每个AST节点都包含两个属性：节点的类型和一个可选值。对于每个叶子节点，“:”被用作类型和值之间的分隔符。对于每个非叶子节点，我们附加一个特殊的EMPTY标记作为其值。例如，考虑图1中的AST节点NameLoad:my_salary，其中NameLoad表示类型，而my_salary是值。唯一类型的数量相对较少(在我们的语料库中有数百种)，使用这些类型编码程序结构，如Indentifier、IfStatement、SwitchStatement等。而编码程序文本的值则有无限的可能。一个值可以是任何程序标识符(如jQuery)、文字(如66)、程序操作符(如+、-、*)等等。

将程序表示为AST，而不是纯文本，使我们能够预测程序的结构，即每个AST节点的类型。再次看到图1中的示例，当下一个标识是关键字for时，对应的下一个AST节点是For(:EMPTY)，它对应以下代码块：

for _ in _ :

  # # for loop body

这样的话，下一个AST节点的成功预测不仅补全了下一个标识，还补全了整个代码块，包括一些微不足道的标识，如in和“:”。这种结构补全可以在不同粒度级别上实现更灵活的代码补全。

为了应用统计序列模型，我们通过深度优先遍历将每个AST转化为节点序列。为了确保序列能被转换回原始的树结构，从而转换回源代码，我们允许每个节点类型对关于AST节点是否有子节点和右兄弟节点的两位额外信息进行编码。定义一个单词w_i = (T_i, V_i)，表示一个AST节点，T_i为类型，V_i为值，则每个程序都可以表示为一个单词序列wⁿ_i=1。因此，我们的代码补全问题被定义为：给定单词w₁，w₂.....，w_t-1序列，我们的任务是预测下一个单词w_t。显然，我们有两种任务：预测下一个节点类型T_t和预测下一个节点的值V_t。我们为每个任务建立一个模型，分别训练它们。我们称之为基于AST的代码补全。

2.2 神经语言模型

代码补全任务可以看作是一个语言建模问题，其中递归神经网络(RNNs)近年来取得了令人瞩目的成功。LSTM [Hochreiter和Schmidhuber, 1997]被提出利用门控机制来缓解RNNs中的梯度消失或梯度爆炸问题。定义一个标准的LSTM单元为h_t = f(x_t,h_t-1)，在每个时间步长t， LSTM单元以当前输入向量x_t和之前的隐藏状态h_t-1作为输入，生成当前隐藏状态h_t，用于在时间步长t内进行预测的计算。

2.3 注意力机制

标准神经语言模型存在隐藏状态瓶颈的问题[Cheng et al.,2016]。为了缓解这一问题，本文提出了一种能够检索和利用相关前隐状态的注意力机制。它被合并到标准的LSTM中，我们在这里称之为注意力LSTM，如图2所示。

上下文注意力机制（Context Attention）：传统的注意力机制利用上下文窗口中的先前隐藏状态[Bahdanau et al.,2014]，我们称之为上下文注意力机制。形式上，我们保留了L个先前隐藏状态的外部记忆，记为M_t = [h_t-L，…，h_t-1] ∈ R^k*L。在时间步长t内，模型使用注意力层来计算h_t与M_t中隐藏状态之间的关系，用注意力分数α_t表示，然后生成总结上下文向量c_t。我们为代码补全任务设计的上下文注意力机制如下：

父级注意力机制（Parent Attention）：除了传统的上下文注意力机制之外，我们还提出了基于AST的代码补全的父级注意力机制。直观上，上下文窗口中的不同隐藏状态应该与当前预测有不同程度的相关性。由于我们的序列是从树（即AST，参见图1）转化而来的，因此父节点应该与子节点有很大的相关性。但是扁平的AST已经失去了父子关系的信息。为了利用这种结构信息，在扁平化AST时，我们记录每个AST节点的父位置pl，即在它之前有多少节点。然后在时间步长t内，我们的模型从外部内存M_t中检索父向量p_t，它是父位置的隐藏状态，即h_t−pl。父代码段的信息可以使我们的模型更有信心地进行预测。

在时间步长t内预测下一个单词时，我们不仅根据当前隐藏状态h_t，还根据上下文向量c_t和父向量p_t来确定决策。输出向量G_t编码下一个标识的信息，然后投射到词汇表空间，再用softmax函数生成最终的概率分布y_t∈R^V：

其中W^g∈R^k*3k和W^v∈R^V*k是两个可训练投影矩阵，b^v∈R^V是一个可训练偏差向量。需要注意的是，V表示词汇表的大小，“;”代表连接操作。

2.4 指针混合网络

受到程序源代码中普遍存在局部重复标识的启发，我们提出在代码补全任务中，利用指针网络来预测OoV词，方法是从前面的输入序列中复制一个标识。具体地说，我们提出了一种结合了标准RNN和指针网络的指针混合网络，如图3所示。

我们的指针混合网络由两个主要组件(全局RNN组件和局部指针组件)和一个切换器来平衡它们。对于全局RNN组件，它是一个注意力LSTM，可以从预定义的全局词汇表预测下一个标识。对于局部指针组件，它根据已知的位置权重在局部上下文中指向以前的位置。我们的指针混合网络通过连接两个组件的输出向量来组合这两个组件。在连接之前，两个单独的输出由一个基于上下文的学习切换器进行缩放，因此我们的模型学习如何在每次预测中选择一个特定的组件。具体来说，切换器生成标量s_t∈[0,1]，它表示使用全局RNN组件的概率，然后1- s_t是使用局部指针组件的概率。

在连接两个缩放后的向量后，我们选择概率最高的一个输出维数。如果此维度属于RNN组件，则从全局词汇表生成下一个标识；否则，将从局部上下文中复制下一个标识。

形式上，在时间步长t内，全局RNN组件根据公式5，为词汇表内的下一个标识x_t生成一个概率分布小波变换w_t∈R^V。局部指针组件根据分布l_t∈R^L，指向内存内的位置，其中L是内存的长度。为了减少参数，加快训练速度，我们在实践中重新使用了注意力分数（即公式2中的l_t）。

切换器是一个基于当前隐藏状态h_t和上下文向量c_t的sigmoid函数：

其中W^s∈R^2k*1和b^s∈R¹是可训练参数，s_t∈[0,1]是平衡w_t和l_t的标量。最后，模型将两种分布结合起来，得到最终的预测结果：

3 评估

3.1 数据集

我们在这两个基准数据集上评估不同的方法：JavaScript (JS)和Python (PY)，表2总结了这两个数据集。这两个数据集都来自GitHub，都是公开的，并在之前的工作中使用过[Bielik et al.,2016;Raychev等人,2016;Liu等,2016]。这两个数据集都包含150,000个程序文件，它们以各自对应的AST格式存储，前100,000个用于训练，其余50,000个用于测试。在按顺序对每个AST进行深度优先遍历之后，我们生成了用于训练和评估的多个查询，每个AST节点一个查询，通过采用从序列中删除某节点(加上右边的所有节点)，然后尝试预测该节点的方法。

JS和PY中唯一节点类型的数量最初分别为44和181。通过添加第2.1节中所讨论的关于孩子和兄弟姐妹的信息，我们可以将数字分别增加到95和330。如表2所示，由于两个数据集中唯一节点值的数量太大，无法直接应用神经语言模型，因此我们在每个训练集中只选择K个最频繁的值来构建全局词汇表，其中K为自由参数。我们还添加了三个特殊值：UNK表示词汇表之外的值，EOF表示每个程序的结束，EMPTY表示非叶子AST节点的值。

3.2 实验设置

配置：我们的基本模型是一个单层的LSTM网络，展开长度为50，隐藏单位大小为1500。为了训练模型，我们使用交叉熵损失函数和带Adam优化器的迷你批处理SGD [Kingma和Ba,2014]。我们将初始学习率设为0.001，在每个时期后乘以0.6衰减。我们划分梯度的标准为5，以防止梯度爆炸。注意力窗口的大小设为50，批大小是128，训练模型时epoch为8。每个实验运行三次，并报告平均结果。

我们将每个程序分成由50个连续的AST节点组成的段，如果最后一个段不够满，则用EOF填充。LSTM隐藏状态和内存状态用h₀、c₀进行初始化，这是两个可训练的向量。如果前一个LSTM段和前一个LSTM段属于同一个程序，则最后一个隐藏状态和内存状态作为初始状态被输入下一个LSTM段。否则，隐藏状态和内存状态将重置为h₀、c₀。我们将h₀, c₀初始化为全零向量，而其他所有变量都以[-0.05,0.05]均匀分布的方式随机初始化。我们使用准确性作为我们的评估标准，即正确预测下一个节点类型和值的比例。

预处理和训练细节：由于每个AST节点由一个类型和一个值组成，为了对节点进行编码并将其放入LSTM中，我们分别为每种类型(300维)和值(1200维)训练一个嵌入向量，然后将两个嵌入连接到一个向量中。由于在两个数据集中唯一类型的数量相对较少，所以在预测下一个AST节点类型时不存在未知的单词问题。因此，我们只应用我们的指针混合网络预测下一个AST节点的值。

对于每个数据集，我们用训练集中K个频率最高的值，构建AST节点值的全局词汇表，并将训练集中和测试集中词汇表外的所有节点值标记为OoV值。在训练之前，如果注意力窗口中出现的OoV值与之前的另一个值完全相同，则我们将该OoV值标注为注意力窗口中对应的位置。否则，OoV值被标记为UNK。如果注意窗口中有多个匹配，我们选择位置标签作为匹配值在窗口中的最后一次出现，它是最近的一次。对于词汇表内的值，我们将它们标记为全局词汇表中相应的id。在训练期间，只要训练查询的真实数据集是UNK，我们就将该查询的loss函数设置为0，这样我们的模型就不去学预测UNK。在训练和评估中，所有目标值为UNK的预测都被认为是错误的预测，降低了整体的准确率。

3.3 实验结果

对于每个实验，我们运行以下模型（在第2节中有介绍）进行比较：

• 标准的LSTM：没有任何注意力或指针机制的标准LSTM网络。

• 带有注意力机制的LSTM：一个带有注意力机制(上下文和父母信息)的LSTM网络，在每个时间步长上持续50个隐藏状态。

• 指针混合网络：我们提出的混合网络，它结合了上述带有注意力机制的LSTM和指针网络。

OvV词的预测：我们首先评估我们的指针混合网络，在预测下一个AST节点值时，缓解未知字问题的能力。对于这两个数据集，我们通过将节点值的全局词汇表大小K设为1k、10k和50k，来创建三个特定的数据集，从而导致不同的词汇表外(OoV)率。我们还测量了OoV值在前一个上下文窗口中出现的频率，从而被标记为相应的位置。我们将此度量称为局部化，这是我们期望从指针组件获得的性能增益的上限。我们在每个特定数据集上运行上述模型。表1列出了相应的统计数据和实验结果。

如表1所示，在每个特定数据集上，标准的LSTM的准确度最低，而带有注意力机制的LSTM较前者有所提高，我们的指针混合网络达到了最高的准确度。此外，可以看到，在JS或PY数据集中增加词汇量，OoV率会降低，而不同模型的总体准确率会因可用信息的增加而提高。我们还注意到，与带有注意力机制的LSTM相比，我们的指针混合网络获得了性能提升，并且在词汇表大小为1k时，该提升是最大的。我们将这种性能增益归因于通过局部指针组件正确预测某些OoV值。结果表明，我们的指针混合网络在预测OoV值方面是有效的，特别是在词汇量小、OoV率大的情况下。

与最先进技术的比较：由于已经有研究对这两个基准数据集进行了代码补全，为了验证我们提出的方法的有效性，我们需要将它们与最先进的技术进行比较。特别是，Liu等人[2016]在JS数据集上采用了标准的LSTM，没有注意力或指针机制。Raychev等人[2016]基于概率语法为代码构建了概率模型，并在两个数据集上实现了最高的代码补全准确度。

具体来说，我们分别对下一个AST节点类型预测和下一个AST节点值预测进行了实验。对于前者，由于类型词汇量较小，不存在单词未知的问题，所以我们只使用了标准LSTM和带有注意力机制的LSTM。对于后者，我们将值词汇大小设置为50k，使结果能够与[Liu et al.,2016]中的工作可做比较，并运用三种模型进行实验。结果如表3所示。

表3的上半部分显示了我们在这项工作中的结果，而下半部分列出了来自先前工作的结果。注意的是，Liu等人[2016]只在JS数据集上应用LSTM，所以他们在PY数据集上没有结果。对于下一个类型的预测，我们的注意力LSTM在两个数据集上都达到了最高的准确度，显著提高了两个数据集的最佳记录。对于JS数据集的下一个值预测，我们的指针混合网络与Raychev等人[2016]的性能相当，Raychev等人的是一种基于领域特定语法的概率模型。但我们的方法优于Liu等人[2016]提出的方法，它也是基于神经网络的。在PY数据集上，我们的指针混合网络预测的下一个值优于以前的最佳记录。因此，我们得出结论：我们的注意力LSTM和指针混合网络可以实现有效的代码补全，达到了四个任务中的三个最优性能。

3.4 讨论

为什么注意力机制能够发挥作用呢？在编写程序时，引用前面声明的变量标识符是很常见的。在本工作中，JS数据集的平均程序长度(即AST节点数)在1000左右，PY数据集的平均程序长度在600左右。因此，在我们的代码补全任务中，我们需要注意力机制来捕获长期依赖关系。此外，我们仅使用上下文注意力来衡量我们提出的父母注意力，对最终预测的影响(见公式4)。从表3可以看出，父母注意力对类型预测有效，但对值的预测的影响很小。

为什么指针混合网络是有效的呢？在合并了指针网络之后，我们通过从局部上下文复制一个值来预测OoV值，而复制的值可能就是正确的预测。因此，我们观察到在我们的指针混合网络的性能增益。然而，有人可能会说，无论指针组件有多强大，只要我们有机会预测OoV值，准确性就一定会提高。

为了验证我们的指针组件的复制能力，我们开发了一个指针随机网络，其中指针分布l_t(见图3)是一个随机分布，而不是重用学习到的注意力分数。我们比较了词汇量大小为1k的JS和PY数据集的值预测。结果如表4所示，其中指针随机网络的准确率低于指针混合网络。这证明了我们的指针混合网络确实学会了何时何地复制一些OoV值。但指针随机网络的性能比带有注意力机制的LSTM更差。我们认为原因在于切换器受到随机噪声的干扰，不能总是选择正确的组件(即RNN组件)，从而影响整体性能。

3.5 用例研究

我们在图4中描述了一个代码补全示例。在本例中，目标预测employee_id是一个相对于整个训练语料库的OoV值。我们展示了每个模型的前5个预测。对于标准LSTM，它只生成EMPTY，这是我们的语料中最常见的节点值。带有注意力机制的LSTM，从上下文中得知目标是UNK的概率较大，但不能产生真实值。指针混合网络在观察前面代码中出现的值后，能够成功地从上下文识别出OoV值。

4 相关工作

统计代码补全：最近有一项研究探索了统计学习和序列模型在代码补全任务中的应用，如n-gram模型[Hindle et al.,2012;Tu等人,2014]，概率语法[Allamanis and Sutton, 2014;Bielik等人,2016年;Raychev等,2016]。最近，神经网络成为建模源代码的一种非常流行的方式[Raychev et al.,2014;White et al.,2015;Allamanis等,2016]。具体来说，Bhoopchand等人[2016]为RNN提出了一种稀疏指针机制，以便在Python源代码中更好地预测标识符。然而，它们的指针组件针对的是Python源代码中的标识符，而不是我们工作中的OoV标记。OoV标记不仅包括标识符，还包括其他类型，如VariableDeclarator。此外，他们直接将每个程序序列化为代码标识序列，而在我们的语料库中，每个程序都表示为AST节点序列，以便更智能地预测代码结构。

神经语言模型：深度学习技术，如RNN已经在语言建模任务中取得了最先进的成果[Mikolov等人，2010]。软注意力或记忆机制[Bahdanau等人，2014;Cheng et al.,2016;Tran等人，2016]提出了缓解标准RNN中梯度消失问题的方法。指针是最近提出的另一种机制[Vinyals et al.,2015]，它赋予RNN“复制”的能力。指针机制在总结[Gu et al.,2016]、神经机器翻译[Luong et al.,2014]、代码生成[Ling et al.,2016]、语言建模[Merity et al.,2016]等任务中都有帮助。

特别地，Gulcehre等[2016]也提出了基于RNN组件和局部指针组件在每个时间步长生成新单词。但是，他们的场景是像神经机器翻译这样的序列到序列的任务，而我们的场景是语言建模。Merity等[2016]也提出了类似的观点，但是它们使用指针组件来有效复制的是全局词汇表中稀有单词的ID，而不是词汇表外的单词。此外，这两篇作品的语料库是自然语言，而我们的语料库是程序源代码。

5 结论

本文在代码补全任务中，运用神经语言模型,并开发一个注意力机制，该机制利用了程序的AST的亲子信息。在代码补全任务中，为了应对OoV值,我们提出一个指针混合网络,学会通过一个RNN组件生成一个新值,或通过一个指针组件从局部上下文复制一个OoV值。实验结果证明了我们的方法的有效性。