[论文解读] TreeGen: A Tree-Based Transformer Architecture for Code Generation

Zeyu Sun, Qihao Zhu+4 more|arXiv (Cornell University)|Nov 22, 2019

Software Engineering Research参考文献 23被引用 18

一句话总结

TreeGen 是一种新颖的基于树的 Transformer 架构，通过整合抽象语法树（AST）结构和注意力机制，解决了长距离依赖和结构建模的挑战，显著提升了代码生成性能。其在 HearthStone 基准测试中准确率提升 4.5 个百分点，达到 SOTA 水平，在 ATIS 和 GEO 语义解析基准测试中分别取得 89.1% 和 89.6% 的准确率。

ABSTRACT

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.

研究动机与目标

• 为解决代码生成中的长距离依赖问题，即远距离代码元素（如变量引用与定义）之间的语义关联。
• 通过将 AST 语法规则和树状结构信息融入神经网络架构，提升代码生成中的结构建模能力。
• 克服标准 RNN 和平面神经网络在捕捉程序层次结构方面的局限性。
• 设计一种混合架构，结合 Transformer（用于长距离依赖建模）与树状卷积（用于结构感知）。
• 评估特定组件设计选择的有效性，特别是结构卷积层在早期 Transformer 块中的位置安排。

提出的方法

• TreeGen 采用三段式架构：自然语言（NL）读取器（编码器）用于输入描述，AST 读取器（前几层 Transformer 解码器块，集成结构卷积）用于部分 AST 编码，解码器（剩余块）用于语法规则预测。
• AST 读取器引入一种新颖的结构卷积子层，将节点表示与结构邻居信息结合，但仅应用于前几层 Transformer 块，以保留原始节点信息。
• 模型将代码生成建模为一系列语法规则分类任务，每一步通过预测的规则扩展 AST 中的一个节点。
• 利用 Transformer 架构中的自注意力机制，对远距离代码元素之间的长距离依赖关系进行建模。
• 采用学习得到的嵌入表示字符、语法规则和标识符，并通过位置编码保持序列顺序。
• 消融实验评估各组件的贡献，包括标识符编码、长距离依赖建模以及结构信息整合。

实验结果

研究问题

• RQ1 将 AST 结构整合到基于 Transformer 的架构中，是否能相比标准模型提升代码生成的准确率？
• RQ2 仅在早期 Transformer 块中应用结构卷积层，是否优于在所有块中均应用？
• RQ3 注意力机制在代码生成任务中捕捉长距离依赖关系的效率如何？
• RQ4 标识符编码与语法规则表示在多大程度上提升了模型性能？
• RQ5 所提出的架构在不同代码生成与语义解析基准测试中是否具备良好的泛化能力？

主要发现

• 在 HearthStone Python 代码生成基准测试中，TreeGen 相较于之前的 SOTA 模型，准确率提升了 4.5 个百分点。
• 在 ATIS 语义解析基准测试中，TreeGen 达到 89.1% 的准确率，是所有基于神经网络方法中的最高水平。
• 在 GEO 语义解析基准测试中，TreeGen 达到 89.6% 的准确率，同样是所有神经网络模型中的最佳表现。
• 消融实验表明，仅在早期块中应用结构卷积显著优于在所有块中应用，验证了该设计选择的有效性。
• 移除标识符编码、通过注意力机制实现的长距离依赖建模或结构信息整合，均导致可测量的性能下降，证实了这些组件的关键作用。
• 尽管某些情况下组件移除后 BLEU 分数有所上升，但 StrAcc（结构匹配准确率）与 Acc+（精确匹配准确率）均下降，表明在代码生成中结构正确性被优先考虑。