TreeGen: A Tree-Based Transformer Architecture for Code Generation

Introduction

Code generation is an important artificial intelligence problem that has the potential to significantly boost the productivity of programmers. Given a specification written in natural language, a code generation system translates the specification into an executable program. For example, if a python programmer gives an instruction “initialize a dictionary, Dict”, the code generator is expected to automatically generates “Dict={ }”.

With the development deep learning techniques, researchers have applied various neural architectures to this problem, such as sequence-to-sequence (Seq2Seq) models or sequence-to-tree (Seq2Tree) models (?; ?; ?; ?; ?; ?). Especially, state-of-the-art approaches generate code by predicting a sequence of grammar rules (?; ?; ?). That is to say, the system keeps a partial abstract syntax tree (AST) of the already-generated code, and predicts the grammar rule to be used to expand a particular node.

In this paper, we propose a novel neural architecture, TreeGen, for the code generation. To address the first challenge, TreeGen adopts the recently proposed Transformer architecture (?), which is capable of capturing long dependencies. However, the original Transformer architecture is not designed for programs, and cannot utilize tree structures, i.e., the second above mentioned challenge. A standard way of utilizing structural information, as in graph- and tree-based convolutional neural networks, is to combine the vector representations of a node and its structural neighbors as the output of a structural convolution sub-layer. However, a standard Transformer architecture does not have such structural convolution sub-layers, and it is not clear where to add them.

It is tempting to add structural convolution sub-layers in all the Transformer blocks. Our core conjecture is that when convolving a node and its structural neighbors, the vector representation should mainly contain the information from the original node. As the vector representation of the nodes is processed by more blocks in the decoder of the Transformer, they gradually mix in more information from other nodes and lose their original information. Therefore, we add the structural convolution sub-layer only to the first several Transformer decoder blocks but not all.

Generally speaking, the TreeGen architecture consists of three parts: (1) a natural language (NL) reader (encoder) encodes the text description; (2) an AST reader (the first several Transformer decoder blocks) encodes the previously generated partial code with the structural convolution sub-layers; (3) a decoder (the rest Transformer decoder blocks) combines the query (the node to be expanded in AST) and the previous two encoders to predict the next grammar rule.

We evaluated our model on an established benchmark dataset for Python code generation, HearthStone (?), which is a Python implementation of a card game HearthStone. The results show that our model significantly outperforms previous models by $4.5$ percentage points. We further evaluated our model on two semantic parsing datasets, ATIS and GEO, which translate natural language sentences into lambda calculus logical forms. The results show that our model has the best accuracy among previous neural models, with $89.1\%$ and $89.6\%$ accuracy, respectively. Our evaluation also shows that adding the structural convolution sub-layer to the first several Transformer blocks significantly outperforms a Transformer with structural convolution in all blocks.

Our Model

We generate code by predicting the grammar rules of the programming language. Figure 2 shows the overall picture of our model, which comprises three parts: an NL reader, an AST reader, and decoder. We introduce them in detail in the following subsections.

AST-based code generation could be thought of as expanding a non-terminal node by a grammar rule. This process is repeated until all leaf nodes are terminal. In Figure 1, “1: root -$>$ Module” is an example of the grammar rules, where the preceding number is the ID of the rules. Following the pre-order traverse, we could obtain the sequence of rules that generate the AST shown in the top right corner.

Formally, the probability can be factorized as the probabilities of the rules generating the code following the order.

where $r_{i}$ is the $i$th rule in the rule sequence. In this way, our task is to train a model to calculate $p(r_{i}\mid\text{NL input},p_{i})$, i.e., given the natural language description and the currently generated partial AST the model calculates the probabilities of the rules to expand this node.

NL Reader

The input description determines the functionality of the code. It can be a semi-structural description as in the HearthStone dataset, or a natural language as in ATIS and GEO semantic parsing datasets.

