Source-linked AI summary
TreeGen: A Tree-Based Transformer Architecture for Code Generation
Zeyu Sun, Qihao Zhu, Yingfei Xiong, Yican Sun, Lili Mou, Lu Zhang
TL;DR
Code generation must address long dependencies and rich program structure that existing neural approaches handle imperfectly. TreeGen combines Transformer attention with an AST reader and structural processing, outperforming prior approaches across Python code generation and semantic parsing benchmarks. Its evaluation and ablations support the importance of these architectural components.
Problem
Code generators face long-dependency and structure-modeling challenges because code elements can depend on distant elements and flat architectures capture tree structure poorly.
Method
TreeGen combines Transformer attention with an AST reader using structural convolution in the first several decoder blocks, alongside an NL reader and grammar-rule decoder.
Results
TreeGen outperforms previous approaches on HearthStone and achieves the best accuracy among previous neural models on ATIS and GEO, reaching 89.1% and 89.6%, respectively.
Takeaways & Limitations
Across Python code generation and semantic parsing, the evaluation indicates that TreeGen’s attention and AST-structure modeling are effective, with ablations showing significant contributions from its components.
Takeaways & Limitations
The rule-definition table lookup treats each grammar rule as an atomic token and loses information about the rule’s content.
Abstract
from arXiv · showhide
A code generation system generates programming language code based on an input natural language description. State-of-the-art approaches rely on neural networks for code generation. However, these code generators suffer from two problems. One is the long dependency problem, where a code element often depends on another far-away code element. A variable reference, for example, depends on its definition, which may appear quite a few lines before. The other problem is structure modeling, as programs contain rich structural information. In this paper, we propose a novel tree-based neural architecture, TreeGen, for code generation. TreeGen uses the attention mechanism of Transformers to alleviate the long-dependency problem, and introduces a novel AST reader (encoder) to incorporate grammar rules and AST structures into the network. We evaluated TreeGen on a Python benchmark, HearthStone, and two semantic parsing benchmarks, ATIS and GEO. TreeGen outperformed the previous state-of-the-art approach by 4.5 percentage points on HearthStone, and achieved the best accuracy among neural network-based approaches on ATIS (89.1%) and GEO (89.6%). We also conducted an ablation test to better understand each component of our model.
Introduction
Code generation must handle both long dependencies between distant code elements and the structural information encoded by program trees. TreeGen combines Transformer attention, structural processing, and separate readers to address these challenges.
- Background: Code generation translates natural-language specifications into executable programs, with neural systems commonly using sequence-to-sequence or sequence-to-tree architectures.State-of-the-art systems predict grammar-rule sequences while maintaining a partial AST.
- Challenges: Long dependencies arise when a code element, such as a variable reference, depends on a far-away definition.The cited example places a reference at line 100 and its definition at line 10.
- Challenges: Flat architectures such as RNNs cannot capture program structure well, although tree-structural information is crucial for code modeling.This motivates incorporating AST structure into the neural architecture.
- TreeGen: TreeGen uses Transformer attention for long dependencies and addresses the Transformer’s lack of program structure through structural convolution in selected decoder blocks.The structural convolution sub-layer is added only to the first several Transformer decoder blocks.
- TreeGen: TreeGen contains an NL reader, an AST reader for previously generated partial code, and a decoder that predicts the next grammar rule.The decoder combines the natural-language and AST representations with the node selected for expansion.
- Evaluation: TreeGen outperforms previous models by 4.5 percentage points on HearthStone and achieves 89.1% and 89.6% accuracy on ATIS and GEO.The evaluation also reports that using structural convolution in the first several Transformer blocks outperforms using it in all blocks.
Our Model
TreeGen models code generation as grammar-rule prediction over an AST, combining natural-language encoding, partial-AST structure modeling, and node-conditioned decoding. Its architecture uses Transformer attention, gating, convolution, and pointer copying to represent dependencies, local context, and user-defined identifiers.
- AST-based generation: TreeGen generates programs by predicting grammar rules that expand non-terminal AST nodes in pre-order until the leaves are terminal.The next-rule distribution is conditioned on the natural-language input and the currently generated partial AST.
- Decoder: The decoder uses the node to be expanded as a query, representing it by the path from the AST root to that node.For example, the path to an Assign node is represented as root, Module, body, Assign before decoding the next rule.
- NL Reader: The NL reader encodes token sequences with embeddings, positional information, self-attention, character-aware gating, and convolutional layers.Self-attention captures long dependencies, while gating integrates character embeddings and convolution extracts local token features.
- AST Reader: The AST reader encodes the partial program by combining grammar-rule representations with tree structures and ancestor information.It first represents generated rules, applies attention, and then uses tree convolution to combine each node with its ancestors.
- Decoder: A pointer network can copy tokens from the natural-language description, producing rules of the form α →a for user-defined identifiers.This mechanism is specifically described as helpful for variable and function names.
Evaluation
TreeGen was evaluated on Python code generation and semantic parsing benchmarks, with comparisons covering preprocessing, efficiency, structural-convolution placement, and component ablations. It improved accuracy over prior approaches, achieved strong neural-model performance on semantic parsing, and ran faster than CNN and RNN baselines.
- HearthStone: TreeGen improved accuracy by 6 percentage points with plain preprocessing and 4.5 percentage points with structural preprocessing on HearthStone.Structural preprocessing performed better than plain preprocessing, and TreeGen also achieved the best BLEU results.
- HearthStone: 18s per epoch made TreeGen faster than CNN at 180s and RNN at 49s on a single Nvidia Titan XP.
- Ablation Test: Adding structural convolution to all Transformer blocks outperformed omitting it but was slightly worse than adding it to the first 7 or 8 blocks.The study compared four settings: all 10 blocks, the first 7 blocks, the first 8 blocks, and none.
- Ablation Test: The ablation study found that identifiers encoding, long-dependency alleviation, and structural information significantly influence accuracy.Compared with a traditional Transformer, TreeGen achieved 21 percentage points higher accuracy and 12 higher BLEU score, with p-value less than 0.001.
- Semantic Parsing: TreeGen achieved the highest accuracy among neural models on ATIS and GEO semantic parsing benchmarks.The evaluation used tree exact match accuracy, allowing child-order changes within conjunction nodes.
Related Work
Code-generation research progressed from templates to neural sequence-to-sequence methods, with later work incorporating ASTs because code contains substantial structural information.
- Early code-generation approaches primarily used templates, before deep learning introduced sequence-to-sequence methods.
- Ling et al. applied sequence-to-sequence modeling to generate code based on tokens.
- Because code contains more structural information than natural language, recent studies increasingly use abstract syntax trees.
- Earlier AST-based studies mainly relied on recurrent neural networks.
Conclusion
TreeGen combines Transformer attention with an AST reader for program generation and was evaluated across code-generation and semantic-parsing benchmarks, where it outperformed existing approaches.
- TreeGen uses Transformer attention to alleviate long dependencies and an AST reader to combine grammar rules with AST structure.
- TreeGen was evaluated on HearthStone, ATIS, and GEO.
- The experimental results show that TreeGen significantly outperforms existing approaches.
- In-depth ablation tests suggest that each TreeGen component plays a significant role.