【论文精度】An AST Structure Enhanced Decoder for Code Generation

An AST Structure Enhanced Decoder for Code Generation

摘要

​ 目前，最主流的神经网络代码生成模型通常配备一个树状结构的LSTM解码器，它输出一系列action，通过前序遍历构建一棵抽象语法树（AST）。然而这种解码器有两个主要的缺点：

• 除了父节点action，其他遥远而重要的历史actions很少对当前的决策做出贡献
• 忽略了未来的actions，这可能对预测当前action至关重要

​ 为解决这些问题，本文提出了一种用于代码生成的新型AST结构增强解码器，在上述两个方面显著扩展了解码器：

• 引入了一种AST信息增强注意力机制来充分利用历史actions，历史actions的影响根据树上的语法距离、动作类型和相对位置进一步区分；
• 通过多任务学习，联合建模当前动作及其重要的未来动作的预测，并利用后者的隐藏状态进一步改进前者。（利用未来兄弟节点信息，解码器先预测与当前节点最近的兄弟节点，此类信息能很好地补充当前节点对全局结构信息的利用）

1. Introduction

​ 给定一个自然语言描述（NL），代码自动生成旨在自动生成一个可执行代码。最典型的用于代码生成的神经模型认为代码生成是一个序列到抽象语法树(seq2AST)建模问题，具体来说，首先使用编码器来学习输入NL描述的语义表示，然后在编码器的语义指导下，解码器输出一系列语法动作，并通过前序遍历生成相应的AST，最后利用确定性生成工具将生成的AST转换为代码。

​ 与已有模型相比，AST的使用很自然地反映了程序设计语言的递归结构，能够有效地生成良好形式的代码；然而利用AST引入了一个树生成问题，通常比标准序列预测更难。因此，如何有效地利用部分生成的AST的丰富特性，这对于确保后续操作的准确预测至关重要。然而这可能超出了标准RNN解码器的能力，①除了父动作，其他遥远而主要的历史动作会极大地影响当前动作的预测；②标准RNN无法利用重要的未来动作，这通常对当前动作的预测也是至关重要的。因此，这些seq2AST模型中使用的部分AST的句法信息是值得研究的。

​ 本文提出了一种增强的AST结构解码器，能够有效地利用部分AST的丰富特征进行动作预测，它是传统树结构解码器的重要扩展，有两个方面的改进：①为解码器增加一种AST信息增强的注意力机制，根据不同历史动作的语法距离、动作类型和相对位置区分题目对当前动作预测的影响；②选择一个未来动作，为当前动作的预测提供重要的上下文信息，并通过多任务学习联合建模当前动作和未来动作的预测，然后可以利用后者的学习隐藏状态来为前者提供信息。 需要注意的是，未来动作是在推理过程中动态预测的。与传统的解码器相比，本文解码器能够更好地利用有益的历史和未来动作来预测当前动作。

总体来说，本文主要贡献有三个方面：

• 研究如何通过AST信息增强注意力，更好地利用历史动作进行动作预测；
• 建议利用对未来动作的预测来有利于预测当前动作，这于通常采用因果解码的大多数解码过程形成了对比；
• 在常用数据集上的实验结果证明了所提出的解码器的有效性、

2. Related work

​ 本文主要研究利用AST信息进行代码生成，对此，Sun等人提出将transformer改进用于代码生成，增加了一个CNN层来处理AST结构，但是他们不仅忽略了不同历史动作对当前动作预测的影响有显著差异的事实，也忽略了未来动作中包含的上下文，而这些上下文对于当前动作的预测可能至关重要。相比之下，本文提出的解码器能够解决这两个缺陷，从而使神经模型在代码生成方面表现更好。

3. Our Model

​ 本文模型是基于TRANX的，该模型因具有竞争力的性能而被广泛应用，所提出的策略也可以应用于其他seq2AST模型。

3.1 使用ASDL文法建模代码生成

​ 在TRANX中，引入AST作为中间表示，采用三步策略对NL描述到代码的转换过程进行建模。如图1所示，对于每个输入NL描述x，首先应用模型产生一系列基于ASDL文法的动作a = a₁,a₂,…,a_T；然后以前序遍历的方式构造AST z；最后直接调用用户指定的函数AST_to_MR(*)将z转换为目标代码y。

