Antecedent Predictions Are More Important Than You Think: An Effective Method for Tree-Based Code Generation

TL;DR

This paper introduces a novel Antecedent Prioritized Loss for tree-based code generation, emphasizing the importance of accurate antecedent predictions to improve overall code generation accuracy.

Contribution

The paper proposes a new loss function and a position-aware AST modeling method to enhance antecedent prediction accuracy in tree-based code generation models.

Findings

Significant performance improvements on four benchmark datasets.

The proposed method outperforms existing Seq2Tree approaches.

Enhanced antecedent prediction leads to more accurate code generation.

Abstract

Code generation focuses on the automatic conversion of natural language (NL) utterances into code snippets. The sequence-to-tree (Seq2Tree) approaches are proposed for code generation, with the guarantee of the grammatical correctness of the generated code, which generate the subsequent Abstract Syntax Tree (AST) node relying on antecedent predictions of AST nodes. Existing Seq2Tree methods tend to treat both antecedent predictions and subsequent predictions equally. However, under the AST constraints, it is difficult for Seq2Tree models to produce the correct subsequent prediction based on incorrect antecedent predictions. Thus, antecedent predictions ought to receive more attention than subsequent predictions. To this end, in this paper, we propose an effective method, named Antecedent Prioritized (AP) Loss, that helps the model attach importance to antecedent predictions by exploiting the position information of the generated AST nodes. We design an AST-to-Vector (AST2Vec) method, that maps AST node positions to two-dimensional vectors, to model the position information of AST nodes. To evaluate the effectiveness of our proposed loss, we implement and train an Antecedent Prioritized Tree-based code generation model called APT. With better antecedent predictions and accompanying subsequent predictions, APT significantly improves the performance. We conduct extensive experiments on four benchmark datasets, and the experimental results demonstrate the superiority and generality of our proposed method.