Paper Title
GraphCodeBert: Pre-Training Code Representations With DataFlow:ICLR(A) 2021, Daya Guo et al.
Abstracrt
当前的预训练模型把代码当成一串tokens,忽略了代码的内在结构,这个内在结构中包含代码语义可以增强代码理解。文章提出了GraphCodeBert,一个预训练模型。由于AST这种语法级结构比较复杂,会带来深度的层次结构AST使得模型效率降低,我们使用了数据流的方式。
文章还引入了两种结构感知预训练任务,一个预测结构边,一个右对齐源码和代码结构的表示。我们使用了一个图引导掩码注意力表示函数去结合代码结构,实现了有效的模型。
Introduction
data flow
数据流是一个图,节点表示变量,边表示变量间的“值从哪来来”的关系
数据流比AST更加简单,不会带来深层次的结构。
数据集
CodeSearchNet
文章使用这个数据集在代码搜索,克隆检测,代码翻译,代码增强四个任务中都取得了良好的效果
贡献
-
首次提出预训练模型利用代码语义结构来学习代码表示
-
引入两个代码感知预训练任务分别从源码和数据流中学习表示
-
在四个下游任务中都有重要的提升
Data Flow
数据流是通过把变量作为节点,值的传递关系作为边,生成一个图。
我们可以用tree-sitter把代码生成AST,使用终端识别变量列表V,把每个变量作为一个节点,变量之间的值传递作为边集合E,生成数据流图(V,E)。
tree-sitter:
tree-sitter/tree-sitter: An incremental parsing system for programming tools (github.com)
GraphCodeBert
结构
把注释,源代码和变量集合作为输入X={[CLS] , W , [SEP] , C , [SEP] , V},[CLS]是在三个片段前的特殊标记,[SEP]是特殊的符号分割不同的数据类型。
GraphCodeBert输入X作为向量,对于每个token,它的输入向量是通过对应的token和位置嵌入相加 来构建的,和Bert的embedding方式类似,GraphCodeBert在对单词进行编码时采用了token embedding和position embedding结合的方式,position embedding也体现了变量在变量序列之中的位置信息,这个位置信息也对应着数据流中的不同结点。
GraphCodeBert模型使用了12个transformer encoder层来组成核心网络结构,采用12个attention head的多头注意力机制,包含Feed Forward层和Layer Normalization层等等,与我们熟知的transformer不同的是,本文中增加了一个Graph-Guided Masked Attention层,这个层与传统的Attention层的区别是在softmax计算权重之前需要增加一个参数M,功能是用来过滤无效的元素(在softmax之前加上负无穷)。
Graph-guided masked attention
定义了一个图引导的掩码注意力函数过滤掉不需要的信号,将图结构导入到transformer.
Pre-train task
MLM
从源码中学习其表示
Edge Predict
首先屏蔽一些变量的数据流边,然后让GraphCodeBert预测这些边。
该预训练任务的目的是让模型学习“值从哪里来”的信息,对应了模型架构图中的蓝色部分。预训练时随机采样20%的结点,通过mask矩阵(如果采样结点中有2个结点之间有边相连,就设为负无穷)来实现边的mask,然后让模型预测被mask的边。
Node Alignment
对齐源代码和数据流之间的表示
该预训练任务的目的是让模型学习数据流图与源代码之间的对应关系,对应模型架构图中的橙色部分。与边预测不同的是,边预测任务学习的是变量序列V中两个结点之间的联系,而变量对齐任务学习的是源代码序列C和变量序列V之间的联系,也就是结点 和单词 之间的对应关系。
如上图所示, 与源代码中"return x"中的"x"对应。与边预测类似,该任务同样随机采样20%的结点,通过mask矩阵来实现结点与单词之间边的mask,模型通过预测被mask的边,学习结点和单词的关系。