code embedding研究系列五-GraphCodeBert

I still …

已于 2023-04-05 14:44:49 修改

阅读量4.2k

点赞数 9

分类专栏： code embedding 程序分析静态代码检测文章标签：机器学习 python 深度学习

于 2021-03-07 11:44:01 首次发布

本文链接：https://blog.csdn.net/qq_44370676/article/details/114384343

版权

程序分析同时被 3 个专栏收录

32 篇文章 11 订阅

订阅专栏

静态代码检测

28 篇文章 25 订阅

订阅专栏

code embedding

14 篇文章 49 订阅

订阅专栏

1.论文概述

近年来，应用于编程语言的预训练模型得到飞速发展，相关任务比如code search, code completion, code summarization 也得到提升。但是，现有的预训练模型是将code snippet（代码片段）视为一个token序列。忽视了代码的结构。

GraphCodeBERT应运而生，GraphCodeBert是基于数据流（data flow）来表示源代码信息。数据流提取的信息是变量之间的信息流（where-the-value-comes-from）。数据流是一个图结构（graph）。图中的结点表示一个变量（variable），边表示变量之间的依赖关系（where-the-value-comes-from）。

GraphCodeBert采用数据流而不是AST，是考虑到数据流图不像AST这么复杂，也不会带来不必要的深层信息。

1.2.预训练任务

作者在这里提出2个预训练任务：

data flow edges prediction
数据流边预测，用来学习代码的结构化表示
variable-alignment across source code and data flow
源代码和数据流之间的变量分配，用于学习数据流结点来自源代码中哪个token
预训练用到的数据集是CodeSearchNet dataset

1.3.下游任务

作者在4个任务上评估预训练模型，分别是：

natural language code search，代码搜索
clone detection，克隆检测
code translation，代码翻译
code refinement，代码细化

1.4.论文贡献

GraphCodeBERT是第一个利用代码的语义结构来学习代码表示的预训练模型
提出2个预训练任务，用于从源代码和数据流学习向量表示
GraphCodeBERT在四个下游任务上提供了显著的改进，即代码搜索、克隆检测、代码翻译和代码细化

2.预训练

2.1.数据流图

数据流图用来表示变量之间的依赖关系，结点代表变量（variable），边代表变量之间的信息流向（ where the value of each variable comes from），对于同一源代码，不同抽象语法下的数据流是相同的。

使用数据流的好处如下：

以v = maxvalue − minvalue 为例，程序员并不总是遵循命名约定，因此很难理解变量的语义，比如v 。而数据流提供了一种在某种程度上理解变量v的语义的方法。数据流图中，v的值来自数据流中的 maxvalue 和 minvalue 。
数据流支持该模型考虑在遥远的地方使用相同的变量或函数所引起的长期依赖性。（dataflow supports the model to consider long-range dependencies induced by using the same variableor function in distant locations），比如下图中 $x^3, x^7, x^9, x^{11}$ ，4个变量有着相同的变量名，但语义信息不同

获取数据流图的流程如下：

在这里插入图片描述

构造AST
给定源代码 $\left\{ c_1, c_2,..., c_n \right\}$ ，先有标准的编译工具将它们解析成AST。这里用到的解析工具是tree-sitter。该解析工具支持论文中提到的6种编程语言，打破了ANTLR的限制。
从AST种提取变量序列（variable sequence）
变量序列标识为 $\left\{ v_1, v_2,..., v_k \right\}$ ，变量序列种的每个元素都会作为数据流图的一个结点。
从AST中提取变量之间的依赖关系，构建数据流图
数据流图的每个结点来自变量序列，边 $\varepsilon = \langle v_i, v_j \rangle$ 表示变量序列中第 $j$ 个变量依赖于第 $i$ 个变量。在赋值语句x = expr #expr为表达式 这句中。x 依赖于 expr 中的所有变量。边集合 $\left\{ \varepsilon_1, \varepsilon_2,..., \varepsilon_l \right\}$ 。数据流图 $G (C) = (V, E)$ 为源代码 $C$ 的数据流图。

2.2.GraphCodeBERT

模型架构如下图所示
在这里插入图片描述该模型以源代码和注释以及相应的数据流作为输入，并用标准的masked language模型应用在2个结构化预训练任务上：

预测变量来源
predict where a variable is identified from (marked with orange lines)
预测边
data flow edgesprediction between variables (marked with blue lines)

2.2.1.模型架构

GraphCodeBert使用Bert作为模型骨架。

