[论文解读] TreeBERT: A Tree-Based Pre-Trained Model for Programming Language

Xue Jiang, Zhuoran Zheng+3 more|arXiv (Cornell University)|Jul 27, 2021

Software Engineering Research被引用 1

一句话总结

TreeBERT 是一种基于树结构的预训练模型，通过将代码的抽象语法树（ASTs）表示为组合路径，并使用节点位置嵌入来建模，从而提升编程语言生成能力。它采用树掩码语言建模（TMLM）和节点顺序预测（NOP）的预训练目标，在代码摘要生成和文档生成任务上超越了现有模型，包括在未见过的编程语言上的零样本迁移。

ABSTRACT

Source code can be parsed into the abstract syntax tree (AST) based on defined syntax rules. However, in pre-training, little work has considered the incorporation of tree structure into the learning process. In this paper, we present TreeBERT, a tree-based pre-trained model for improving programming language-oriented generation tasks. To utilize tree structure, TreeBERT represents the AST corresponding to the code as a set of composition paths and introduces node position embedding. The model is trained by tree masked language modeling (TMLM) and node order prediction (NOP) with a hybrid objective. TMLM uses a novel masking strategy designed according to the tree's characteristics to help the model understand the AST and infer the missing semantics of the AST. With NOP, TreeBERT further extracts the syntactical structure by learning the order constraints of nodes in AST. We pre-trained TreeBERT on datasets covering multiple programming languages. On code summarization and code documentation tasks, TreeBERT outperforms other pre-trained models and state-of-the-art models designed for these tasks. Furthermore, TreeBERT performs well when transferred to the pre-trained unseen programming language.

研究动机与目标

• 通过在预训练中引入树状结构的 AST 来提升编程语言生成能力。
• 解决现有代码预训练模型中未充分利用树结构的问题。
• 设计一种混合预训练目标，以同时捕捉代码 AST 中的语义和句法结构。
• 通过鲁棒的结构表征实现对未见编程语言的有效零样本迁移。
• 通过树感知的表征学习提升代码摘要生成和文档生成性能。

提出的方法

• 将代码 AST 表示为组合路径，以保留其层次结构并支持树感知建模。
• 引入节点位置嵌入，以编码节点在 AST 结构中的相对位置。
• 设计一种树掩码语言建模（TMLM）目标，采用树特定的掩码策略，以预测缺失的 AST 节点。
• 提出节点顺序预测（NOP），通过预测 AST 中节点的正确顺序来学习句法约束。
• 将 TMLM 和 NOP 结合为混合预训练目标，以联合优化语义和结构理解。
• 在多语言代码数据集上预训练 TreeBERT，以提升跨编程语言的泛化能力。

实验结果

研究问题

• RQ1 一种显式建模 AST 结构的预训练模型是否能提升摘要生成和文档生成等代码生成任务的性能？
• RQ2 采用树特定掩码策略的树掩码语言建模（TMLM）在捕捉 AST 语义方面有多有效？
• RQ3 节点顺序预测（NOP）在多大程度上增强了模型对代码句法结构的理解？
• RQ4 TreeBERT 是否能通过零样本迁移泛化到未见的编程语言？
• RQ5 TreeBERT 在代码生成基准测试中与现有预训练模型及 SOTA 方法相比表现如何？

主要发现

• TreeBERT 在多种编程语言的代码摘要生成和文档生成任务中均优于现有预训练模型。
• TMLM 与 NOP 的混合预训练目标相比单独使用任一目标，均能带来更优的性能。
• TreeBERT 在未见编程语言上实现了强大的零样本迁移性能，表明其具有鲁棒的结构泛化能力。
• TMLM 中的树特定掩码策略显著提升了模型推断 AST 中缺失语义内容的能力。
• 节点顺序预测（NOP）有效捕捉了句法约束，增强了模型对代码结构的理解。
• TreeBERT 在代码生成基准数据集上取得了最先进结果，证实了基于树的预训练方法的有效性。