For an input description, we first tokenize it into tokens $n_{1},n_{2},\cdots,n_{L}$, where $L$ denotes the length of the input. Each token $n_{i}$ is then split to characters $c^{(n_{i})}_{1},c^{(n_{i})}_{2},\cdots,c^{(n_{i})}_{S}$, where $S$ is the number of characters in $n_{i}$. All the tokens and characters are represented as real-valued vectors $\bm{n}_{1},\bm{n}_{2},\cdots,\bm{n}_{L}$ and $\bm{c}^{(n_{i})}_{1},\bm{c}^{(n_{i})}_{2},\cdots,\bm{c}^{(n_{i})}_{S}$ by embeddings.

Character Embedding. It often happens that similar words have similar characters (e.g., “program” and “programs”). To utilize this property, we represent a token by character embeddings with a fully-connected layer

where $W^{\text{(c)}}$ are the weights and the character sequence is padded to a pre-defined maximum length $M$. After the fully-connected layer, we also apply layer normalization (?). These vectors are then fed to the NL reader, and are integrated with the word embeddings by a gating sub-layer.

Neural Structure of NL Reader.

The NL reader is composed of a stack of blocks ($N_{d}$ blocks in total). Each block contains three different sub-layers (namely, self-attention, gating mechanism, and word convolution) to extract features, which we introduce in detail in the following subsections. Between two sub-layers, we employ a residual connection (?) followed by a layer normalization.

Self-Attention. The self-attention sub-layer follows the Transformer’s architecture (?), and uses multi-head attention to capture long dependency information.

For a sequence of input tokens $n_{1},n_{2},\cdots,n_{L}$, we represent them as an embedding $\bm{n}_{1},\bm{n}_{2},\cdots,\bm{n}_{L}$ by a look-up table. We also use position embeddings to encode the information of word positions. In particular, we adopt the variant in ? (?), and compute the position embedding for the $i$th word in the $b$th Transformer block as

where $p_{i,b}[\cdot]$ indexes a dimension of the vector $\bm{p}_{i,b}$, and $d$ is the number of dimensions (i.e., embedding size).

where $H$ denotes the number of heads and $W_{h}$ is the weight. An attention layer is applied in each head $head_{t}$, computed by

where $d_{k}=d/H$ denotes the length of each features vector. $Q$, $K$ and $V$ are computed by

Gating Mechanism. After the features are computed by self-attention, we further incorporate with the information of character embeddings. This is given by a gating mechanism based on softmax. For the $i$th word, we compute a control vector $\bm{q}_{i}$ from $\bm{y}_{i}^{\text{(self)}}$ by a linear transformation. The softmax weight $\bm{k}_{i}^{\text{(c)}}$ for character embedding is given by a linear transformation from $\bm{n}_{i}^{\text{(c)}}$ in Equation 2. The softmax weight $\bm{k}_{i}^{\text{(y)}}$ for Transformer’s output is given by another linear transformation from $\bm{y}_{i}^{\text{(self)}}$. Then, the gate is computed by

They are used to weigh the feature of the Transformer’s layer $\bm{v}_{i}^{\text{(y)}}$ and the feature of character embeddings $\bm{v}_{i}^{\text{(c)}}$, linear transformed from $\bm{y}_{i}^{\text{(self)}}$ and $\bm{n}_{i}^{\text{(c)}}$, respectively.

Similar to Equation 5, the output of our gating mechanism is $Y^{\text{(gate)}}=(\bm{h}_{i,t})_{i,t}$, where $(\cdot)_{i,t}$ represents a block matrix with the elements being $h_{i,t}$.

Word Convolution. Finally, two convolutional layers are applied to the output of the gating mechanism $\bm{y}^{\text{(gate)}}_{1},\cdots,\bm{y}^{\text{(gate)}}_{L}$ and to extract the local features around each token $\bm{y}^{\text{(}\text{conv},l\text{)}}_{1},\cdots,\bm{y}^{\text{(}\text{conv},l\text{)}}_{L}$, where $l$ denotes the layer of convolutional layers. The $\bm{y}^{\text{(conv},l\text{)}}_{i}$ is computed by

In summary, the NL reader has a few Transformer blocks of self-attention, the gating mechanism, and word convolution. The natural language description is encoded as features $\bm{y}_{1}^{\text{(NL)}},\bm{y}_{2}^{\text{(NL)}},\cdots,\bm{y}_{L}^{\text{(NL)}}$.

