TreeGen: A Tree-Based Transformer Architecture for Code Generation

Conference: AAAI 2020

Autors:

Zeyu Sun,† Qihao Zhu,† Yingfei Xiong,∗† Yican Sun,† Lili Mou,‡ Lu Zhang† †Key Laboratory of High Confidence Software Technologies (Peking University), MoE; Software Institute, Peking University, 100871, P. R. China
{szy , zhuqh, xiongyf, sycpku, zhanglucs}@pku.edu.cn
‡University of Alberta, Edmonton, AB, Canada [email protected]

Link: https://arxiv.org/abs/1911.09983

摘要

目前关于代码生成的最新研究是基于神经网络的生成。

但目前的研究仍然面临两个问题：

• 长依赖问题；如变量的定义和引用往往会相距较远
• 模型结构问题；程序往往会包含丰富的结构信息。

这篇文章提出TreeGen：

• 使用Transformer的注意力机制来缓解长依赖问题
• 使用AST reader(encoder) 来融合语法规则和AST结构信息

Evaluation:

• Python数据集：HearthStone
• 语义解析：ATIS & GEO

简介

代码生成就是根据一个自然语言描述作为输入，生成特定的可执行程序的任务。

随着深度学习的发展，研究人员也将各种网络结构应用到了这个问题上，比如Seq2Seq和Seq2Tree模型。目前SOTA方法是通过预测语法规则的序列的生成代码，这种方法保留了部分AST的信息，以此来预测语法规则并扩展特定的节点。

但语法规则的分类面临着两个挑战：

• 长依赖问题
• 代码结构的表示问题，“flat”的神经网络结构很难比较好的捕捉模型的结构信息

这篇文章使用Transformer来解决第一个挑战，但是原始的Transformer并不是为了处理代码程序而设计，也不能很好的利用树形结构，也就是上面提到的第二个挑战。在基于图和树的卷积神经网络中，利用结构信息的标准方法是将节点及其结构邻居的向量表示组合起来作为结构卷积子层的输出。然而，标准的Transformer架构没有这样的结构卷积子层，也不清楚在哪里添加它们。

本文的核心猜想是，当对一个节点及其结构邻居进行卷积时，向量表示应该主要包含来自原始节点的信息。在Transformer的decoder中，由于节点的矢量表示被更多的块处理，它们逐渐地混合了来自其他节点的更多信息，从而丢失了原来的信息。因此，结构卷积子层只对前几个Transformer的decoder块而不是全部。

TreeGen可以分为3部分：

• 自然语言的编码（NL reader）
• AST的编码（AST reader）
• 解码器（decider）

模型

如上所属，TreeGen分为了3各部分，整体的结构如下图所示：

图2: 模型的整体结构

Grammar Rule Prediction

在这个小节中主要介绍如何将代码生成建模为一系列语法规则的分类问题。

程序可以分解成几个上下文无关的语法规则，并被解析为一个AST。图1中展示了关于代码“length=10”的AST，其中虚线框是终结符，实心框是非终结符。

图1: 代码“length=10”的AST

基于AST的代码生成可以看作是通过语法规则扩展非终结符。并重复这个过程，直到所有叶节点都是终结符。

在图1中，“1:root->Module”是语法规则的一个示例，前面的数字是规则的ID。按照预排序遍历，我们可以获得生成AST的规则序列，如右上角所示。

形式上看，概率可以视为规则按照顺序生成代码的概率。

其中 r i r_i ri​表示在规则序列当中的第 i i i个规则。这样，我们的任务是训练一个计算 p ( r i ∣ N L   i n p u t , p i ) p(r_i | NL\ input,p_i) p(ri​∣NL input,pi​)的模型，也就是说，给定自然语言描述和当前生成的部分AST，模型计算规则扩展该节点的概率。

NL Reader

输入的自然语言用于描述代码的功能，它可以是半结构化的描述，也可以是完全的自然语言描述。

输入文本的表示

使用了Character Embedding方法，因为往往相类似的单词具有相类似的字符，为了利用这种属性，所以将每个字符进行embedding后，再由一个全连接层得到单个NL输入的表示结果，公式如下图所示：

其中 n i n_i ni​表示第i个输入的描述。

NL Reader的网络结构

如图2所示，NL Reader是由 N d N_d Nd​个block堆叠而成，对于每一个block，其中包含了3个不同的sub-layers，分别是：