形式上，每个ASDL文法由两个组件组成：类型(type)、构造函数(constructors)，其中，类型定义为复合的或者原子的；每个构造函数使用其字段(field)指定特定类型的语言组件。每个字段都指定了其可以保存值的类型，并包含一个基数，表示它所保存的值的数量。在每个具有复合类型的字段中，在其类型下存在几个构造函数；具有原子类型的字段则直接存储值。

以图1中a2的ASDL语法为例，最左边的复合类型"expr"表示它的构造函数"Tuple(expr* elts)“只能应用于形成一个类型为"expr"的AST节点。当特定于此构造函数时，类型为"expr"且基数为”*"的字段"elts"将被进一步扩展，以形成类型为"expr"的子AST节点；

a3的ASDL文法的字段"id"的类型是原子的，这里使用a4: GenToken[result]组成一个AST节点，直接存储值为"result"的字段"id"。

在每个时间步，考虑三种基于ASDL文法的动作用于构造一个AST节点

• APPLYCONSTR[c]：使用这个动作，构造函数c被应用到与c类型相同的父节点的开口复合字段，使用c的字段填充一个节点，这里父节点的复合字段也被称为前沿字段
• REDUCE：表示可选或多个基数(optional or multiple cardinalities)展开字段的完成
• GENTOKEN[v]：这个动作使用一个token v填充一个原子边界字段

正式地，代码y的生成概率可以分解为对应AST z中的action序列的概率：$p(z|x)=\underset{t=1}{\overset{T}{\Pi }}p(a_t|a_{<t},x)$ ，其中，a_t是第t个action；a_<t 是a_t前面的action序列，使用神经网络预测$p(z|x)$的概率。

3.2 模型框架

​ encoder部分与TRANX是一样的：对于一个NL描述$x=x_1,x_2,\dots ,x_N$，使用一个BiLSTM编码器去学习词级别的隐藏层状态 { h i } \{h_i\} {hi​} ，其中，正向LSTM从左到右读取NL描述，反向LSTM的操作方向相反：

其中，$E(x_i)\in R^{dw}$ 是NL词$x_i$的嵌入。

词级别表示通过将前向和反向隐藏状态连接得到，即$h_i = [\mathop{h_i}\limits ^{\rightarrow}: \mathop{h_i}\limits ^{\leftarrow}] $

​ 本文解码器是TRANX解码器的增强版，它也是一个LSTM网络，在每一个时间步t上，其隐藏状态由$s_t=f_{LSTM}([E(a_{t-1}):{\tilde{s}}_{t-1}:p_t],s_{t-1})$，其中，$E(a_{t-1})$ 是前一个action $a_{t-1}$的嵌入表示；$p_t$是父节点复合字段嵌入向量与父action的解码器隐藏状态向量的连接；临时隐藏状态${\tilde{s}}_t=tanh(W_1[s_t:c_t:h{c}_t:f{c}_t])$ ，其中，$c_t$是由编码器隐藏状态 { h i } i = 1 N \{h_{i}\}_{i=1}^N {hi​}i=1N​ 引导的注意力上下文向量 ,$h{c}_t$ 代表通过注意力机制来自历史actions的上下文信息，$f{c}_t$表示从预测模块中引导出的未来action的上下文，$W_1$是一个参数矩阵。

接下来，与TRANX中一样，根据父节点的类型去计算action a_t的概率：

• Composite：采用APPLYCONSTR来扩展字段、REDUCE来结束字段。

  $p(a_t=APPLYCONSTR[c]/REDUCE|a_{<t},x)=softmax(E(c{)}^T{W}_2{\tilde{s}}_t)$

  $E[c]$ 表示构造函数c的嵌入表示。

• primitive：采用GENTOKEN操作来生成一个token v，它从词汇表中生成或者从输入NL描述中复制。因此，

3.2.1 AST Information Enhanced Attention Mechanism

与传统仅考虑父action和先前action的TRANX解码器不同，本文decoder使用注意力机制充分利用所有历史action，如图2所示，

当预测时间步t₅的action时，TRANX decoder只关注t₂和t₄的历史动作，本文decoder能够利用t₁到t₄的历史动作。

