GraphCodeBert：基于数据流的代码表征预训练模型

GraphCodeBert

https://arxiv.org/abs/2009.08366

1 模型结构

使用多层双向 Transformer

变量定义 ：

• C：源码集合
• W：文本集合
• V：变量集合
• E：变量间的边的集合

输入：把注释，源代码和变量集连接为序列 X = { [CLS]，W，[SEP]，C，[SEP]，V }，其中 [CLS] 是三个段前的特殊token，[SEP]是分割两种数据类型的特殊token。

GraphCodeBERT 将序列 X 作为输入，然后把序列转换为输入向量$H^0$，对于 X 中的每一个 token，它的输入向量是通过将其对应的token和位置嵌入（position embedding）相加构建而来的。对变量序列中的所有变量使用特殊的 position embedding来表示他们是数据流的节点。

该模型在输入向量上应用 n 个 Transformer 来产生上下文表示。每个transformer层包含一个架构相同的 transformer block， transformer block应用多头自注意力机制。对于第 n 个 transformer 层，多头自注意力的输出$H^n$由下面的公式计算：

其中前一层的输出$H^{n-1}$分别使用不同的模型参数 W 线性投影到 Q K V，u是头数，$d_k$ 是头的维度。M（|X|×|X|）是mask矩阵，当第i个token和第j个token允许计算 attention score时，M=0，否则M为一个极大的负值。

2 Graph 引导的 Masked Attention

将图结构整合到Transformer中。

通过给注意力分数加一个非常大的负值，softmax之后，注意力分数会变成0，从而可以避免q查询k。

在变量序列中，如果两个变量有一条边直接相连或者两个节点是同一变量，那么就允许q对k进行查询。

为了表示源码token和数据流节点之间的关系，定义一个新的集合E’，如果变量 $v_i$ 是由源码token $c_j$ 所定义，那么就表示为<vi,cj>∈E’。如果存在这种关系，那么就允许计算attendion

• [CLS], [SEP] 可以和其它序列中所有的元素自由attention
• W , C 中的元素之间可以自由attention。
• 如果一个变量 $v_i$ 是由源码token $c_j$ 所定义，比如 int c = 10; 中的 code token c 和变量 c 对应，那么vi 可以和 cj 计算attention，否则不可以。
• 变量序列中的两个变量之间只有存在数据流关系的情况下才可以计算attention

3 预训练任务

• MLM

  MLM用于源码表示学习，MLM目标是预测随机抽样Masked token的原始token，如果源码上下文不足以推出 masked code token，可以使用注释上下文信息，从而促进模型对齐 NL-PL 表示。

• Edge Prediction

  Edge Prediction 用于数据流的表示学习。其目的是为了让模型可以学习结构感知表示，这种表示对“值来自哪里”的关系进行编码，以便更好地理解代码。在数据流中随机抽取20%的节点v，通过在mask矩阵（M）中添加一个无限负值来mask连接这些被采样节点的直接边，然后预测这些masked边。

• 预测 code token 和 node之间的边。其动机是鼓励模型根据数据流来对齐变量和源代码。