Graph为Code Intelligence赋能
一、简介
代码智能(Code Intelligence)是计算机科学和软件工程中的一个研究领域,旨在使机器能够理解、生成和改进源代码。代码智能的目标是通过机器学习和深度学习技术,帮助开发者提高编程效率、代码质量和维护性,其应用覆盖多个方面,包括代码补全、代码搜索、文档生成、程序合成、类型推断、漏洞检测等。
程序代码的局部性(Locality)是指程序的语义和逻辑通常受限于局部的上下文环境。例如,函数内的局部变量和参数仅在函数体内具备特定的作用。这种强局部性的特性可以通过图结构来直观地表达。代码中不同部分的相互关系,例如函数调用、变量定义与使用、控制流与数据流等,常常表现为复杂的依赖结构,而图结构能够很好地捕捉这些复杂的关系和层次。例如,抽象语法树 (AST) 用于描述代码的语法结构,控制流图 (CFG) 描述了程序执行的控制流逻辑,数据流图 (DFG) 展示了数据在程序中的流动,程序依赖图 (PDG) 则展示了程序之间控制或数据依赖关系。通过这些图的表示,代码中的局部关系可以以一种结构化的方式被捕捉和分析,从而提高效率和准确性。
本文从可解释性和表达能力两个角度,展示利用图结构团提升代码智能的研究:CFExplainer 和 G2SC,分别在漏洞检测和代码搜索中展示其效果。
二、Graph Neural Networks for Vulnerability Detection: A Counterfactual Explanation(ISSATA 2024)
图神经网络(GNN)因其捕捉源代码底层语义结构的能力,成为了漏洞检测中的一种重要代码嵌入方法。然而,由于其内在的“黑箱”特性,GNN 在可解释性方面面临显著挑战。目前的研究主要集中在基于事实推理的解释方法(factual reasoning-based explainer),其核心思想是识别输入数据中对最终预测起重要作用的关键特征(如代码图中的子图)。然而,这些基于事实推理的解释无法回答这样的 what-if 问题:“如果改变这些代码段的结构,检测系统的决策会发生什么变化?”
基于上述动机,本文提出了 CFExplainer,这是一个新颖的用于GNN 漏洞检测的反事实解释器(counterfactual reasoning based explainer)。CFExplainer 旨在通过对输入代码图的最小扰动来改变检测结果,从而生成反事实解释。与基于事实推理的解释器不同,CFExplainer 能够提供更具行动性的信息,帮助开发者找到漏洞的根本原因,并采取适当的修复措施。
2.1 Code Graph Perturbation
首先,原始代码从纯文本转化为代码图,如ASTs,CFGs,DFGs和PDGs.
扰动的目标:漏洞通常源于源代码中不正确或不一致的结构关系(例如控制流和数据流缺陷)。对于给定的代码图 ,CFExplainer 聚焦于扰动其图结构(即边),而不是节点特征。
边掩码方法:CFExplainer 使用边掩码技术,将寻找反事实扰动的任务转化为一个边掩码的学习任务。具体地,通过对原始代码图的边进行屏蔽,可以生成一个扰动后的代码图 。这一过程通过 来实现,其中 是扰动后的邻接矩阵, 是一个二值边掩码矩阵,表示逐元素相乘。如果 ,表示边在扰动中被移除。
连续化学习:由于直接学习二值边掩码矩阵 不具有可微性,作者将其松弛为连续实数值矩阵 ,然后通过 sigmoid 函数将其映射到 范围内,使得边的存在与否可以进行平滑过渡。最终,边掩码矩 可以通过梯度下降进行优化,从而更快更精确地确定最小的反事实扰动。
2.2 Counterfactual Reasoning Framework
论文构建了一个反事实推理框架,用于生成基于 GNN 的漏洞检测系统预测的解释。该框架的核心思想是识别对代码图的最小扰动,以改变检测系统的预测。这是通过解决一个反事实优化问题来实现的,具体如下所示。
假设已经有一个预先训练好的 GNN 模型(其权重参数 是固定且不可访问的),以及用于目标代码片段 的代码图 。首先将边掩码 应用于代码图 以生成一个扰动图,即 。接下来,如图所示,将原始和扰动后的代码图输入到 GNN 模型中,以生成各自的估计标签:
其中 表示代码图 中节点的特征。为了识别最小的反事实扰动,我们基于反事实推理问题的优化目标来学习边掩码。即优化公式:
这里约束部分确保新的预测 与原始预测 不同,而目标部分旨在确保扰动后的邻接矩阵 尽可能接近原始邻接矩阵 。由于优化公式的目标和约束部分都是不可微的。论文设计了两个可微的损失函数项,以使两个部分分别可优化。
Prediction Loss Item:为了满足优化公式的约束条件,预测损失项鼓励检测系统在原始代码图 被扰动为 时产生不同的预测,具体如下:
这个损失项旨在最小化扰动后的图 保持原始预测 的可能性,从而最大化实现改变预测结果的机会。
Distance Loss Item:为了解决优化公式的目标部分,距离损失项利用二元交叉熵作为可微距离函数,以量化原始和扰动后的邻接矩阵之间的差异,公式如下:
选择这个距离函数是因为它在衡量两个概率分布的差异方面的有效性。将图中的边的存在和不存在视为二元类别,在优化过程中, 确保 尽可能接近 ,从而确定对代码图 的最小反事实扰动。
总体损失函数:将上述两个损失项集成到一个总体损失函数中,以协同优化它们:
其中 是一个超参数,用于调节预测损失项和距离损失项之间的权衡。
2.3 Counterfactual Explanation Generation
生成最优反事实解释:优化后,得到了最优的边掩码 。在这个掩码矩阵中,较高的值表示其对应的边应当保留,而较低的值表示其对应的边应当被移除,以使得检测系统的决策发生反转。为了形成最终的解释,我们使用一个超参数 来控制要扰动的边的数量,即选取掩码值最低的 条边。接着,我们通过移除选定的 条边得到最优的反事实扰动图 ,并导出一个子图:
因此,最优的反事实解释的形式如下:导出的子图 是对检测结果最关键的部分,若将其移除,则代码不会被预测为有漏洞。
生成多样化的反事实解释:在现实场景中,开发者可能需要多样化的反事实解释,以探索并理解检测漏洞的上下文。为了实现这一点,论文基于子图 构建了一个缩小的搜索空间。在该空间中使用穷尽搜索来系统地探索并筛选各种边组合,这些组合的移除会改变检测系统的预测结果。这个过程生成了一组多样化的反事实解释,每个解释都提供了对检测到的漏洞的不同视角。此外,这种多样性为开发者提供了多个可操作的选择,以应对检测到的漏洞。
实验效果
与其他模型进行比较,结果表明,CFExplainer在大多数场景下都优于基线模型,在准确率、精确率、召回率和F1分数上表现更好。
三、Graph Neural Networks for Vulnerability Detection: A Counterfactual Explanation(FSE 2022)
3.1 背景介绍
语义代码搜索(Semantic Code Search)可以极大地促进软件的复用,帮助用户找到与自然语言查询高度匹配的代码片段。这一任务并不简单,因为自然语言查询中的高层次意图与源代码中的低层次实现细节之间往往存在语义差距。文章举例说明了这一点(Figure 1):例如,查询“如何从 XML 读取对象?”返回的代码片段中可能完全不包含查询中提到的关键词,例如“读取”和“对象”。因此,基于关键字匹配或文本相似度的单模态代码搜索方法(如传统的信息检索方法)无法很好地解决这一挑战。
为缩小这种语义差距,近年来的研究提出了多种基于深度学习的多模态模型。这些模型通过在高维空间中学习查询和代码的语义表示,成功地提高了搜索的有效性。这些工作可以分为两大类:
多模态预训练模型:这类模型基于 Transformer 架构,通过若干预训练任务学习多模态数据的通用表示,然后对一系列代码相关的下游任务(如代码搜索、代码摘要生成等)进行微调。例如,CodeBERT 就是一种典型的预训练模型。
非预训练的多模态模型:这些模型则是专门为代码搜索任务端到端训练的,如 DCS(Deep Code Search)等。
虽然这些模型能够捕捉到更多的语义信息,但它们在代码结构信息的利用上仍存在不足之处。
代码图(如控制流图和程序依赖图)由于其表达能力丰富,已被主流的多模态模型和预训练模型用于代码建模。然而,这些方法仍存在一些局限性:在搜索有效性方面仍有改进空间,同时它们未能充分考虑代码图的独特特性。为此,本文提出了一种图到序列转换器 G2SC,将代码图转换为无损序列,使得通过序列特征学习来处理小图成为可能,并捕获代码图中的节点和边属性信息。通过实验,本文验证了 G2SC 对现有模型的性能提升,尤其是提升代码搜索的准确性。
3.2 代码搜索的主流模型
在代码搜索任务中,现有的解决方案可以分为两类:多模态模型和预训练模型。本文选择了这两类中的四个代表性模型进行比较,分别是 DCS(Deep Code Search)、MMAN(Multi-modal Attention Network)、CodeBERT 和 GraphCodeBERT。
A. 多模态模型
Deep Code Search (DCS):DCS 是最早基于深度学习的代码搜索模型之一。DCS 通过提取代码中的三方面信息来进行代码表示学习:
方法体中的 Token 使用多层感知机(MLP)进行嵌入;
方法名则通过循环神经网络(RNN)进行处理;
API 序列也通过 RNN 进行建模。然而,DCS 没有考虑代码的结构化信息,仅对代码文本进行简单嵌入。
Multi-modal Attention Network (MMAN):MMAN 是一种多模态模型,集成了代码结构信息,如抽象语法树(AST)和控制流图(CFG)。MMAN 使用 TreeLSTM 来嵌入 AST,使用 GGNN(门控图神经网络)来嵌入 CFG,并通过 LSTM 嵌入代码序列,然后通过注意力层将这些特征融合为一个单一向量表示整个代码。然而,MMAN 忽略了代码图中信息的多样性,如控制依赖和数据依赖。
B. 预训练模型
CodeBERT:CodeBERT 是第一种用于代码搜索的预训练模型,由 Transformer 架构支持。它通过两种预训练任务进行训练:掩蔽语言模型(MLM)和替换 Token 检测(RTD),从而学习代码与自然语言之间的通用表示。CodeBERT 并非专门为代码搜索设计,因此需要进行额外的微调。
GraphCodeBERT:GraphCodeBERT 是 CodeBERT 的增强版本,引入了捕捉数据流的图结构特性。除了 CodeBERT 的 MLM 任务之外,GraphCodeBERT 还使用了两个特定的预训练任务:边预测和节点对齐,此外引入了基于图的掩蔽注意力函数来进一步捕捉代码结构信息。然而,GraphCodeBERT 只考虑了变量之间的关系,忽略了代码中其他丰富的语义信息,如控制流关系。
3.3 图到序列转换器(G2SC)
本文提出的 G2SC 的核心理念在于将代码图(如控制流图和程序依赖图)转换为等价的序列表示,以使代码的结构信息可以通过现有的序列学习模型得到更好的捕捉。传统图学习方法,特别是基于图神经网络(GNN)的方法,在处理小规模代码图时往往表现欠佳,而序列特征学习更适合对这些小图进行处理。通过 G2SC,能够将代码图中的节点和边的信息转换为序列,使得代码图特征能够通过如 LSTM 或 Transformer 等序列模型进行学习。
首先,G2SC 使用特定的图遍历策略将代码图中的所有节点以唯一的顺序进行遍历,以生成唯一对应的节点序列。遍历策略的选择至关重要,它必须保证遍历顺序的确定性,即相同的代码图总能得到相同的节点序列。
为了保证转换的正确性和无损性,G2SC 采用了一套精心设计的图遍历策略,确保转换过程是一对一且可逆的。这意味着,G2SC 转换的序列不仅能够完整保留图的结构信息,还能够通过逆向过程重构出原始的代码图。
算法 1 详细介绍了图到序列转换器(Graph-to-Sequence Converter,G2SC)算法的步骤。在第一步中,给定程序依赖图(PDG),需要确定用于生成 PDG 序列的入口节点(entry node)。由于 PDG 只有一个根节点,PDG 的根节点 可以直接作为整个 PDG 序列的入口 (见算法 1 第 1 行)。之后,从节点 开始,递归选择边进行扩展,按照以下规则依次执行:
回溯边优先(回溯边是指边的终点已经被遍历过):回溯边优先于前向边(前向边是指终点尚未被遍历过的边)被选择。 多条边的选择:当存在多条回溯边或多条前向边时,可以选择控制依赖边或数据依赖边优先进行遍历。两者的优先级都会产生唯一的结果。在本文的图到序列算法中,设置控制依赖边具有更高的优先级(见算法 1 第 4-8 行)。 选择优先级最高的边:如果在第二步后仍然无法确定边,则选择终点优先级最高的边进行扩展。在我们的方法中,节点的优先级由其对应代码片段在代码中的顺序决定(见算法 1 中的函数 Edge_Check)。 无连边处理:如果没有连接边,则遍历继续回退到具有至少一条未遍历边的最近的祖先节点(见算法 1 第 9-11 行)。
在遍历 PDG 的过程中,我们可以按照以下两个步骤生成我们提出的 PDG 序列(见算法 1 中的函数 Sequence_Expand):
如果遍历的边 是控制依赖边,则只将由边 指向的节点的属性添加到生成的 PDG 序列 中。 如果遍历的边 是数据依赖边,则将边的属性和由边 指向的节点的属性依次添加到序列 中。
按照算法,可以将Figure 5中的 PDG 转换为其唯一对应的序列。还需要注意的是,图到序列算法的时间复杂度仅为,这意味着它在进行 PDG 预处理中不需要太多时间。
3.4 G2SC Enabled Multi-modal Model
首先,我们介绍 G2SC 支持的多模态模型,即 GSMM。具体来说,GSMM 利用多层感知器(MLP)来学习代码主体(Body Tokens),利用 BiLSTM 学习方法名称和 PDG 序列,并采用注意力机制来结合这三种代码特征。GSMM 与其他多模态模型的主要区别在于 GSMM 使用 G2SC 转换的 PDG 序列来进行代码图结构学习。文本及表示:GSMM 的文本级表示模型与 DCS 类似。它将方法名称和代码主体中的标记作为文本级表示,将其作为token序列,使用BiLSTM结合最大池化(maxpooling)总结为嵌入向量序列。
图级表示:给定代码PDG序列,该序列已经由G2SC处理,进而采用另外一个BiLSTM结合最大池化嵌入它们。
最后,通过注意力融合层将三个方面的向量融合为一个向量:
其中, 表示每个向量的注意力权重,、、 分别表示方法名、代码主体、PDG序列,输出向量 表示代码片段的最终嵌入。
在 GSMM 模型训练过程部分,论文将三元组 输入到模型中。对于每个代码片段 ,有一个正描述 (正确描述 )以及一个负描述 (错误描述 ,从所有描述池中随机选择)。对于给定的 三元组,GSMM 预测 ⟨C,D+⟩ 和 ⟨C,D−⟩ 两对的余弦相似度,并最小化以下合页损失函数:
其中, 表示模型参数; 表示训练数据集; 是机器学习中常用的松弛变量,通常设置为 0.05; , 和 是嵌入向量。排序损失的直觉在于:推动模型预测代码片段与其正确描述之间的高相似性,以及代码片段与错误描述之间的低相似性。
在模型训练完成后,对于具有给定代码库 和查询 的多模态模型,目标是根据代码片段与查询 的相似度对这些代码片段进行排序。我们首先将代码片段 输入到 GSMM 中,并将查询 作为描述输入,以获得它们相应的表示,记为 和 。然后,计算排名得分如下:
其中, 和 分别是代码和查询的向量。相似度越高,代码与查询的语义相关性就越高。
实验效果
Table 1 显示了 GSMM、DCS 和 MMAN 在 CodeSearchNet 上的性能。结果表明,GSMM 在所有评估指标上始终优于 MMAN 和 DCS。
为了进一步研究 GSMM 的性能改进,论文将没有注意机制的 GSMM 版本 GSMM-w/o.A 与其他模型进行比较。从结果中可以看到,即使没有注意机制,GSMM仍然优于 DCS、MMAN 等其他基线。
3.5 G2SC Enabled CodeBERT Model
接下来,我们介绍 G2SC 支持的 CodeBERT 模型,即 GSCodeBERT。和其他预训练模型类似,CodeBERT 包含两个阶段:预训练和微调。由于预训练 CodeBERT 的源代码未发布,论文仅在 CodeBERT 的下游任务中使用 G2SC。
Figure 7 展示了 GSCodeBERT 的整体流程。每个输入代码片段被分解为两对 和 (D,P)。其中,D=⟨d1,⋯,dn⟩ 是从代码片段 的自然语言描述中提取的标记序列 是从程序语言代码序列中提取的标记序列 是从 G2SC 转换的 PDG 序列中提取的标记序列 。所有的 和 对分别被输入到预训练的 CodeBERT 模型中,使用一个特殊的标记 表示 与 (或 )之间的语义相关性作为下游微调任务的输入。这使得 GSCodeBERT 模型在保留 CodeBERT 对自然语言和代码序列理解能力的同时,能够提高对代码结构信息的理解。
GSCodeBERT 使用带有二元分类损失函数的模型进行微调,其中使用 softmax 层将输出向量 ClassLabel 映射到 之间的得分。得分越高表示代码序列 PDG 序列与自然语言之间的相似性越大。GSCodeBERT 的微调任务包括两个部分:一个是描述 与 PDG 序列 之间的相似性,另一个是描述 与代码序列 之间的相似性。二者均使用类似的损失函数,公式如下:
其中, 表示模型参数; 表示训练集; 表示自然语言描述, 可以是代码序列 或 PDG 序列 ; 是松弛变量; 是一个特殊标记,指示 是正样本还是负样本。所有的 对被视为正样本。所有 和 都是负样本,这些负样本具有平衡的实例数量,且通过随机将 中的 或 替换为 或 生成。函数 用于计算相似性。我们的损失函数旨在增加(减少)序列 与其正确(错误)描述 之间的相似性得分。
在模型训练完成后,对于每个查询 ,将 与代码库中所有代码序列 配对输入到模型中。输入的格式化形式为 。在 CodeBERT 处理之后,返回与 对应的标记嵌入。进一步地,我们通过训练好的下游分类器计算相似性,具体如下:
其中, 和 q 分别是代码和查询的序列。 是对应于每对 的特殊标记。 表示下游任务,给定由 CodeBERT 处理的类标签进行评分。此外,相似度越高,代码与查询的匹配程度就越强。
实验效果
Table 4 展示了 CodeBERT、GraphCodeBERT 和论文方法 GSCodeBERT 在CodeSearchNet数据集上的表现比较。结果表明GSCodeBERT各个指标上都优于其他两个模型。潜在的原因可能是CodeBERT在G2SC的支持下,从PDG序列从捕捉到了更丰富的信息如代码片段的控制依赖性和数据依赖性,可以进一步利用代码片段中的结构信息来学习代码片段和自然语言之间更好的映射。
四、总结
本文介绍的两份工作,分别在漏洞检测和代码搜索两个任务中的多个指标上得到了显著提升。图结构及其相关技术为处理代码复杂性提供了有效的手段,同时也为提升模型的可解释性和表达能力带来了全新的可能性。这表明图神经网络不仅是代码理解和改进中的有力工具,更是未来代码智能领域不可或缺的重要方向。