从标准注意机制延伸，本文注意从几个方面有效区分历史动作对当前动作预测的影响：

• Syntactic distance $d_{tj}$ ：历史action $a_j$对当前action$a_t$的影响主要取决于他们在AST上的语法距离$d_{tj}$
• Action type $p_j$：如果历史action $a_j$是一个APPLYCONSTR或REDUCE action，它包含AST的结构信息；否则它主要包含AST的语义细节。因此，引入两种类型标记来进一步区分历史动作的影响: APPLYCONSTR/REDUCE 和 GENTOKEN，其中REDUCE被视为一种特殊的APPLYCONSTR
• Relative position $r_{tj}$： 如果历史action $a_j$是当前action$a_t$的祖先或者兄弟节点，它往往对当前动作的预测具有更大的影响，为此，在本文注意机制中将历史动作分为两类：密切相关的动作和其他动作。如图1所示，$a_5$的密切动作主要包含$a_1,a_2,a_3$.

正式地，从历史action得到的上下文向量$h{c}_t$ 为：

其中，$s_j$是第j个decoder隐藏状态，$M_{t,j}$是一个非递增函数，将语法距离$d_{t,j}$转换为[0,1]之间的某一个值；具体地，$M_{t,j}$被定义为：

其中，m是一个控制柔软度的超参数，$S_m$是最大范围。因此，可以动态地控制注意力操作的范围，以减少噪声上下文的负面影响。例如，GENTOKEN动作的范围不需要大跨度的上下文，因此GENTOKEN动作的$z_t$通常较小；相比之下，APPLYCONSTR的$z_t$值通常很大，因为它需要结构信息。（在后续实验部分的案例研究中验证这一假设）

3.2.2 Prediction Module for Important Future Action

以图2(a)为例，对于t₅时间步的action预测，如果decoder可以提前预测其右侧的action $a_7$为一个REDUCE操作，那么它可以更确定当前action为一个APPLYCONSTR action。为了对这个信息进行建模，本文通过多任务学习对当前action和其重要的未来action的预测进行联合建模，然后利用学习到的隐藏状态来预测重要的未来动作，以细化对当前action的预测。

一个关键问题是如何在时间步t确定重要的未来action $a_t^f$ ，本文根据以下标准来选择 $a_t^f$ 。对于每个时间步的action预测，如果它有一个右兄弟节点，则将这个兄弟节点设置为重要的未来action，否则将其重要的未来动作设置为与其父动作相同。以图1为例，$a_5$的主要未来动作是$a_7$，而$a_6$因为没有右侧的兄弟动作，所以它的重要未来动作也是$a_7$；唯一的两个例外包括（使用一个特殊标志去表示他们的重要未来动作）：

• 动作$a_1$为AST的根结点
• REDUCE action

预测重要的未来动作 $a_t^f$ ：

其中，$f{p}_t$ 编码未来动作的前沿字段信息、$c_t^{‘}$ 是一个注意力上下文向量；$\tilde{f}{s}_t$ 用于细化当前action的预测，具体地，使用一个基于元素的门控操作⊙去生成重要未来action的向量表示$f{c}_t$

3.2.3 Model Training

给定训练数据集$D=(x,a)$，使用以下目标函数训练模型：

其中，$a_t^f$ 表示时间步t的重要的未来动作，λ是一个超参数

4. Experiments

4.1 Setup

使用四个benchmark 数据集进行实验：

• DJANGO：总共包含18805行从Django Web框架提取的Python源码，每一行都配对有一个NL描述。总共数据集中代码主要包含不同的现实场景的Python用例，包括I/O操作、字符串变换、异常处理等
• CONALA：包含2879个人工标注的NL问题及其在STACK OVERFLOW上的Python代码；与DJANGO相比，CONALA样例中的目标语义表示覆盖范围更广、组合度更高、使用难度明显增加
• ATIS：被广泛用于语义解析，它包含5410个航班信息查询
• GEN：包含880个美国地理问题，其中每个示例的输入是一个NL描述，输出是一小段与ATIS类似的lambda逻辑形式的代码