AST Reader

We design an AST reader to model the structure of the partial AST that has generated. Although our programs are generated by predicting the sequence of grammar rules, these rules alone lack a concrete picture of the program and are insufficient for predicting the next rule. Therefore, our AST reader considers heterogeneous information, including the predicted rules and the tree structures.

To incorporate such program-specific information, we first represent the code as a sequence of rules, then encode the rules with attention mechanism, and finally use a tree convolution layer to combine the encoded representation of each node with its ancestors.

Rule Sequence Embedding. To encode the rule information, we use the ID of the rules. Suppose we have a sequence of rules $r_{1},r_{2},\cdots,r_{P}$ that are been used to generate the partial AST in a decoding step, where $P$ denotes the length of the sequence. We represent these rules as real-valued vectors $\bm{r}_{1},\bm{r}_{2},\cdots,\bm{r}_{P}$ by table-lookup embeddings.

Rule Definition Encoding. The above table-lookup embedding treats a grammar rule as an atomic token, and loses the information of the rule’s content.

To alleviate this problem, we enhance the representation of a rule with the encoding of rule definition.

For a grammar rule $i:\alpha\rightarrow\beta_{1}\cdots\beta_{K}$, where $\alpha$ is the parent node and $\beta_{1}\cdots\beta_{K}$ are child nodes. They can be either terminal or non-terminal symbols. The index $i$ is the ID of the rule.

Similar to Equation 2, we encode the rule content as a vector $\bm{r}^{\text{(c)}}$ by a fully-connected layer with input being the table-lookup embeddings $\bm{\alpha},\bm{\beta}_{1},\cdots,\bm{\beta}_{K}$ of respective symbols. It is noted that the sequence is also padded to a maximum length.

Then the rule definition features $\bm{y}_{1}^{\text{(rule)}},\cdots,\bm{y}^{\text{(rule)}}_{P}$ are computed by another fully-connected layer as

where $\bm{r}_{i}$ is the table-lookup embedding of the rule $r_{i}$, $\bm{r}_{i}^{\text{(c)}}$ is the content-encoding rule representation, and we emphasize the parent node $\alpha$ again. After that, a layer normalization is followed.

Position and Depth Embeddings. Since our AST reader would use self-attention mechanisms, we need to represent the position where a grammar rule is used.

We first adopt the position embedding as in Equation 4, representing when a rule is used in the sequence $r_{1},\cdots,r_{P}$. The position embeddings are denoted by $\bm{p}^{\text{(r)}}_{1}\cdots,\bm{p}^{\text{(r)}}_{P}$

However, such position embedding does not capture the position of a rule in the AST. We further encode such information by a depth embedding. If we expand a symbol $\alpha$ by the rule $r:\alpha\rightarrow\beta_{1}\cdots\beta_{K}$, we represent the depth of the rule by its parent node, i.e., $\alpha$. In this way, we associate another sequence of table-lookup depth embeddings $\bm{d}_{1},\cdots,\bm{d}_{P}$ to the sequence of used grammar rules $\bm{r}_{1},\cdots,\bm{r}_{P}$.

Neural Structure of AST Reader.

The AST reader is also composed of a stack of blocks ($N_{1}$ blocks in total). Each block is decomposed into four sub-layers (namely, self-attention, a gating mechanism, NL attention, and tree convolution). We employ a residual connection around each sub-layer except the layer of tree convolution. After each sub-layer, we apply a layer normalization.

Self-Attention. To capture the information of AST, we build a Transformer-like self-attention layer, where the input is sum of the rule embedding, position embedding, and depth embedding, i.e., $\bm{r}_{i}+\bm{d}_{i}+\bm{p}^{\text{(r)}}_{i}$. The self-attention sub-layer extract features $\bm{y}_{1}^{\text{(ast-self)}},\bm{y}_{2}^{\text{(ast-self)}},\cdots,\bm{y}_{P}^{\text{(ast-self)}}$ of AST input, using the same mechanism as Equations 4, 5, 6 with different weights but add an additional depth embedding to $\bm{p}^{\text{(r)}}_{i}$.

