Exploring Dynamic Selection of Branch Expansion Orders for Code Generation

TL;DR

This paper introduces a reinforcement learning-based method to dynamically select branch expansion orders in Seq2Tree models for code generation, improving handling of multi-branch nodes and enhancing overall performance.

Contribution

It proposes a context-based Branch Selector optimized via reinforcement learning to adaptively determine expansion orders in code generation models, a novel approach for this task.

Findings

Improved accuracy on multiple code generation datasets

Demonstrated effectiveness of dynamic branch selection

Showed generality across different models and datasets

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.