使用NLT对所有NL描述进行分词，特别地，为DJANGO执行简单的规范化，例如用占位符替换输入的带引号的字符串。

除了TRANX之外，本文还将得到的模型（ASED）与以下基线进行比较：

• SEQ2TREE：第一个用于代码生成的seq2tree模型，通过对输入NL语句的编码向量进行调节输出树生成逻辑代码
• LPN：这个模型允许同时选择条件上下文和生成的粒度，例如characters或token
• ASN：
• YN17
• COARSE2FINE
• SZM19
• TREEGEN

使用与TRANX相同的配置，具体地，设置action和field嵌入大小为128，隐藏层维度为256，dropout为0.5（CONALA数据集上为0.3）。使用精确的匹配精度作为评估标准（在CONALA数据集上使用corpus-level BLEU作为评估标准）

4.2 Effect of λ

​ λ为一个很重要的超参数，平衡了传统损失和未来预测损失。为研究其影响，使用不同的λ在验证集上进行实验，结果如图3所示，本文模型在λ=0.6时在所有验证集上实现了最佳性能，因此选用λ=0.6进行后续实验。

4.3 Main Results

Table II 给出了实验结果，

4.4 Performance by the Lengths of Action Sequences

根据action序列的长度分别将两个相对较大的数据集(DJANGO和ATIS)分成不同的组，并分别在图4(a)和(b)中展示了模型的平均性能。可以看到，在大多数组中，ASED总是取得更好的性能，这些结果进一步证明了模型的通用性。

4.5 Ablation Study（消融实验）

为了研究不同组件对模型的影响，提供了模型在DJANGO和ATIS上的简化版模型的性能：

• ASED(AHA)：只配备了所提出的作用于历史动作的注意机制。进一步研究其四种变体：ASED(AHA-M), ASED(AHA-SD), ASED(AHAAT) and ASED(AHA-RP)，其中$M_{tj}$函数、语法距离、action类型、历史actiond的相对位置分别被移除
• ASED(PFA)：简化模型仅包含所提出的未来重要action的预测模块。有一个变体：ASED(PFAML) ，其中$\tilde{f}{s}_t$ 被移除

从Table III可看出，删除任何模块都会导致模型性能下降。实验结果表明，这些构件在利用AST结构信息方面是有效的。此外，四种ASED(AHA)的变体出现了进一步的性能下降，尤其是ASED(AHA-RP)的性能下降幅度最大，这表明相对位置是所提出的注意力机制中最重要的特征。最后，ASED(PFAML)的性能要劣于ASED(PFA)，这表明与传统的多任务学习方法相比，所提出方法能够更有效地挖掘未来上下文。

4.6 Case Stud

图5和图6显示了DJANGO数据集产生的两个action示例。

• 例1中，可视化了AST信息增强注意力机制的注意力权重。可以观察到，所提出模型在每个时间步的动作预测中总是更多地关注父或者兄弟历史动作，特别地，在注意力范围的动态建模的影响下，一些较远的历史动作会被直接忽略。

• 例2中，比较了不同模型产生的最佳action序列。可以发现TRANX在第2步的动作预测是错误的，而本文模型在第2步的动作预测是正确的，这是因为TRANX首先正确地预测了第4步的action，从而帮助模型预测"type"字段的值。

5. Conclusion

​ 本文提出了一个AST结构增强的decoder用于代码生成，在两个方面显著扩展了TRANX的decoder：

• 引入了一个AST信息增强的注意力以充分利用历史actions
• 通过多任务学习，为decoder配备了一个预测重要未来动作的模块，这使得decoder能够在预测当前action期间利用未来action

​ 与以往的模型相比，本文的增强模型能够更好地利用历史和未来工作来进行代码生成，在几个数据集上的实验结果验证了所提出的decoder的有效性。

thinking

这篇文章是在 TRANX 上进行了一些改良，他们聚焦于抽象语法树的生成步骤，模型中他们采用了双向 LSTM 作为 Encoder，这与 TRANX 类似；相较于 TRANX ，他们的 Decoder 在构成AST时多考虑了 hc_t 和 fc_t，分别代表注意力机制中从历史行为得到的上下文信息、从他们的预测模型中获取的关于将来行为的上下文信息。