Gating Mechanism. We would like to incorporate the content-encoding rule representation $\bm{y}_{i}^{\text{(rule)}}$ into the Transformer-extracted features. We adopt a gating mechanism as in Equations 8, 9, and the fused features becomes $\bm{y}_{1}^{\text{(ast-g)}},\bm{y}_{2}^{\text{(ast-g)}},\cdots,\bm{y}_{P}^{\text{(ast-g)}}$ after this sub-layer.

NL Attention. During the decoding step, we should be informed of the input NL description. This is given by a multi-head NL attention, similar to the Transformer decoder’s attention to its encoder (?). The extracted features are denoted by$\bm{y}_{1}^{\text{(ast-nl)}},\bm{y}_{2}^{\text{(ast-nl)}},\cdots,\bm{y}_{P}^{\text{(ast-nl)}}$.

Tree Convolution. Should we consider only the above sub-layers, it would be hard for the reader to combine the information of a node with its ancestors. A node can be far away from its ancestors in the rule sequence but is close in structure. Therefore, it is difficult for a traditional Transformer to extract such structural features.

We integrate the features of a node with those of its ancestors. We treat the AST as a graph and use an adjacency matrix $M$ to represent the directed graph. If a node $\alpha_{i}$ is the parent of $\alpha_{j}$, then $M_{ji}=1$. Suppose all the nodes are presented by features $\bm{f}_{1},\cdots,\bm{f}_{n}$, their parents’ features can be given by the multiplication with the adjacency matrix:

where $\bm{f}^{\text{(par)}}_{i}$ denotes the parent of the $i$th node. For the father of the root node, we pad it with the feature vector of the root node itself.

The tree-based convolution window, applied to the current sub-tree, is given by

In summary, the AST reader has $N_{1}$ blocks of these four sub-layers, and yields the features $\bm{y}_{1}^{\text{(ast)}},\bm{y}_{2}^{\text{(ast)}},\cdots,\bm{y}_{P}^{\text{(ast)}}$.

Decoder

Our final component is a decoder that integrates the information of the generated code with the NL description, and predicts the next grammar rule. Similar to the AST reader, a stack of blocks ($N_{2}$ blocks in total) each with several sub-layers is used in the decoder as follows. A residual connection is also employed around each sub-layer followed by a layer normalization.

The decoder takes the non-terminal node to be expanded as a query. Inspired by a previous approach (?), the querying node is represented as a path from the root to the node to be expanded. For example, if we are going to expand node “Assign” in Figure 1, the path should be root, Module, body, Assign. We represent the nodes in this path as real-valued vectors. Then we apply a fully-connected layer like Equation 2 to these vectors and the output of the path (querying node) is $\bm{q}^{\text{(path)}}_{i}$.

We then apply two attention layers to integrate the outputs of the AST reader and the NL reader.

We first apply an AST attention layer over the output of the AST reader with queries and extract features $\bm{f}^{\text{(tree)}}_{1},\cdots,\bm{f}^{\text{(tree)}}_{P}$. In this layer, $Q$ is computed from queries $\bm{q}^{\text{(path)}}_{1},\cdots,\bm{q}^{\text{(path)}}_{P}$; $K$ and $V$ are computed from the code features $\bm{y}_{1}^{\text{(ast)}},\cdots,\bm{y}_{P}^{\text{(ast)}}$. We further integrate the features from the input description. This integration is also implemented with an NL attention, where $Q$ is computed by feature $\bm{f}^{\text{(tree)}}_{1},\cdots,\bm{f}^{\text{(tree)}}_{P}$; and $K$ and $V$ are computed by the input description $\bm{y}_{1}^{\text{(NL)}},\cdots,\bm{y}_{L}^{\text{(NL)}}$.

Finally, a set of two fully-connected layers, where the first layer has a $GELU$ activation function, are followed to extract features for prediction.

Training and Inference

We predict the next grammar rule, among all possible candidates, by softmax based on the decoder’s last layer features.

