TreeBERT：一种面向编程语言的基于树的预训练模型
机器学习
2021-07-16 v2 编程语言
摘要
源代码可根据定义的语法规则解析为抽象语法树(AST)。然而，在预训练过程中，很少有学者考虑将树结构纳入学习过程。在本文中，我们提出TreeBERT，一种基于树的预训练模型，用于改进面向编程语言的生成任务。为利用树结构，TreeBERT将代码对应的AST表示为一组组合路径，并引入节点位置嵌入。该模型通过树掩码语言建模(TMLM)和节点顺序预测(NOP)以混合目标进行训练。TMLM使用一种根据树的特征设计的新型掩码策略，以帮助模型理解AST并推断AST缺失的语义。通过NOP，TreeBERT通过学习AST中节点的顺序约束来提取句法结构。我们在覆盖多种编程语言的数据集上预训练了TreeBERT。在代码摘要和代码文档任务上，TreeBERT优于其他预训练模型以及为这些任务设计的SOTA模型。此外，当迁移到预训练时未见的编程语言时，TreeBERT表现良好。
引用
@article{arxiv.2105.12485,
title = {TreeBERT: A Tree-Based Pre-Trained Model for Programming Language},
author = {Xue Jiang and Zhuoran Zheng and Chen Lyu and Liang Li and Lei Lyu},
journal= {arXiv preprint arXiv:2105.12485},
year = {2021}
}
备注
Accepted by UAI2021