论文笔记：GraphCodeBert:Pre-Trainng Code Representions with Data Flow

首先，为了融入代码结构，开发者并没有考虑AST（抽象语法树），而是采用了DFG（数据流向图）。他们认为AST代表的是各个token之间的句法关系，是syntactic层面的关系，而DFG代表的变量之间的来源关系，是semantic层面的关系。DFG在一定程的上能够避免AST中一些不必要的联系从而提升了模型的有效性；

其次，GraphCodeBert也是基于Transformer开发的；

最后，GraphCodeBert是通过三项预测任务来训练模型的：MLM（掩语言模型）、EP（数据流向图的边预测）和NA（数据流向图节点和代码token之间的关系预测）。

GraphCodeBert是第一个加入语义级别的代码结构的预训练模型，并且提高了四个下游任务（代码搜索，克隆检测，代码翻译和代码改错）的处理能力。

数据流向图

数据流向图DFG是一个有向图，代表的是每个变量的值从哪里来到哪里去的关系。

下图代表了同一个python函数（功能是输入两个数，返回两数中较大的那个数）分别转换成AST和DFG后的结果。

AST在这里不多做解释（目前已经有一个比较完善的AST生成器），来看左边的源码-DFG。源码中所有的数值和变量名都会作为一个节点出现在DFG中，在源码中存在值传递关系的两个变量（或变量与数值），在DFG中相关的两个节点也会通过有向边来连接。以上图为例，函数的返回值 $x^{11}$ 会更多的关注 $x^7$ 或者 $x^9$ ,而不会关注 $x^3$ 。

不同于AST，数据流向在不同的抽象语法中是固定。

和普通的图数据结构一样，DFG可以表示为 $G(C)=(V,E)$ :

C代表的是源码序列 $C=\{c_1,c_2,...,c_n\}$ ;

V为图中所有节点的集合 $V=\{v_1,v_2,...,v_k\}$ ;

E为图中所有边的集合 $E=\{\varepsilon _1,\varepsilon _1,...\varepsilon _l\}$ ，其中 $\varepsilon =<v_i,v_j>$ 。

模型设计

模型架构

输入序列

模型的输入序列由源代码序列 $C=\{c_1,c_2,...,c_n\}$ 、评论序列 $W=\{w_1,w_2,...,w_m\}$ 和从代码中抽取的变量序列 $V=\{v_1,v_2,...,v_k\}$ 构成，即输入序列可以表示为:

$X=\{[CLS],W,[SEP],C,[SEP],V\}$

输入序列X中每个元素通过word embedding并且加上位置编码(变量序列V使用特殊的位置编码来标识这是DFG的节点)后，得到作为编码器输入的向量序列 $H_0$ 。

Transfomer

原文还贴心地帮我们回忆了Transfomer里面的一些前馈过程。

第n-1个block输出的上下文表征向量 $H^{n-1}$ ，先计算多头自注意力：

$MultiAttn(H^{n-1})=[head_1;head_2;...;head_u]W_n^O$

其中

$head_i=softmax(\frac{(H^{n-1}W_i^Q)(H^{n-1}W_i^K)^T}{\sqrt {d_k}}+M)(H^{n-1}W_i^V)$

这里M是一个掩码矩阵,当输入序列中第i个token允许同第j个token做注意力的时候， $M_{ij}=0$ ，否则为负无穷大。

多头自注意力以残差相加的方式依此进入Layer正规化层（LN），全连接层（FFN）和Layer正规化层（LN）后就得到了第n个block的输出 $H^n$ :

$G^n=LN(MultiAttn(H^{n-1})+H^{n-1})$

$H^n=LN(FFN(G^n)+G^n)$

基于图的掩码注意力机制

原文是GRAPH-GUIDED MASKED ATTENTION，不知道我翻译的准不准确。

主要的设计思想就是在掩码矩阵赋予 $M_{ij}$ 负无穷大的值来避免一些不必要的注意：

1、如果第j个DFG节点通过一条有向边指向第i个DFG节点( $<v_j,v_i> \in E$ )者两个节点重合( $i=j$ )，那么 $q_{v_i}$ 是允许和 $k_{v_j}$ 作注意力计算的，即 $M_{ij}=0$ ，否则为负无穷大；

2、如果一个DFG节点是从一个代码token中提取出来的，那么这DFG节点和代码token可以互相做注意力计算；

3、注释序列和代码序列看作一个整体，其中的token可以互相做注意力计算；

4、[CLS]和[SEP]认为是可以和所有的输入token做注意力计算。

预训练任务

同bert一样，GraphCodeBert也只训练编码器。GraphCodeBert通过MLM、EP和NA三个任务来训练编码器。

MLM 掩码语言模型

这个任务也是参考bert的训练工作。

从代码序列和注释序列中选取15%的token，其中80%会被掩盖，10%会替换成一个随机的token，剩下的10%不做变化。MLM的任务就是由扰乱后的序列输出扰乱之前的序列。

EP 边预测

加入这项任务是为了使得模型能够学习到data flow中的信息。