We also introduce the pointer network (?) (essentially, an attention) that can directly copy a token $a$ from the NL description. In this case, the resulting grammar rule is $\alpha\rightarrow a$, where $\alpha$ is a non-terminal symbol to be expanded and $a$ is a terminal symbol. Such pointer mechanism is helpful for user-defined identifiers (e.g., variable and function names).

The choice between softmax rule prediction and the pointer network is given by another gating mechanism $p_{g}$, also computed from the decoder’s last feature. The overall predicted probability of the next grammar rule is

where $i$ denotes the ID of the rule, $\mathbf{D}$ is the set of predefined rules, and $\mathbf{C}$ denotes the set of rules in the form of $\alpha\rightarrow a$, where $a$ is a terminal token that occurs in the NL description.

$p_{g}$ is the probability of using the type of predefined rules, and the $p(r_{i}|\cdot)$ (the probability of each predefined rules) are computed by two single-layer perceptrons with the sigmoid and softmax activation functions, respectively, and the input of these layers are the features $\bm{h}^{\text{(dec)}}$.

where $\bm{h}^{\text{(dec)}}$ denotes the decoder’s last feature. The model is optimized by maximizing negative log likelihood loss against the reference program.

The inference starts with a start rule, $start:\textit{snode}\longrightarrow\textit{root}$, expanding a special symbol snode to the root symbol. The recursive prediction terminates if every leaf node in the predicted AST is a terminal. During prediction, we use beam search with a size of 5. Invalid rules are excluded during beam search.

Evaluation

We evaluated our approach on two types of benchmarks: (1) a Python code generation benchmark, HearthStone, and (2) two semantic parsing benchmarks, ATIS and GEO.

We first evaluated our approach on the HearthStone benchmark (?). The benchmark contains Python code that implements 665 different cards of HearthStone. Each card is composed of a semi-structural description and a groundtruth Python program. The Python programs have a length of 84 tokens on average. The description comes with several attributes such as card name, card type, as well as a natural language description for the functionality of the card. A Python program is mainly decided by the natural language description where the attributes decide the constants or identifier names. A sample description and its corresponding Python program are shown in Figure 3. When preprocessing the card description into token sequences, existing approaches consider two methods. The first (?; ?) (called plain preprocessing) treats the whole description as plain text and delimit the tokens by standard separators such as space or periods. The second (?) (called structural preprocessing) treats the descriptions as semi-structural and always treat an attribute as one token. In this experiment, we consider both methods and denote the results corresponding to the plain preprocessing as TreeGen-A and that corresponding to the structural preprocessing as TreeGen-B. We followed the train-dev-test split in ? (?), and the statistic is listed in Table 2.

We measured the performance following the metrics in ? (?). We computed the StrAcc, which is the percentage of programs that has exactly the same token sequence as the ground truth; the BLEU score, which is used to measure the similarity between the generated code and the reference code at the token level; and the Acc+, which is evaluated manually, allows variable renaming on top of StrAcc, for every test case.

For neural networks, we set the number of NL reader layers $N_{d}=6$, and $N_{1}=N_{2}=5$ for the AST reader as well as the decoder. The size of all embedding is 256. The hidden sizes were all set to the 256 except each fully-connected layers, except the first layer was 1024 dimensions. We applied dropout after each layer (including attention layers, gating mechanism layers, convolutional layers, and fully-connected layers, where the drop rate is 0.15). The model is optimized by Adafactor (?) with default parameters.

We show the results in Table 1. In this table, the structural preprocessing has a better performance compared with the plain preprocessing.

As shown, our model achieves 6 percentage points accuracy improvement with plain preprocessing and 4.5 percentage points accuracy improvement with structural preprocessing. For the BLEU score, our model also achieves the best results. These boosts in performance indicate that TreeGen successfully alleviates the long dependency problem and effectively encodes the structural information in the code generation.

We further evaluated the complexity of our model on the HearthStone, and the result shows that our model is faster than the previous ones. It takes 18s for an epoch on a single Nvidia Titan XP, whereas 180s for the CNN (?) and 49s for the RNN (?).

