【论文阅读 - AAAI 2020】TreeGen: A Tree-Based Transformer Architecture for Code Generation

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 doublepower.mou@gmail.com

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$表示在规则序列当中的第$i$个规则。这样，我们的任务是训练一个计算$p(r_i|NLinput,p_i)$的模型，也就是说，给定自然语言描述和当前生成的部分AST，模型计算规则扩展该节点的概率。

NL Reader

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

输入文本的表示

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

其中$n_i$