给定源代码 $\left\{ c_1, c_2,..., c_n \right\}$ ，对应的注释 $\left\{ w_1, w_2,..., w_m \right\}$ 。相应的数据流图 $G (C) = (V, E)$ ， $\left\{ v_1, v_2,..., v_k \right\}$ 为变量序列， $\left\{ \varepsilon_1, \varepsilon_2,..., \varepsilon_l \right\}$ 为边集合，其中每条边代表数据流向。

2.2.1.1.输入序列

输入包括以下部分：

源代码 $\left\{ c_1, c_2,..., c_n \right\}$
对应的注释 $\left\{ w_1, w_2,..., w_m \right\}$
变量序列 $\left\{ v_1, v_2,..., v_k \right\}$

最终输入的序列 $X$ 为上面3个序列的连接
$\left\{ [CLS],W,[SEP],C,[SEP],V \right\}$

[CLS]是三段前面的一个特殊标记，[SEP]是2个序列间的分隔符。

2.2.1.2.向量化

输入序列 $X$ 会被转化为向量 $H^0$ 。包括了token和position embedding。并对变量序列 $V$ 使用了一种特殊的position embedding 来标识它们是数据流图的一个结点。模型应用了 $N$ 个transformer层来产生 contextual representations，论文里 $N$ 的值设为12。
$H^n=transformern(H^{n−1}),n∈[1,N]$

$H^n$ 是 $\times d_h$ 维度的向量，论文里隐层向量维度为768。

该transformer层内部如下
$G^n=LN(MultiAttn(H^{n−1}) +H^{n−1})$
$H^n=LN(FFN(G^n) +G^n)$

其中

$M u lt i A tt n$ 是 multi-headed self-attention mechanism，论文设置了12个attention head。
$FFN$ 是 2层前向反馈网络
$L N$ 是 layer normalization操作

$M u lt i A tt n$ 内部计算如下,设 $\hat{G^n} = MultiAttn(H^{n - 1})$

$Q_i=H^{n−1}.W^Q_i, K_i = H^{n−1}.W^K_i, V_i = H^{n−1}.W^V_i$
$head_i = softmax(\frac{Q_i.K^T_i}{\sqrt{d_k}}+ M).V_i$
$\hat{G^n} = [head_1;...;head_u].W^O_n$

其中

$∣ X ∣$ 表示输入序列的长度，包括 token序列, 注释序列，变量序列。
$H^n$ 是 $\times d_h$ 维度的向量。
模型参数 $W^Q_i, W^K_i, W^V_i$ 是 $d_h \times d_k$ 维度向量。
模型参数 $W^O_n$ 是 $d_h \times d_h$ 维度向量。
$M$ 是 Graph-Guided Masked Attention 矩阵（GraphCodeBert相比于Bert的特色之处）， $\times |X|$ 维度向量。引用原文：where $M_{ij}$ is 0 if i-th token is allowed to attend j-th token otherwise $-\infty$ 。大概功能就是在softmax时，如果序列第 $i$ 个和第 $j$ 个token之间没有数据流关联，softmax结果为0。

2.2.2.Graph-Guided Masked Attention

这里用

