本文正在参加arxiv.org/abs/2009.08…
1 简介
NLP的飞速发展促进了Code LLM的广泛研究，然而，以前的工作只利用源代码进行预训练，而忽略了代码的固有结构。以表达式v = max_value−min_value为例，v由max_value和min_value计算得到。程序员并不总是遵循命名约定，因此仅从变量v的名称很难理解其语义。
本文提出了GraphCodeBERT模型，它考虑了代码的固有结构。利用代码的语义级信息(即Data Flow，数据流)进行预训练。数据流是一个图，其中节点表示变量，边表示变量之间“值从哪里来”的关系。与AST相比，数据流不那么复杂，没有深层层次结构，这一特性使得模型更加高效。同时，为了从源代码和代码结构中学习代码表示，本文引入了两个新的结构感知预训练任务:
- 数据流边缘预测，用于从代码结构中学习表示;
- 跨源代码和数据流的变量对齐，用于在源代码和代码结构之间对齐表示。
GraphCodeBERT基于Transformer架构，通过引入graph-guided masked attention来扩展它，以纳入代码结构。通过在CodeSearchNet数据集上预训练，GraphCodeBERT在四个下游任务上都达到了最先进的性能。进一步分析表明，代码结构和新引入的预训练任务可以改善GraphCodeBERT，并且模型对加入的数据流具有一致的偏好。
2 相关工作
略
3 数据流
数据流是表示变量之间依赖关系的图，其中节点表示变量，边表示每个变量的值从何而来。以图1为例，有四个变量名称相同(即x3, x7, x9和x11)，但语义不同。图中的图表显示了这些变量之间的依赖关系，如x11更加关注x7和x9，而不是x3。
图1展示了如何从源代码中提取数据流图。首先，需要将原始代码转为一个AST树，然后将树中每个叶子节点（几变量）作为一个图中的节点，如果一个变量j来源自变量i，那么将节点i到j进行连边。
4 GraphCodeBERT
4.1 模型架构
图2展示了GraphCodeBERT的模型架构，输入由三部分组成，分别是代码C，注释W和数据流图G(C)=(V,E).模型最终的输入序列为：X = { [CLS], W, [SEP], C, [SEP], V }
整个模型由12层Transformer-layer组成。
4.2 Graph-guied Masked Attention
考虑到AST是一种图结构，为了让Transformer能适应其与一般序列结构的差异，作者修改了其注意力机制，主要是通过调整Attention Mask缩小感受野。
作者将其称为Graph-Guided Masked Attention，其中E代表的是数据流图的边，E’代表的是数据流图的节点和代码的对应关系边。
4.3 预训练任务
除了MLM任务外，作者还引入两个新的任务：
- Edge Prediction：即数据流图中的连边预测，如图2中的x11和x9、x7、x3之间的连边（蓝色），通过预测数据流图的边，帮助模型学习代码的结构信息；
- Node Alignment：即变量对齐，如图2中的红色边，让模型学习数据流图中的某个node来自输入代码中的哪个code token。
5 实验
5.1 代码搜索
所有模型都通过计算代码和查询编码的内积作为相关性得分来对候选代码进行排名。评价指标为MRR。
5.2 代码克隆检测
5.3 代码翻译
5.4 代码优化
代码优化旨在自动修复代码中的错误，这有助于减少错误修复的成本。
5.5 模型分析
表5显示了消融实验的结果，当分别去除节点对齐和边缘预测预训练任务时，模型性能有所下降。在完全去除数据流图之后，性能显著下降。
[CLS]是可以和一个输入序列中所有的token做注意力，因此[CLS]在一定程度上也代表了这个代码的表征向量。如表6所示，第一行表示了代码和数据流图中节点的数量分布；第二行表示了[CLS]对于代码和数据流节点的注意力分布。可以看到，更少的数据流节点比例获得了更高的注意力分布，这表明了数据流图的重要性。
图4显示了代码搜索任务中Ruby编程语言验证数据集上MRR相对于输入序列长度的得分。 AST Pre-order Traversal通过先序遍历来序列化一个AST；AST Subtree Masking是一个transfomer模型，但是AST上的每一个节点只能注意到他的子代。结果显示，当输入序列长度较低时，AST相关的模型的MRR评分低于参照，而无论输入长度如何GraphCodeBert表现效果都明显高于AST。作者认为，这是由于AST层次化过于复杂导致。
同时，作者也给了几个Case Study来展示GraphCodeBERT的优越性。
6 总结
本文提出GraphCodeBERT，它利用数据流来学习代码表示。这是第一个为预训练代码表示考虑代码结构的预训练模型。作者介绍了两个结构感知的预训练任务，并展示了GraphCodeBERT在四个与代码相关的下游任务上达到了最先进的性能，包括代码搜索、克隆检测、代码翻译和代码优化。进一步的分析表明，代码结构和新引入的预训练任务提高了性能。此外，在代码搜索任务中的案例研究表明，在预训练的模型中应用数据流可以提高代码理解能力。
