Exploring Dynamic Selection of Branch Expansion Orders for Code Generation (2021)

Abstract: Due to the great potential in facilitating software development, code generation has attracted increasing attention recently. Generally, dominant models are Seq2Tree models, which convert the input natural language description into a sequence of tree-construction actions corresponding to the pre-order traversal of an Abstract Syntax Tree (AST). However, such a traversal order may not be suitable for handling all multi-branch nodes. In this paper, we propose to equip the Seq2Tree model with a context-based Branch Selector, which is able to dynamically determine optimal expansion orders of branches for multi-branch nodes. Particularly, since the selection of expansion orders is a non-differentiable multi-step operation, we optimize the selector through reinforcement learning, and formulate the reward function as the difference of model losses obtained through different expansion orders. Experimental results and in-depth analysis on several commonly-used datasets demonstrate the effectiveness and generality of our approach. We have released our code at https://github.com/DeepLearnXMU/CG-RL .

## introduction code generation aims at automatically generating a source code snippet given a natural language (nl) description, which has attracted increasing attention recently due to its potential value in simplifying programming. instead of modeling the abstract syntax tree (ast) of code snippets directly, most of methods for code generation convert ast into a sequence of tree-construction actions. this allows for using natural language generation (nlg) models, such as the widely-used encoder-decoder joint work with pattern recognition center, wechat ai, tencent inc, china. generally, during the generation of dominant seq2tree models based on pre-order traversal, branches of each multi-branch nodes are expanded in a left-to-right order. figure 1 gives an example of the nl-to-code conversion conducted by a seq2tree model. at the timestep t 1 , the model generates a multi-branch node using the action a 1 with the grammar containing three fields: type, name, and body. thus, during the subsequent generation process, the model expands the node of t 1 to sequentially generate several branches in a left-toright order, corresponding to the three fields of a 1 . the left-to-right order is a conventional bias for most human-beings to handle multi-branch nodes, which, however, may not be optimal for expanding branches. alternatively, if we first expand the field name to generate a branch, which can inform us the name 'e', it will be easier to expand the field type with a 'exception' branch due to the high co-occurrence of 'e' and 'exception'. to verify this conjecture, we choose tranx @xcite to construct a variant: tranx-r2l, which conducts depth-first generation in a right-to-left manner, and then compare their performance on the django dataset. we find that about 93.4% of asts contain multi-branch nodes, and 17.38% of ast nodes are multi-branch percentage only tranx 8.47 only tranx-r2l 7.66 in this paper, we explore dynamic selection of branch expansion orders for code generation. specifically, we propose to equip the conventional seq2tree model with a context-based branch selector, which dynamically quantifies the priorities of expanding different branches for multi-branch nodes during ast generations. however, such a non-differentiable multi-step operation poses a challenge to the model training. to deal with this issue, we apply reinforcement learning to train the extended seq2tree model. particularly, we augment the conventional training objective with a reward function, which is based on the model training loss between different expansion orders of branches. in this way, the model is trained to determine optimal expansion orders of branches for multi-branch nodes, which will contribute to ast generations. to summarize, the major contributions of our work are three-fold: ## background as shown in figure 1 , the procedure of code generation can be decomposed into three stages. based on the learned semantic representations of the input nl utterance, the dominant seq2tree model @xcite first outputs a sequence of abstract syntax description language (asdl) grammar-based actions. these actions can then be used to construct an ast following the preorder traversal. finally, the generated ast is mapped into surface code via a user-specified function ast to mr( * ). in the following subsections, we first describe the basic asdl grammars of seq2tree models. then, we introduce the details of tranx @xcite , which is selected as our basic model due to its extensive applications and competitive performance @xcite @xcite ). foot_0 ## dynamic selection of branch expansion orders in this section, we extend the conventional seq2tree model with a context-based branch selector, which dynamically determines optimal expansion orders of branches for multi-branch ast nodes. in the following subsections, we first illustrate the elaborately-designed branch selector module and then introduce how to train the extended seq2tree model via reinforcement learning in detail. ## related work with the prosperity of deep learning, researchers introduce neural networks into code generation. in this aspect, @xcite first explore a seq2seq model for code generation. then, due to the advantage of tree structure, many attempts resort to seq2tree models, which represent codes as trees of meaning representations @xcite @xcite @xcite @xcite . typically, @xcite propose tranx, which introduces asts as intermediate representations of codes and has become the most influential seq2tree model. then, @xcite ) respectively explore cnn and transformer architectures to model code generation. unlike these work, shin et al. (2019) present a seq2tree model to generate program fragments or tokens interchangeably at each generation step. from another perspective, xu et al. ( 2020 ) exploit external knowledge to enhance neural code generation model. generally, all these seq2tree models generate asts in pre-order traversal, which, how-ever, is not suitable to handle all multi-branch ast nodes. different from the above studies that deal with multi-branch nodes in left-to-right order, our model determines the optimal expansion orders of branches for multi-branch nodes. some researchers have also noticed that the selection of decoding order has an important impact on the performance of neural code generation models. for example, alvarez-melis and jaakkola (2017) introduce a doubly rnn model that combines width and depth recurrences to traverse each node. dong and lapata (2018) firstly generate a rough code sketch, and then fill in missing details by considering the input nl description and the sketch. @xcite present an insertionbased seq2seq model that can flexibly generate a sequence in an arbitrary order. in general, these researches still deal with multi-branch ast nodes in a left-to-right manner. thus, these models are theoretically compatible with our proposed branch selector. finally, it should be noted that have been many nlp studies on exploring other decoding methods to improve other nlg tasks @xcite @xcite @xcite . however, to the best of our knowledge, our work is the first attempt to explore dynamic selection of branch expansion orders for tree-structured decoding. ## conclusion and future work in this work, we first point out that the generation of domainant seq2tree models based on pre-order traversal is not optimal for handling all multi-branch nodes. then we propose an extended seq2tree model equipped with a context-based branch selector, which is capable of dynamically determining optimal branch expansion orders for multi-branch nodes. particularly, we adopt reinforcement learning to train the whole model with an elaborate reward that measures the model loss difference between different branch expansion orders. extensive experiment results and in-depth analyses demonstrate the effectiveness and generality of our proposed model on several commonlyused datasets. in the future, we will study how to extend our branch selector to deal with indefinite branches caused by sequential field.