$v_i$ 表示变量序列 $V$ 第 $i$ 个变量
$c_i$ 表示源代码token集合 $C$ 第 $i$ 个token
$E^{'}$ 定义为，如果变量 $v_i$ 与token序列第 $j$ 个token $c_j$ 相关联，那么 $v_i,c_j〉/〈c_j,v_i〉∈E^′$

为了将图结构引入transformer，这里提出Graph-Guided Masked Attention 来过滤不相关signal。graph-guided masked attention用矩阵 $M$ 表示。

$M_{ij}=\left\{ \begin{array}{rcl} 0 & & {if (q_i∈{[CLS],[SEP]}) or( q_i,k_j∈ W∪C) or (〈q_i,k_j〉∈ E∪E^{'}}) \\ -\infty & & {otherwise}\\ \end{array} \right.$

[CLS], [SEP] 可以和其它序列中所有的元素自由attention。（ $q_i \in <CLS>, <SEP>$ ）
自然语言和code token序列 $W, C$ 中的元素之间可以自由attention。（ $<q_i, k_j> \in W \cup C$ ）
如果一个变量 $v_i$ 在token $c_j$ 处定义，比如 int c = 10; 中code token c 和变量 c 对应。那么 $v_i$ 可以和 $c_j$ attention，反之不行。（ $<q_i, k_j> \in E^{'}$ ）
变量序列中的2个变量 $v_i, v_j$ 只有在存在数据流关系的情况下可以进行attention（ $<q_i, k_j> \in E$ ）。

可以通过如下代码获取Mask矩阵:

#calculate graph-guided masked function，初始化为0，设置为true表示可以自由attention
attn_mask=np.zeros((self.args.max_source_length,self.args.max_source_length),dtype=np.bool)
#calculate begin index of node and max length of input
# [CLS] W [SEP] C [SEP] 部分的position idx 从1开始，而 V [EOS] 部分的position idx全为0，因此node_idx为前半部分不包括data flow的所有结点
node_index = sum([i > 1 for i in self.examples[item].position_idx])
# 除了[CLS]之外的所有结点，包括data flow
max_length = sum([i != 1 for i in self.examples[item].position_idx])

#sequence can attend to sequence
# W, C, [SEP]部分之间的元素可以自由attention
attn_mask[:node_index, :node_index]=True
#special tokens attend to all tokens

# [CLS], [SEP], [EOS] 可以和其它token自由attention
for idx,i in enumerate(self.examples[item].source_ids):
    if i in [0, 2]:
        attn_mask[idx,:max_length]=True

#nodes attend to code tokens that are identified from
# vi和cj如何对应可以自由attention
for idx,(a,b) in enumerate(self.examples[item].dfg_to_code):
    if a < node_index and b < node_index:
        attn_mask[idx+node_index, a:b]=True
        attn_mask[a:b,idx+node_index]=True

#nodes attend to adjacent nodes
for idx, nodes in enumerate(self.examples[item].dfg_to_dfg):
		# a为idx的其中1个前驱结点
    for a in nodes:
        if a+node_index<len(self.examples[item].position_idx):
            attn_mask[idx+node_index,a+node_index]=True

dfg_to_code 的类型为 List[Tuple[int, int]]，长度和变量序列长度一致。为每个变量对应token在sub-token序列的开始和结束索引。
dfg_to_dfg 的类型为 List[List[int]]，长度和变量序列长度一致。为每个变量对应的DFG中的前驱变量索引列表。

2.2.3.预训练任务

预训练任务使用Masked Language Modeling，Edge Prediction和Node Alignment。

2.2.3.1.Mask Language Modeling

对于序列 $CLS>, w_1, ..., <MASK>, ..., w_n, <SEP>, c_1, ..., <MASK>, ..., c_m, <SEP>, v_1, ..., v_l, <EOS>]$ ，还原 $< M A S K >$ 部分的token。这里不会对变量序列进行mask。

2.2.3.2.Edge Prediction

数据流边预测，目的在于让模型学习"where-the-value-comes-from"的信息，对应架构图中蓝色部分。

学习方式如下：
在这里插入图片描述预训练时随机采样20%的node，记为集合 $V_s$ ，mask的方式就是在mask矩阵 $M$ (shape = $\times |X|$ ) 中， $V_s$ 集合中的2个node如何有边相连，那就把值设为 $-\infty$ 。mask的边集合为 $E_{mask}$ 。

以 $x^{11}$ 为例，预训练时会mask 边 $\langle x^{7}, x^{11} \rangle$ 和 $\langle x^{9}, x^{11} \rangle$ 。让模型预测这2条边。

记边集合 $E_c = V_s \times V \cup V \times V_s$ ( $\times$ 为笛卡尔积) 为需要预测的边集合。

定义

$\delta(e_{ij} \in E)=\left\{ \begin{array}{rcl} 1 & & {if \langle v_i, v_j \rangle \in E} \\ 0 & & {otherwise}\\ \end{array} \right.$

训练loss的表达式为，引入了负采样。
$loss_{EdgePred} = - \sum\limits_{e_{ij} \in E_c}[\delta(e_{ij} \in E_{mask}).log(p_{e_{ij}}) + (1 - \delta(e_{ij} \in E_{mask})).log(1 - p_{e_{ij}})]$

$p_{e_{ij}}$ 表示结点 $i$ 和结点 $j$ 之间存在边的概率，由GraphCodeBert模型中2个结点的向量表示内积 sigmoid所得

$p_{e_{ij}} = sigmoid(repr_i. repr_j)$

$repr_i, repr_j$ 分别为GraphCodeBert关于两个变量输出的向量表示。

2.2.3.2.Node Alignment

该任务是为了学习数据流图与源代码之间的对应关系，与边预测不同的是，边预测学习的是变量序列 $V$ 中2个结点之间的联系，而变量分配任务学习的是源代码token序列 $C$ 和变量序列 $V$ 之间的联系，也就是学习变量结点 $v_i$ 和 token $c_j$ 的对应关系。以下图为例， $x_{11}$ 与return x中的x对应。

在这里插入图片描述
与边预测相同，变量分配预测任务同样随机采样20%的node，记为集合 $V_s^{'}$ ，不过这个边连接的是token和变量，mask的方式就是在mask矩阵 $M$ (shape = $\times |X|$ ) 中， $V_s^{'}$ 集合中的2个node如何有边相连，那就把值设为 $-\infty$ 。mask的边集合为 $E_{mask}$ 。

$E_c = V_s^{'} \times C$ ( $C$ 为token序列)

训练loss的表达式为，引入了负采样。
$loss_{NodeAlign} = - \sum\limits_{e_{ij} \in E_c^{'}}[\delta(e_{ij} \in E_{mask}^{'}).log(p_{e_{ij}}) + (1 - \delta(e_{ij} \in E_{mask}^{'})).log(1 - p_{e_{ij}})]$

3.下游任务

3.1.Natural Language Code Search

给定一种自然语言作为输入，代码搜索(Code Search)的目的是从一组候选代码中找出语义上最相关的代码。

这里用到的数据集是CodeSearchNet的语料库。采用的方式是用代码文档的第一段作为query。这与code-nn类似。作者对数据集代码中无关文本（比如http://..）进行了过滤，过滤后数据集统计信息如下

在这里插入图片描述
对比实验结果如下，评估指标为Mean Reciprocal Rank

3.2.Code Clone Detection

代码克隆检测的目的是度量两个代码片段之间的相似性，作者使用了BigCloneBench数据集。

该任务模型输入为2个代码片段，输出它们的相似度。

比如，如下2段java代码，相似度98.3%。
在这里插入图片描述对比实验结果如下

3.3.Code Translation

代码翻译旨在将遗留软件从平台中的一种编程语言迁移到另一种编程语言。用到的数据集是一些开源项目，包括

这些项目都有Java和C#的实现。作者基于文件名和方法名对2种编程语言进行pair，除去重复和函数体为空的方法后，method pairs的总数为11800，作者从中拆分出500对作为development set， 1000作为test set。

该任务中模型输入为Java(C#)代码，输出与之对应的C#(Java)代码。

在这里插入图片描述
对比实验结果如下

3.4.Code Refinement

代码优化旨在自动修复代码中的bug，作者用到了Java数据集。

模型输入Java代码，输出修复后的代码。

在这里插入图片描述
对比实验结果如下

在这里插入图片描述

4.总结

4.1.论文贡献

作者提出了GraphCodeBert，是CodeBert的升级版，与CodeBert相比引入了数据流结构。

CodeBert参考：CodeBERT: A Pre-Trained Model for Programming and Natural Languages

4.2.Graph-Guided Masked Attention矩阵

当然，引入数据流图不代表应用图神经网络，实际上最精髓的部分在于在transformer的Muti-head attention的公式中加入Graph-Guided Masked Attention 矩阵 $M$ 。

普通的head计算公式
$softmax(\frac{Q.K^T}{\sqrt{d_k}}).V$

加入Graph-Guided Masked Attention后
$softmax(\frac{Q.K^T}{\sqrt{d_k}} + M).V$

所以此时有得好好了解下 $M$ 矩阵了。

整个模型的输入序列为 $\left\{ [CLS],W,[SEP],C,[SEP],V \right\}$

$W$ 为注释序列
$C$ 为代码token序列
$V$ 为变量序列（一个变量对应数据流图一个结点）

$M$ 的维度为 $\times |X|$ ，相当于邻接矩阵的改版，定义如下
$M_{ij}=\left\{ \begin{array}{rcl} 0 & & {if (q_i∈{[CLS],[SEP]}) or( q_i,k_j∈ W∪C) or (〈q_i,k_j〉∈ E∪E^{'}}) \\ -\infty & & {otherwise}\\ \end{array} \right.$

大概的含义就是被mask的部分不参与softmax计算（ $M_{ij}$ 为0则不影响softmax，而 $e^{-\infty} = 0$ ，所以 $M_{ij} = -\infty$ 相当于softmax的时候被忽视了）

下图是草草画出的 $M$ 矩阵。
在这里插入图片描述根据定义

白色部分的值均为0
橙色部分，如果token $c_i$ 与变量 $v_j$ 有对应关系（比如上面例子中的return x中的token x 和 $x^{11}$ 就是对应的，其它的token（包括其它的x）和 $x^{11}$ 没有对应关系），那么 $M_{c_iv_j} = 0$ ，否则为 $-\infty$ 。
青色部分，如果变量 $v_i$ 与 $v_j$ 有数据流关系。那么 $M_{v_iv_j} = 0$ ，否则为 $-\infty$ 。