从DFG中随机选取20%的节点，通过给上文提到过的掩码矩阵赋予一个无穷大的值来去除所有指向这些节点的边。EP的任务就是预测这些被掩盖的边 $E_{mask}$ 。如上图EP就是为了预测边 $<x^9,x^{11}>$ 和 $<x^7,x^{11}>$ 。EP的损失函数计算如下：

$\small loss_{EdgePred}=-\sum_{e_{ij} \in E_c}[\delta (e_{ij}\in E_{mask})logp_{e_{ij}}+(1-\delta (e_{ij}\in E_{mask}))log(1-p_{e_{ij}})]$

其中 $\small E_c$ 代表的所有采样出来用来预测的候选边；

如果 $\small e_{ij} \in E_{mask}$ 那么 $\small \delta (e_{ij} \in E_{mask})=1$ ，否则等于0。

NA 节点对齐

NA的目的是为了预测节点与代码之间的来源关系，想让模型学习到一个DFG节点到底来自于code序列中的那个token。上图就是来预测DFG节点的 $\small x^{11}$ 到底来同code中的那个x是对齐的。

随机选取20%的DFG节点，通过赋予掩码矩阵一个无穷大的值来去除这些节点于code的联系，并预测这些被掩盖后的边 $\small E_{mask}'$ 。损失函数的计算类似于EP任务：

$\small loss_{NodeAlign}=-\sum_{e_{ij} \in E_c'}[\delta (e_{ij}\in E_{mask}')logp_{e_{ij}}+(1-\delta (e_{ij}\in E_{mask}'))log(1-p_{e_{ij}})]$

下游任务

为了显示GraphCodebert的有效性，作者进行了四项下游任务，包括了：代码搜索、克隆检测、代码翻译和代码改错。

代码搜索

这个实验的任务是通过输入一段自然语言描述，模型输出从代码库中选出一段最符合描述的代码。这个实验所用的数据集为经过预处理之后的CodeSearchNet，并且以Mean Reciprocal Rank (MRR)作为评价标准。

实验结果对比

以上模型主要通过计算代码和查询描述之间的内积来对候选的代码进行排名的。第一组分别代表了词袋模型、卷积神经网络、双向RNN和自注意力机制；第二组代表的一些预训练模型，RoBerta是一个在文本语料库上通过MLM任务训练的模型，RoBerta（code）则是在代码上的预训练模型，CodeBert是在文本-代码对上通过MLM训练的模型。

案例

上图内容显示了两个案例，发现GraphCodeBert相比于其他模型能够很好的匹配出最优的代码。

克隆检测

任务的目的是给出两个输入代码的相似度。实验使用的数据集是BigCloneBench

实验结果对比

Deckard通过AST计算代码的向量并且以局部敏感哈希对向量聚类；

RtvNN通过训练递归自编码器来学习AST的表征；

CDLH是以AST为基础的LSTM，并用汉明距离计算两个AST表征向量之间的距离；

ASTNN通过RNN来对代码片段编码AST子树，一次来学习整个代码的表征；

FA-AST-GMN是一个基于AST的GNN。

案例

上图显示了对于两个字面相似但是语义不同的函数，GraphCodeBert给出了正确的结果而其他两个模型预测错误。

代码翻译

代码翻译的任务就是能够将一种语言的代码转换成另一种语言。

实验吗结果对比

Navie是将源语言原封不动输出用来做参照；

PBSTMT代表的是以短语为基础的统计学机翻模型。

案例

代码改错

实验结果对比

上图显示了不同模型分别在短和中等长度的代码中的能力评价。

案例

模型分析

消融实验

对比去除任意一个组件在代码搜索任务上的表现：

注意力分布

[CLS]被认为是可以和一个输入序列中所有的token做注意力，因此[CLS]在一定程度上也代表了这个代码的表征向量。

上图第一行表示了代码和DFG节点的数量分布；第二行表示了[CLS]对于代码和节点的注意力分布。此图表说明了模型对于单个节点具有更高的注意力。

AST与DFG的对比

AST Pre-order Traversal通过先序遍历来序列化一个AST；AST Subtree Masking是一个transfomer模型，但是AST上的每一个节点只能注意到他的子代。

上图显示，当输入序列长度较低时，AST相关的模型的MMR评分低于参照，而无论输入长度如何GraphCodeBert表现效果都明显高于AST。作者给出的解释是：AST中包含了大量不必要的联系。

上图则显示了data flow的有效性，通过对比有没有dataflow这个结构的模型对自然语言描述和代码的匹配度。

模型局限

首先模型不能处理一些特有的库，如tensorflow；其次模型还会出现一些类似于标识符未定义先使用、括号没闭合等问题。

个人感想

很赞同导师的一句话，如果你看完一篇论文，那么你下一篇就应该看看这篇论文的参考论文，或者说学习了一个模型就应该去看这个模型的baseline，因为往往是baseline才能展现出这个模型的设计思想，而最新的模型只是在baseline上加入一些特殊的结构。

从这篇论文中我看到了非常多的Transfomer和Bert的先验知识，因此建议读者很有必要去阅读Attention is all your need和Bert: Pre-training of deep bidirectional transformers for language understanding。如果有必要我也会给出自己对这两个模型的理解。

目前code intelligence还仍然是一个待研究和实践的新领域，我的研究也在探索提高生成的准确率和效率。