TreeBERT：基于树的编程语言预训练模型。

TreeBERT

https://arxiv.org/abs/2105.12485

Comments: Accepted by UAI2021
Subjects: Machine Learning (cs.LG); Programming Languages (cs.PL)
Cite as: arXiv:2105.12485 [cs.LG]

1 Introduction

现有挑战：

• 设计适当的机制来学习程序的语法结构

  代码是强结构化的，代码的语义依赖于要表示的具有不同语法结构的程序语句和表达式的组合，不能仅仅采用类似于自然语言的处理方法（简单的将代码建模为单词序列）。

  如何使用AST作为预训练模型的输入？

• 树形结构的预训练任务探索

  面向序列的任务直接应用于非顺序结构化AST中存在一些不恰当的问题，因此，需要为树设计新的预训练任务，使预训练模型能够同时从AST中提取语法和语义信息。

主要贡献：

• 提出了TreeBERT，一个面向pl的基于树的预训练模型。
• 遵循Transformer编码器-解码器架构。为了使Transformer能够利用树形结构，将代码片段对应的AST表示为根节点到终端节点路径的集合，然后引入节点位置嵌入以获得节点在树中的位置。
• 提出了一种适用于AST的混合目标来学习语法和语义知识，即树掩码语言建模(TMLM)和节点顺序预测(NOP)。
  • 在TMLM中，在编码器侧输入具有屏蔽节点的AST路径集，并在解码器侧使用AST中的上下文信息预测完整的代码片段。
  • 由于路径中节点的顺序表达了程序结构信息，NOP通过预测节点是否无序来提高模型捕获语法结构信息的能力。

2 Model

2.1 Model Architecture

采用transformer的编码器-解码器结构。修改了Transformer的编码器端，只添加了一个完全连接的层来调整输入的维度。

2.2 Input Representation

AST Representation

AST以树的形式展示了程序的语法结构。树中的每个节点表示代码中的一个结构。

AST节点分为两类：

• 类型节点（type node），用 v 表示。
• 值节点（value node），用 x 表示。值节点几乎都是叶节点（终端节点），除了函数名是非叶子节点但是使用了 value 属性。

用从根节点到叶节点的路径集合表示AST，A = { p1,p2,…,pN }，N表示AST中路径的个数。

Code Representation

与AST相应的代码片段被分割成一系列 tokens，[LT] 和 [CLS] 分别被添加在序列的开始和结尾处。

C = [LT]，x1，x2，…，x3，[CLS] ，其中 [LT] 是 [LT] 的向量表示， [CLS] 是 [CLS] 的向量表示，M是代码片段的长度。

C 被使用在解码器的输入。[EOS] 是解码器端的句尾标识符。

[LT] 不仅作为解码器端的句子开头标识符，它的值还表示目标编程语言的类型。例如， [LT] = [PLT] 表示语言类型是 Python，[LT] = [JLT] 表示语言为Java，当 [LT] = [UNK] 时表示编码器生成的语言是在预训练阶段未见过的语言。这样定义是因为在将代码片段转换为AST时，隐藏了不同类型语言的实现细节，我们需要提示语言类型，以便模型了解不同编程语言之间的差异。

使用 [CLS] 作为NOP的聚合表示。

Path Representation

每一个 path 是一个 nodes 序列，$p_i=v_1^i{v}_2^i...v_{L-1}^i{x}_t^i$，path 上的叶子节点 $x_t^i$ 是对应的代码片段的一个 token，L是 path 的长度。

我们将路径上的节点向量连接起来以表示路径：

$p_i=Concat[v_1^i;v$