One of the keys of our approach is to add the structural convolution sub-layers only to part of the Transformer blocks in the decoder. To evaluate whether this design decision is effective, we evaluate four competing settings: 1) adding the structural convolution sub-layers to all Transformer blocks (i.e., $N_{1}=10$); 2) adding the structural convolution sub-layers to the first 7 blocks in AST reader (i.e., $N_{1}=10(7)$); 3) adding the structural convolution sub-layers to the first 8 blocks in AST reader (i.e., $N_{1}=10(8)$); 4) the other adds to none (i.e., $N_{1}=0$). As we can see, from Table 1 our approach adding the sub-layer to all transformer blocks ($N_{1}=10$) significantly outperforms the last setting ($N_{1}=0$), but slightly worse than the other two settings.

We ablated our model (TreeGen-B was used) to analyze the contribution of each component, results also shown in Table 1. First, we compared our model with the traditional Transformer, which is a Transformer without effective structure modeling. We achieved 21 percentage points higher accuracy (p-value is less than 0.001) and 12 higher BLEU score. This result provides strong evidence of the effectiveness of the AST reader in our model and the importance of the structural information. Next, we replaced the tree convolutional layers in the AST Reader with two layers of fully-connected layers, and we removed the char embedding, rule definition encoding, self-attention layers in turn. The experimental results show the identifiers-encoding, alleviating long-dependency and structural information significantly influence the accuracy. Please note that in some cases BLEU increases while StrAcc and Acc+ decrease. Here we consider StrAcc and Acc+ more important as they guarantee the correctness of the generated programs and correctness is usually crucial in code generation.

Experiment II: Semantic Parsing

We further evaluated our approach on the semantic parsing tasks. Our experiment was conducted on two semantic parsing datasets, ATIS and GEO. The input of these datasets is a natural language description, while the output is a short piece of code in lambda calculus. We followed the standard train-dev-test split of these datasets, and the statistics are listed in Table 2.

In this task, we follow the evaluation of the previous approaches (?) and use accuracy as the metric, where the tree exact match was considered to avoid spurious errors. In other words, the order of the children can be changed within conjunction nodes. We followed all the settings in the HS experiment except that we changed the embedding size and the hidden sizes to 128 compared with the setting of the HS experiment.

Table 3 shows the performance of our TreeGen. As seen, the accuracy of our approach is sightly worse than the traditional approach WKZ14 (?), which is based on the CCG parser and uses a large number of templates. This traditional approach is hard to generalize new datasets. However, our model was directly adopted from the HS dataset, and achieved the highest accuracy, among all neural models (?; ?; ?; ?; ?; ?). This experiment shows the effectiveness and generalizability of TreeGen.

Related Work

Code generation achieves significant progress in recent years. The early approaches are mainly based on templates (?; ?; ?; ?). With the prosperity of deep learning, the sequence-to-sequence framework has shown to be effective in various tasks (?). ? (?) applied this framework to generate code based on tokens. Unlike natural languages, it is shown that the code contains much more structural information. Thus, the abstract syntax tree (AST) was used in more recent works (?; ?; ?; ?; ?). However, these studies mainly use recurrent neural networks (RNNs) from the long dependency problem (?). ? (?) proposed to use the convolutional neural network (CNN) to handle the long dependency problem. Our approach addresses this problem by Transformer’s intensive attention mechanism (?). To incorporate the structural information and the idea of self-attention, we propose a tree-based Transformer architecture for code generation.

Conclusion

In this work, we propose TreeGen for program generation. TreeGen uses the attention mechanism of Transformers to alleviate the long-dependency problem and introduces the AST reader to combine the grammar rules and the AST structure.

The evaluation was conducted on a Python dataset, HearthStone, and two semantic parsing datasets, ATIS and GEO. The experimental results show that our model significantly outperforms existing approaches. We also conducted in-depth ablation tests, which suggests that each component in our model plays a significant role.

Acknowledgments

This work is sponsored by the National Key Research and Development Program of China under Grant No. 2017YFB1001803, and National Natural Science Foundation of China under Grant Nos. 61672045, 61529201, and 61922003. Lili Mou is an Amii Fellow; he is supported by the CCAI Chair Program; and he also thanks AltaML for support.

References