• self-attention
• gating mechanism
• word convolution

在两个sub-layer之间都添加了残差连接和正则化层（Add & Norm）

AST Reader

本文设计了一个AST Reader来对生成的部分AST的结构进行建模。虽然程序是通过预测语法规则的序列生成的，但是这些规则本身缺乏对程序的具体了解，不足以预测下一个规则。因此，我们的AST Reader考虑异构信息，包括预测的规则和树结构。

为了整合这些特定于程序的信息，本文首先将代码表示为规则序列，然后使用注意机制对规则进行编码，最后使用树卷积层将每个节点的编码表示与其祖先相结合。

AST表示

• Rule Sequence Embedding

  • 使用规则的ID表示规则的序列
  • 用实数值表示规则向量，然后使用table-lookup的embedding的方式

• Rule Definition Encoding

  • 上面基于表查找的嵌入将语法规则视为原子标记，并丢失规则内容的信息。

  • 为了解决这个问题，我们用规则定义的编码来增强规则的表示

  • 对于语法规则 i : α → β 1 ⋅ ⋅ ⋅ β K i:α→β_1···β_K i:α→β1​⋅⋅⋅βK​，其中α是父节点， β 1 ⋅ ⋅ ⋅ β K β_1···β_K β1​⋅⋅⋅βK​是子节点，它们可以是其他终结符或非终结符，索引i是规则的ID。

  • 与公式2类似，通过一个全连接层将规则内容编码为向量 r ( c ) r^{(c)} r(c)，输入为各符号的查表嵌入量α、β1、···、βK。值得注意的是，序列也被填充到最大长度。

  • 然后规则定义的feature也通过全连接层进行处理：
    

    • 其中 r i r_i ri​表示使用table-lookup进行embedding方式对规则i进行的编码， r i ( c ) r_i^{(c)} ri(c)​表示对规则内容的编码， α \alpha α表示对父节点的编码。在这之后增加规则化层，完成对规则定义的编码。

• Position and Depth Embeddings

  • 因为使用了self-attention所以要进行位置编码
  • 但这种位置编码无法获取AST中规则的位置，所以结合了一种table-lookup的深度嵌入方法。

Decoder

最后一个组件是一个解码器，它将生成的代码的信息与NL描述相结合，并预测下一个语法规则。与AST Reader类似，解码器中使用一block（总共有 N 2 N_2 N2​个块），每个块有几个子层。在每个子层前后还使用了一个残差连接，然后进行了层标准化.

解码器将要展开的非终端节点作为查询。受先前方法的启发（Sun等人，2019），查询节点表示为从根到待扩展节点的路径。例如，如果我们要展开图1中的节点“Assign”，那么路径应该是root、Module、body、Assign。我们将此路径中的节点表示为实值向量。然后我们对这些向量应用到一个全连接层，如等式2，路径的输出（查询节点）是 q i ( p a t h ) q_i^{(path)} qi(path)​。

然后使用两个attention层分别计算AST reader 和NL reader的输出。

首先在带有查询的AST Reader的输出上应用AST注意层并抽取特征 f 1 ( t r e e ) , ⋯   , f P ( t r e e ) f_1^{(tree)},\cdots, f_P^{(tree)} f1(tree)​,⋯,fP(tree)​。在这一层中，Q是根据查询 q 1 ( p a t h ) , ⋯   , q P ( p a t h ) q_1^{(path)},\cdots, q_P^{(path)} q1(path)​,⋯,qP(path)​，K和V是根据代码特征 y 1 ( a s t ) , ⋯   , y P ( a s t ) y_1^{(ast)},\cdots, y_P^{(ast)} y1(ast)​,⋯,yP(ast)​计算的。进一步整合了输入描述中的特征，这种集成也是通过NL注意来实现的，其中Q是由特性 f 1 ( t r e e ) , ⋯   , f P ( t r e e ) f_1^{(tree)},\cdots, f_P^{(tree)} f1(tree)​,⋯,fP(tree)​计算的；K和V由输入描述 y 1 ( N L ) ) , ⋯   , y L ( N L ) y_1^{(NL))},\cdots, y_L^{(NL)} y1(NL))​,⋯,yL(NL)​计算。

最后，一组两个全连接，其中第一层有一个GELU**函数，用来提取特征用于预测。