连接拓扑空间与特征空间:一种用于漏洞检测的双通道图神经网络
《Array》:Bridging topological and feature spaces: A dual channel graph neural network for vulnerability detection
【字体:
大
中
小
】
时间:2026年05月16日
来源:Array 4.5
编辑推荐:
安博·黄|海城曲辽宁工程技术大学软件学院,葫芦岛,125000,辽宁,中国
摘要
基于序列的深度学习方法往往难以捕捉源代码的丰富结构特征,而现有的图神经网络(GNN)在建模长距离依赖关系方面能力有限,并且倾向于忽略代码图中特征空间中编码的宝贵信息。为了解决这些限制,我们提出了T
安博·黄|海城曲辽宁工程技术大学软件学院,葫芦岛,125000,辽宁,中国
摘要
基于序列的深度学习方法往往难以捕捉源代码的丰富结构特征,而现有的图神经网络(GNN)在建模长距离依赖关系方面能力有限,并且倾向于忽略代码图中特征空间中编码的宝贵信息。为了解决这些限制,我们提出了TFDG(拓扑-特征双空间图神经网络),这是一种基于GNN的新框架,可以联合建模软件漏洞检测的拓扑结构和特征空间。
在TFDG中,首先使用Joern将源代码转换为CPG(代码程序图),然后使用Word2Vec生成初始节点表示。在拓扑空间通道中,通过TOPK算法逐步简化图结构,该算法利用图卷积网络(GCN)的输出作为节点重要性分数,并相应地更新邻接矩阵和节点特征。这一过程通过降低图复杂性有效地缩短了长距离依赖关系。同时,特征空间通道专注于通过执行特征压缩和多层特征关系建模来提取潜在的语义信息。为了处理大小可变的图,我们引入了一种掩码算法来明确区分有效节点和填充节点,从而防止批量处理过程中由填充引起的噪声。最后,一个门控机制动态融合来自两个通道的表示,以产生最终的检测结果。
在三个公开可用的数据集上进行的广泛实验表明,TFDG在准确性、精确度、召回率和F1分数方面始终优于大多数基于GNN的基线方法。进一步分析表明,拓扑空间和特征空间通道捕获了互补的信息,使得漏洞检测更加稳健和全面。这些结果表明,双空间建模为解决现有基于图的漏洞检测方法的局限性提供了有效的解决方案。
1. 引言
截至2023年,像GitHub这样的平台已经吸引了超过1亿用户,托管了大量的开源项目。根据Synopsys Black Duck在2024年对1067个代码仓库进行的审计,96%的分析项目依赖于开源组件。令人担忧的是,84%的仓库至少包含一个已知的开源漏洞,74%的仓库受到高风险漏洞的影响,与2022年相比有显著增加。这些漏洞不仅仅是理论上的安全缺陷;它们通常是严重安全事件的初始攻击向量,包括大规模数据泄露。随着人工智能技术的快速发展、海量数据的广泛传播以及日益严格的全球数据治理法规,数据已成为一种战略资产,其安全性至关重要。根据IBM的数据,2023年全球数据泄露的平均成本从445万美元上升到2024年的488万美元,增长了10%。这些统计数据显示,软件漏洞——尤其是那些隐藏在广泛重用的开源代码中的漏洞——可以放大级联的安全风险,最终导致巨大的财务损失和对国家安全、企业及个人用户的严重威胁。因此,有效和准确的漏洞检测是防止数据泄露和保护现代软件生态系统的重要前提。
基于检测过程中使用的技术,源代码漏洞检测方法大致可以分为两类:基于规则的漏洞检测方法和基于学习的漏洞检测方法。
基于规则的源代码漏洞检测方法(例如某些开源或商业漏洞检测工具)依赖于专家定义的规则。然而,专家的强主观性导致这些规则存在局限性和不完整性,从而导致误报或遗漏漏洞。此外,随着安全漏洞数量的不断增加,这些方法通常需要大量的手动工作,增加了劳动力成本。
相比之下,基于学习的漏洞检测可以从历史数据中自动学习漏洞模式,而无需依赖专家知识。深度学习在计算机视觉(CV)和自然语言处理(NLP)等领域取得了显著的成功,为漏洞检测提供了宝贵的见解。例如,由于编程语言(PL)和自然语言(NL)都依赖于语法规则来组织语句/代码并表现出上下文依赖性,一些研究人员利用NLP技术进行漏洞检测。然而,这些基于NLP的方法往往表现不佳,因为它们无法充分建模源代码中固有的丰富程序语义和结构信息。
为了解决这一限制,一些研究人员转向程序分析,从源代码中提取程序语义并将其表示为图结构,然后进行基于图的漏洞检测。与由规则排列在规则2D/3D网格中的图像或由离散的1D单词、字符或子词组成的NL不同,图数据是由节点(实体)和边(关系)组成的非欧几里得结构,具有不规则和复杂的关系模式。这些属性使图能够准确表示源代码,从而提高检测效率。
随着大规模预训练语言模型(如CodeBERT)和其他基于Transformer的模型的快速发展,在NLP和代码理解任务方面取得了显著进展。这些模型擅长从大量代码数据中捕获丰富的语义信息,将程序视为令牌序列。然而,尽管它们具有强大的表示能力,基于序列的语言模型在显式建模程序结构方面仍然存在固有的限制。源代码的语义不仅由令牌级别的上下文决定,还受到控制流、数据流和依赖模式等结构关系的严重影响,这些关系自然地以图的形式表达。基于图的方法通过提供明确的结构归纳偏差来解决这一限制。此外,基于图的方法并不与现代NLP技术相矛盾。相反,预训练的语言模型可以用来初始化程序图中的节点特征表示,从而将强大的语义先验注入图学习过程。通过这种方式,从大规模代码语料库中学到的语义信息在嵌入阶段被纳入,而GNN通过利用节点之间的结构关系进一步细化这些表示。这种组合使基于图的方法能够联合建模文本语义和程序结构,提供更全面的表示。然而,现有的基于图的漏洞检测方法仍然存在固有的建模限制。首先,由于GNN中的消息传递机制,浅层GNN架构难以捕捉大型程序图中远距离节点之间的长距离依赖关系。其次,大多数现有的基于GNN的检测方法主要关注利用图拓扑,而来自源代码图的特征空间的内在特性尚未得到充分探索,这限制了节点表示的表达能力并进一步影响了检测性能。
这些限制突显了需要更有效的漏洞检测方法,这些方法能够联合捕捉程序图中的结构依赖关系和特征空间关系。具体来说,它们提出了两个重要的研究问题:(1)如何有效缓解由大型和复杂程序图引起的长距离依赖问题;(2)如何更好地利用来自源代码图的特征空间信息,而这通常是现有以拓扑为中心的方法所忽略的。解决这些问题对于提高基于图的漏洞检测模型的表示能力至关重要。
为了解决上述挑战,我们提出了TFDG,这是一种新的漏洞检测框架,它从拓扑和特征空间的角度联合建模程序图。遵循Devign的方法,首先使用Joern将源代码转换为CPG,然后通过Word2Vec初始化节点表示。与主要关注CPG上结构消息传递的现有方法不同,TFDG将拓扑空间学习和特征空间建模明确分离为两个互补的通道。在拓扑空间通道中,TFDG使用基于GCN的TopK算法根据学习到的重要性分数选择信息节点来简化图结构。这一过程有效地减少了图规模并缩短了节点距离,从而缓解了传统GNN在建模长距离依赖关系方面的限制。同时,TFDG引入了一个专门的特征空间通道,直接操作节点特征矩阵,旨在捕捉拓扑空间通道所忽略的潜在语义关系。该通道执行特征压缩和多层特征关系建模,使得独立于图连接性更深入地利用特征信息。为了处理大小可变的图,我们引入了一种掩码算法来明确区分有效节点和填充节点,从而防止批量处理过程中由填充引起的噪声。最后,采用了一个门控机制来动态整合拓扑空间和特征空间通道的输出,使TFDG能够动态平衡结构和语义信息,用于漏洞检测。
本文的主要贡献总结如下:
- 我们提出了TFDG,这是一种基于双空间概念设计的新型漏洞检测框架,它从拓扑空间和特征空间两个角度联合建模漏洞检测。这种设计超越了现有的以拓扑为中心的GNN方法,明确捕获了结构依赖关系和特征级语义关系。
- 我们在拓扑空间中设计了一种基于GCN的TopK图简化策略,以减少图规模并缩短有效节点距离,从而缓解了在大型程序图中建模长距离依赖关系的难度。
- 我们引入了一个专门的特征空间通道,直接操作节点特征矩阵以捕获潜在的语义关系。为了支持这一过程,我们进一步开发了一种掩码算法来处理大小可变的图。实验结果还表明,将所提出的特征空间建模策略纳入基于拓扑的模型中可以持续提高其性能。
- 在三个公共基准数据集上进行的广泛实验表明,所提出的TFDG框架在F1分数方面与几种代表性的基线方法相比具有竞争力,验证了所提出的双空间设计的有效性。
2. 相关工作
近年来,深度学习技术在漏洞检测领域取得了显著进展。现有方法可以根据其模型架构大致分为两类:基于序列程序表示的模型和基于抽象语法树或图程序表示的模型。
2.1 基于序列程序表示的深度学习模型
这些方法通常将源代码视为序列数据,并利用各种神经网络架构进行特征学习和漏洞检测。因此,首先需要将源代码转换为适当的序列表示,以有效建模代码中固有的丰富语法和语义信息。一些研究人员提出了将源代码或代码段表示为令牌序列的方法。这些方法通常使用N-gram模型或Word2Vec将代码令牌转换为向量表示,然后将其输入到序列神经网络中。例如,Russell等人使用Word2Vec进行代码嵌入,并分别使用CNN和RNN架构提取局部和全局特征。同样,Zhen Li等人提出了VulDeepecker、SySeVR和μVulDeePecker,它们使用各种方法将源代码分割成包含不同级别语法和语义信息的段,并使用RNN的变体(如GRU、LSTM)作为提取漏洞特征的主干网络。
为了更好地捕获代码中的上下文信息,许多研究采用了双向架构,如BiGRU或BiLSTM进行代码表示学习。与单向对应物相比,BiGRU和BiLSTM结合了前向和后向层,在每个时间步骤基于两个方向的组合隐藏状态产生输出。这种架构在学习代码序列中的上下文依赖关系(包括前后关系)方面特别有效。
最近在基于深度学习的漏洞检测方面的进展突显了基于Transformer的架构的潜力,因为它们在建模长距离依赖关系方面具有出色的能力。Thapa等人进行了系统实验,证明像BERT和CodeBert这样的Transformer模型在使用序列化程序表示的C/C++漏洞检测任务中优于传统的RNN变体(如BiLSTM和BiGRU),主要是因为它们在捕获代码中的长距离依赖关系方面具有更好的能力。在此基础上,一些研究人员利用大型语言模型和基于Transformer的改进的词嵌入来增强代码理解,从而提高了特征表示能力并取得了优异的结果。
尽管这些方法取得了显著的成功,但它们仍存在固有的限制:PL包含比简单令牌序列更严格的逻辑结构。例如,循环内外的语句可能在序列上相邻,但在语义上无关。总之,虽然基于序列的深度学习模型在学习序列语义特征和建模令牌级依赖关系方面有效,但它们从根本上无法捕获源代码中固有的结构语义和语法信息。
2.2基于抽象语法树/图程序表示的深度学习模型
基于抽象语法树/图程序表示的深度学习模型指的是首先将源代码转换为结构化表示,然后使用相应的神经架构进行漏洞检测的方法。这些方法与基于序列的方法根本不同,因为它们明确保留并利用了源代码中嵌入的丰富结构和语义信息。基于语法树的代码表示方法将源代码解析为抽象语法树(AST),并对其进行建模,有效地保留了语法结构信息,并已广泛应用于各种与语义相关的任务中[5]。然而,实证证据[9]、[30]、[?]表明,尽管这些方法能够完全捕获代码的语法结构,但它们在表示程序数据依赖关系的控制逻辑以及漏洞检测的关键方面表现不佳。因此,仅依赖原始AST作为中间代码表示通常会导致检测性能不佳。为了解决这一限制,最近的研究趋势侧重于通过添加程序分析信息(如控制流图(CFG)和数据流图(DFG)来丰富代码表示,从而提高检测的准确性和鲁棒性。
经典的方法是将CFG和DFG集成到AST中(即构建CPG)。例如,Zhou等人[9]提出了使用门控图神经网络(GGNN)[31],结合GRU单元来加强CPG中结构和序列信息的联合学习。Wang等人[32]进一步发展了这种方法,使用关系门控图神经网络(R-GGNN)[33]来更有效地处理代码图中的异构关系。图学习中的基础模型GCN[34]经常被用作漏洞检测的基线方法[35]、[36]。许多研究[9]、[37]结合了GCN和CPG,最近,Cho Do Xuan等人[38]结合了CPG、GCN和数据重新平衡技术,在Verum[39]和FFMPeg+Qemu数据集上取得了最先进(SOTA)的结果[38]。
最近,图注意力网络(GATs)[40]因其出色的选择性特征聚合能力而越来越受到青睐。此外,随着(大型语言模型)LLM的发展,一些研究人员[41]、[42]在基于图的漏洞检测的嵌入阶段使用了LLM,然后利用GNN进行特征提取和检测,取得了相当令人印象深刻的结果。总的来说,这些研究为基于图的漏洞检测研究提供了宝贵的参考。这些方法直接对代码的结构语义进行了建模。然而,它们在捕获长距离依赖关系方面仍然面临挑战,并且忽略了代码图中编码的宝贵信息,这可能在某些情况下限制了它们的有效性。
3. 方法
在本节中,我们首先定义问题,然后详细描述TFDG的结构和细节。TFDG的框架如图1所示。
下载:下载高分辨率图像(589KB)
下载:下载全尺寸图像
图1. TFDG的架构图
3.1. 问题定义
我们将函数级别的源代码漏洞检测视为一个二分类问题,其中分类结果0表示图的相应源代码不包含漏洞,1表示图的相应源代码包含漏洞。受到Devign[9]的启发,我们使用Joern[11]来获取源代码的CPG。数据集定义为{(ci,yi)∣ci∈C,yi∈Y}i=1N,其中C和Y={0,1}分别表示源代码集合及其标签。N表示集合的数量。然后,源代码C通过Joern转换为CPG G。我们的目标是通过最小化损失函数f:G?Y来学习这个映射,损失函数如下:(1)min∑L(f(Gi,yi∣ci)),其中L是交叉熵损失函数。
3.2. 框架
本节详细介绍了TFDG的架构,它包括四个关键组成部分:(1) 源代码处理部分,将源代码转换为模型兼容的输入表示;(2) 拓扑空间通道,通过TOPK算法显式处理图拓扑结构,执行GCN和简化操作以优化拓扑结构并高效提取长距离依赖关系;然后使用不同的池化方法在不同尺度上获取信息,并使用残差连接不同简化层的结果以生成图的全局表示;(3) 特征空间通道,通过压缩特征和建模多层特征关系来提取潜在的语义特征;同时,使用掩码算法在处理变长图数据时避免填充节点和引入噪声;(4) 读出模块,利用门控机制动态调整两个通道的输出以获得检测结果。
3.2.1. 源代码处理
在本节中,我们首先使用Joern[11]将源代码转换为适合TFDG处理的图结构表示。受到Devign[9]的启发,我们采用CPG来全面编码源代码的结构和语义,其中代码实体(例如类、函数和变量)表示为节点,程序语义(例如数据流和控制流关系)表示为边。所有层次关系都被整合到一个统一的图表示中。值得注意的是,这个图不是直接从原始文本构建的。相反,CPG中的边显式定义了邻接关系,这些边是基于确定性程序分析规则生成的。将源代码转换为CPG后,我们使用Word2Vec[12]生成一个节点特征矩阵X∈Rn×d,其中n表示图中的节点数量,d表示每个节点的特征维度。然后,根据CPG中定义的连接关系构建一个邻接矩阵A∈Rn×n,其中每个元素表示两个节点之间是否存在边以及边的类型。
3.2.2. 拓扑空间通道
拓扑空间通道采用基于重要性分数的TopK算法,这些重要性分数来自GCN的输出。我们选择GCN作为该通道的骨干网络,因为其结构简单且聚合行为稳定,非常适合生成可靠的节点重要性分数。在重要性分数的驱动下,通道通过各种池化方法在不同层次上保留特征信息,从而有效提取拓扑特征。然后,通过多层设计和残差连接,逐步优化拓扑特征以获得最终的图表示。由于提出的TopK算法显著减少了节点数量并缩短了有效图距离,它在很大程度上缓解了GCN在建模长距离依赖关系方面的局限性。这种设计在简化过程中保持了关键的拓扑结构,同时保持了网络层之间的特征多样性。TopK算法如表1所示。
我们首先将拓扑图G=(A,X)输入到GCN中,其中A∈Rn×n是图的邻接矩阵,X∈Rn×d是图的节点特征矩阵。GCN的第l+1层的输出可以表示为:(2)H(l+1)=α(D??12A?D??12H(l)),其中α(.)表示ReLU激活函数,A?∈Rn×n是图的邻接矩阵加上单位矩阵,D?∈Rn×n是A?的度矩阵,H(l)是GCN的第l层的输出,W(l)是GCN的第l层的权重。D??12A?D??12和H(l)W(l)分别用于图归一化和特征转换。
表1. 拓扑空间通道中的TOPK算法。
当输出维度设置为1时,GCN在激活函数后为每个节点产生一个标量值,得到输出S∈R1×n,这个输出可以解释为节点重要性分数,因为GCN从局部邻域聚合信息并编码了节点特征和拓扑上下文。这样,单层GCN就足以学习单个节点的重要性及其局部结构信息。(3)S=Tanh(H(l+1)),其中Tanh被用作激活函数。获得节点重要性分数S后,通过降序排序分数得到相应的排名指数I:(4)I=Argsort(S,descending),其中Argsort(.)表示排序函数。根据排名指数I,确定要保留的节点数量为:(5)K=max(1,?ρn?),其中ρ∈(0,1]表示图中要保留的节点比例,K是要保留的节点数量。基于k,我们生成一个二进制掩码Mask∈R1×n:(6)Mask∈{0,1}n (7)∑i=1nMaski=k,其中Mask表示节点是保留(1)还是丢弃(0)。设I=i∣Maski=1表示保留的节点索引集,|I|=k。节点特征和邻接矩阵更新为:(8)X′=X[I],A′=A[I,I],其中X′∈Rk×d,A′∈Rk×k。在图简化后,应用全局读出操作来聚合图级表示:(9)g1=concat(avg(X′),max(X′)),其中avg(?)捕获节点特征的整体分布,max(?)突出最显著的特征。在后续的GCN层及其相应的池化步骤中,进一步处理简化后的图(A′,X′)以获得另一个图级表示g2。最后,将多尺度表示融合并输入到分类头中:(10)Z1=MLP(g1+g2),其中MLP表示线性分类头,Z1∈(0,1)是拓扑空间通道的输出。
3.2.3. 特征空间通道
特征空间通道通过三种关键机制增强特征表示能力:掩码机制、特征映射和层次特征关系建模。为了实现可变大小图的高效批量训练,我们实施了一种结构化的填充策略:对于每个小批量,所有图都进行零填充以匹配批量中最大图的维度。然后应用掩码算法,将填充节点从特征聚合中排除,通过逐元素掩码保持计算正确性,并确保梯度传播仅限于有效节点。掩码算法如表2所示。
其中feature是批量图的张量列表,M是批量中图的最大节点数,padded_value是需要填充到图中的值,padded_features∈(B,M,D)表示批量大小、图的最大节点数和节点特征维度。
表2. 特征空间通道中的掩码算法。
节点数少于M的图沿节点维度用padding_value进行填充,并生成一个mask∈RB×M。在mask中,填充位置设置为0,有效节点位置设置为1,生成最终的padded_features。
首先计算每个图的特征以在特征空间通道中获得全局图表示。(11)y=(Xi?(maski.unsqueeze(?1))).sum(1)∑maski,这里Xi∈RM×D表示第i个图的填充节点特征矩阵,maski∈RM表示批量图中第i个图的掩码,maski∈RM通过unsqueeze(?1)转换为maski∈RM×1,Xi?(maski.unsqueeze(?1))∈RM×D实现了图中节点特征和掩码的逐节点乘法,仅保留有效节点特征,而∑maski表示用于归一化的有效节点数量。随后,沿节点维度聚合掩码后的节点特征并通过有效节点数量进行归一化,以获得图的全局语义表示。后续操作在纯化的全球特征y∈R1×D上进行,这排除了权重分配过程中填充节点的任何影响。接下来通过线性层进行特征压缩以获得h:(12)h=yW1,其中W1∈RD×d′是可学习参数,d′表示隐藏维度,h∈R1×d′。然后进行线性矩阵组合。(13)A1=Expand(Softmax(W2h)),其中Expand(?)扩展向量以实现所有成对特征维度的动态调整。(14)A=(I⊙A1)+A2,其中⊙表示逐元素乘法,W2∈Rd′×d′是可学习参数,Softmax函数确保权重归一化,Expand扩展维度以实现特征空间中所有维度对的动态调整。操作I?A1在保持原始特征结构的同时根据学习到的权重进行适应性调整,其中A1∈Rd′×d′,I∈Rd′×d′是单位矩阵,确保基本特征的保留。引入可学习的偏置项A2∈Rd′×d′通过适应性特征调整增强了模型的灵活性。在线性矩阵组合之后,进行特征重构和加权。(15)h′=hA (16)Xout=Xi?σ(h′W3),其中h′∈R1×d′为后续映射提供更具表现力的低维特征,W3∈Rd′×d将特征映射到原始维度,σ是激活函数,将权重限制在[0,1]范围内,最后将加权特征输入分类头以产生最终预测。(17)Z2=MLP(Xout),其中MLP表示用作分类头的多层感知器,Z2∈(0,1)表示来自特征空间通道的预测漏洞概率。
3.2.4.读出模块
在获得拓扑空间通道和特征空间通道的输出后,使用门控机制动态调整以获得模型的最终预测结果。(18)Z=αZ1+(1?α)Z2,其中α是一个可学习的参数,用于动态调整两个通道的输出,Z是模型的最终输出。
4. 实验设置
在本节中,我们提出了研究问题、数据集信息、基线以及实验的相关设置。
4.1. 研究问题
为了评估所提出的TFDG框架是否能够有效解决现有基于GNN的漏洞检测方法的局限性——特别是捕捉长距离依赖关系的难度以及在代码图中 insufficient 利用特征空间信息的问题——我们设计了五个研究问题来评估该框架的不同方面,包括整体检测性能、特征空间通道的贡献以及双空间表示在建模复杂程序图中的有效性。我们的评估旨在回答以下五个研究问题:
• RQ1:TFDG与现有基线的性能相比如何?
RQ1旨在评估所提出模型的整体检测性能与现有方法相比如何。
• RQ2:拓扑空间通道和特征空间通道如何贡献于检测性能?
RQ2旨在调查所提出的特征空间通道是否在双空间设计下对模型性能有所贡献,以及从这个通道提取的特征是否为从拓扑空间通道提取的特征提供了补充信息。
• RQ3:掩码算法如何影响检测结果?
RQ3旨在评估所提出的掩码算法对漏洞检测性能的影响。
• RQ4:使用门控机制融合如何影响检测结果?
RQ4旨在评估所提出的基于门控的融合机制对漏洞检测性能的影响。
• RQ5:超参数如何影响TFDG的检测性能?
RQ5检查不同的超参数设置如何影响TFDG的检测性能。
4.2. 数据集
为了评估TFDG的有效性,我们遵循了AMPLE [8] 的工作,在三个来自实际项目的数据集上进行了TFDG的评估。所有数据集都按照8:1:1的比例进行划分。
FFMPeg+Qemu数据集是从四个大型开源C项目(Linux Kernel、Qemu、Wireshark和FFmpeg)中收集并手动注释的。它包含12,460个易受攻击的函数和14,858个非易受攻击的函数,形成了一个相对平衡的数据集[8]。为了透明度和可重复性,该数据集可以通过此链接公开访问。
Saikat等人[43]从Chromium和Debian项目的开发者或用户报告的漏洞中编译了一个真实世界的数据集。该数据集包括2240个易受攻击的函数和20,494个非易受攻击的函数,更准确地反映了真实世界软件中易受攻击代码的分布。为了透明度和可重复性,该数据集可以通过此链接公开访问。
Fan等人[44]分析了2002年至2019年的348个开源GitHub项目,收集了包含91种不同类型常见弱点枚举(CWEs)的188,636个C/C++函数的数据集。为了透明度和可重复性,该数据集可以通过此链接公开访问。
表3. TFDG的运行时环境配置
配置名称 发布版本
CPU Intel(R) Xeon(R) Gold 6133 CPU @ 2.50 GHz
GPU NVIDIA RTX 3090 24GB
RAM 64GB
硬盘 512GB M.2 SSD
操作系统 Ubuntu 18.04桌面工作站
Python 3.9.2
pytorch 2.2.1
Dgl 2.0.0
4.3. 基线
(1) VulDeePecker [18]:VulDeePecker将源代码转换为代码小工具,并使用Word2vec生成向量表示,然后使用BilSTM学习代码的上下文特征。
(2) Russell [13]:Russell将源代码视为一系列令牌,并使用Word2vec生成向量表示,然后使用CNN提取局部特征和模式。
(3) SySeVR [19]:SySeVR使用BiLSTM学习代码的上下文语义特征。与VulDeePecker相比,SySeVR更关注代码的语义完整性以及对漏洞上下文的系统分析。
(4) Devign [9]:Devign是一种基于GNN的自动化漏洞检测方法,通过分析CPG来识别软件中的潜在安全漏洞。
(5) Reveal [43]:Reveal使用GGNN作为核心模型,将源代码转换为CPG,以便在图的节点和边上促进信息传播和特征提取,以检测代码中的潜在安全漏洞。
(6) AMPLE [8]:一种代码表示简化方法,用于解决节点的长距离依赖问题,并利用GNN增强图特征学习。
(7) EFVD [45]:一种漏洞检测模型,结合了基于序列的CodeBERT进行全局代码语义分析,以及基于边缘注意力的GGNN来动态融合异构图结构,以捕获局部和远程依赖关系,同时使用焦点损失来解决数据不平衡问题。
(8) DGVD [46]:通过结合双图神经网络和卷积特征增强,有效提高了漏洞检测的准确性。
(9) Vul-LMGNNs [42]:一种混合漏洞检测框架,将预训练的代码语言模型与图神经网络集成在代码属性图上。它采用在线知识蒸馏和联合隐式-显式融合来促进语义和结构信息的跨层传播。
(10) GRD [41]:一种基于CPG的漏洞检测模型,结合了基于GCN的结构特征提取和基于RoBERTa的语义建模,并通过Dropout应用数据重新平衡来减轻类别不平衡并提高分类性能。
4.4. 实施细节
具体的实验环境总结在表3中。为了确保不同方法之间的公平和一致评估,所有实验都遵循相同的数据划分策略,其中训练集、验证集和测试集的比例为8:1:1。这种实验场景旨在评估TFDG的整体性能,并在相同条件下与现有基线方法进行直接比较。批量大小设置为64,采用RADAM优化器,学习率为0.0001,最大训练周期固定为100。这些设置旨在平衡训练稳定性和计算效率,并在所有实验场景中一致应用,以确保可重复性和公平性。此外,为了便于可重复性和透明度,本研究中使用的源代码和实验配置可以在https://github.com/habtop3/TFDG-sourcecode公开获取。
我们在评估中使用了以下四个广泛使用的性能指标:
(19) 准确率 = TP + TN
准确率表示正确分类的样本数占总样本数的比例。
(20) 精确度 = TP / (TP + FP)
精确度表示TFDG不将负样本分类为正样本的能力。
(21) 召回率 = TP + FN
召回率表示模型预测为正样本的真实阳性案例的比例。
(22) F1分数 = (精确度 × 召回率) / 2
F1分数是精确度和召回率的调和平均值;F1分数的值越大,模型的鲁棒性越高(见表4)。
表4. TFDG在不同数据集和指标上的性能与基线比较。
5. 实验结果
针对第4.1节提出的研究问题,本章进行了相关实验并提供了实验结果的分析。所有模型都在相同的实验条件下进行训练和评估,包括相同的数据划分策略和超参数设置,以确保公平比较。训练集、验证集和测试集的比例为8:1:1。
5.1. RQ1:TFDG与现有基线的性能相比如何?
本节旨在验证TFDG在不同数据集和评估指标上是否比现有方法具有更好的性能。在基线中,我们在相同的实验条件下重新实现了Devign [9]、AMPLE [8]和DGVD [46],以确保公平比较。对于其余方法,我们采用了AMPLE [8]论文中报告的结果或其原始出版物中的结果。
从基于CPG的图学习方法的角度来看,Devign [9]、GRD [41]和DGVD [46]代表了三种代表性的基于图的漏洞检测方法。Devign将源代码转换为CPG并应用GNN来传播结构信息,这证明了基于图的建模在漏洞检测中的有效性。然而,Devign依赖于单一的拓扑空间并采用标准的消息传递机制,难以有效捕捉大型和复杂程序图中的长距离依赖关系。GRD [41]通过在图嵌入阶段应用NLP技术进一步增强了图表示学习,但它仍然主要关注拓扑信息,并没有明确利用嵌入在特征空间中的潜在语义。最近,DGVD [46]通过结合双图神经网络和基于卷积的特征增强机制进一步改进了基于图的漏洞检测。这种设计增强了模型捕捉程序图中的结构模式和语义信号的能力,从而在多个基准测试中提高了检测准确性。
从基于序列和NLP驱动的方法的角度来看,Vul-LMGNNs [42]和EFVD [45]代表了利用大型预训练语言模型进行漏洞检测的最新方法。Vul-LMGNNs [42]利用强大的LLM从源代码序列中捕获细粒度的语义模式,实现了较高的精确度。同样,EFVD [45]将图神经网络与大型预训练语言模型结合。这种设计增强了源代码的语义表示,并使模型在漏洞检测过程中能够捕获更丰富的上下文信息。这表明基于NLP的方法在学习令牌级语义方面具有优势。然而,这些方法本质上将源代码视为线性序列,缺乏明确建模结构依赖关系(如控制流和数据流关系)的机制,这对于漏洞检测至关重要。因此,当漏洞与复杂的程序结构密切相关时,它们的性能可能会下降。
相反,TFDG同时关注这两种类型的信息,在拓扑空间通道中,TFDG结合了GCN [34]和TOPK算法来有效简化图结构。具体来说,TOPK算法选择性地保留最有信息的节点,同时丢弃不太相关的节点,从而显著减少了图中的节点总数。因此,剩余节点之间的有效最短路径距离缩短,将原本的长距离依赖关系转化为更易于管理的短距离交互。这种结构简化使GCN能够更有效地聚合信息,并缓解了其在建模大型和复杂程序图中的长距离依赖关系时的固有局限性[8]。
为了进一步解决第二个挑战——即现有方法很少仅从特征空间探索潜在表示——TFDG引入了特征空间通道。该通道通过特征压缩和多层特征关系建模捕获节点属性之间的潜在语义关系,从而更深入地利用嵌入在节点特征矩阵中的信息。此外,还结合了掩码算法来有效消除处理不同大小图时引入的噪声,确保学习到的特征表示不会被无关或人为值污染。通过独立于图拓扑进行特征空间学习,该通道补充了拓扑空间通道,使TFDG能够更好地建模之前方法忽略的特征级语义。
值得注意的是,TFDG在所有三个数据集上都优于经典基线方法Devign [9]。在第一个数据集上,TFDG取得了显著的改进:准确率提高了7.06%,精确度提高了6.27%,召回率提高了9.65%,F1分数提高了7.73%。与GRD [41]和Vul-LMGNNs [42]相比,TFDG在准确率和F1分数方面取得了更好的结果。然而,在精确度和召回率指标上,TFDG的表现不如这两种方法。
这种差异可能与不同模型采用的节点嵌入策略有关,实验结果见附录。在我们的框架中,TFDG使用Word2Vec生成节点特征表示,而一些最新方法利用从大型预训练语言模型派生的语义表示,这可以为源代码提供更丰富的上下文信息。为了进一步研究嵌入策略的影响,我们还通过用CodeBERT [6](一种广泛使用的预训练代码表示模型)替换原始的Word2Vec嵌入进行了实验,同时保持其余架构不变。结果表明,TFDG可以从更强的嵌入表示中受益,并提高检测性能。
总体而言,比较结果表明,仅基于图的方法或仅基于序列的方法都不足以完全捕捉源代码漏洞的复杂特征。
这些观察进一步激发了TFDG的概念设计。具体来说,所提出的双空间框架明确地建模了程序图的拓扑结构和特征空间关系,实现了更全面的漏洞表示学习。通过结合两种范式的优势并减轻各自的局限性,TFDG在各种数据集上实现了更强大和稳定的漏洞检测性能。
5.2. 拓扑空间通道和特征空间通道如何贡献于检测性能?
本节旨在验证所提出的双通道架构的有效性,并展示每个通道如何分别有助于解决长距离依赖关系建模和特征空间表示学习。为了研究TFDG中每个组件的贡献,我们通过选择性地移除或修改拓扑空间通道和特征空间通道来进行消融实验。具体来说,我们评估了三种配置:仅使用拓扑空间通道的模型、仅使用特征空间通道的模型,以及同时使用两种通道进行表示学习的模型。结果展示在表5中。从表5可以看出,单独使用任一通道时,模型的准确性、精确度、召回率和F1分数都低于双通道模型。这表明拓扑空间通道和特征空间通道提供了互补的视角。它们的整合显著增强了模型区分易受攻击代码样本和不易受攻击代码样本的能力。为了进一步验证特征空间通道的有效性,我们将其扩展到了Devign [9],这是一个最初仅依赖拓扑信息的基线模型。如表6所示,添加特征空间通道后,Devign [9]在准确性、精确度、召回率和F1分数上分别提高了4.34%、3.31%、1.48%和2.53%,突显了特征空间通道带来的额外区分能力。为了研究两种通道之间的语义差异和互补性,我们对它们各自的嵌入分布进行了可视化分析。如图2所示,使用t-SNE将两个空间的高维嵌入投影到二维空间中。图表显示了明显的聚类模式和显著的分离,表明这两种通道以不同的分布结构编码信息,因此无法展示出对齐的聚类趋势。此外,我们计算了相同样本的拓扑空间通道和特征空间通道嵌入之间的余弦相似度,并将结果以热图的形式展示在图3中。低相似度值和缺乏强烈的对角线模式表明这两种通道的表示在很大程度上是正交的,进一步强化了它们的语义独立性和互补性。总之,定量评估和定性分析的结果都证实了拓扑空间通道和特征空间通道捕获了输入图的独立且互补的特征。它们解耦但又协同的性质丰富了模型的能力,并显著提高了检测性能。
表5. 在FFMPeg+Qemu数据集上对拓扑空间通道和特征空间通道贡献的消融研究。
模型 准确性 精确度 召回率 F1分数
仅拓扑空间通道 61.32 58.11 58.60 58.35
仅特征空间通道 60.41 57.44 55.55 56.48
TFDG 63.58 58.84 70.70 64.22
表6. Devign与添加了所提出的特征空间通道的Devign在FFMPeg+Qemu数据集上的性能比较。
模型 准确性 精确度 召回率 F1分数
Devign 56.52 52.57 61.05 56.49
Devign+特征空间 60.86 55.88 62.53 59.02
下载:下载高分辨率图像(118KB)
下载:下载全尺寸图像
图2. 拓扑空间和特征空间通道嵌入的t-SNE可视化。
下载:下载高分辨率图像(436KB)
下载:下载全尺寸图像
图3. 拓扑空间和特征空间表示之间的余弦相似度热图。
5.3. RQ3:掩码算法如何影响检测结果?
本节旨在评估掩码机制对漏洞检测性能的影响。掩码机制被引入特征空间通道中,以有效处理在源代码表示中常见的变长图。实际上,图通常会被填充到统一的大小以便进行批量训练,这可能会引入无效或噪声节点。为了解决这个问题,我们的模型结合了掩码算法,该算法明确记录了有效节点的位置。在随后的全局池化和特征空间通道的图级表示生成过程中,只考虑有效节点的嵌入。这确保了模型仅关注有意义的结构和语义信息,同时梯度仅针对信息节点进行计算,从而避免了由填充引起的噪声干扰。为了评估掩码算法的影响,我们在保持所有其他架构和训练设置不变的情况下禁用了掩码功能,并进行了消融研究。模型掩码版本和未掩码版本之间的性能比较总结在表7中。结果显示,包含掩码算法在多个评估指标上都有了一致的改进。这些结果证实了该算法在稳定训练和提升特征空间通道中学习到的表示质量方面发挥了作用。掩码算法的有效性突显了在基于图的模型中处理噪声或填充数据的重要性,特别是在处理具有不同图大小的真实世界代码数据集时。通过这种方式,模型可以更有效地学习到稳健的语义模式,最终有助于实现更准确和可靠的漏洞检测。
表7. 掩码算法对FFMPeg+Qemu数据集检测性能的影响。
模型 准确性 精确度 召回率 F1分数
无掩码 63.12 59.02 66.30 62.44
有掩码 63.58 58.84 70.70 64.22
5.4. RQ4:门控机制融合的使用如何影响检测结果?
本节研究了不同的特征融合策略如何影响TFDG的检测性能。为了回答RQ4,我们对比研究了两种特征融合策略:(1)一种门控机制,用于动态控制拓扑空间通道和特征空间通道嵌入的贡献;(2)一种静态加权求和,其中两个分支的输出通过固定系数组合。结果显示在表8中。结果清楚地表明,与静态融合相比,门控融合策略在准确性、精确度、召回率和F1分数上都有轻微的改进。与将拓扑空间通道和特征空间通道嵌入的贡献视为固定的静态融合不同,门控机制引入了一个可学习的系数,可以动态平衡这两种信息来源。这使得模型能够根据每个输入的特点动态调整其对拓扑空间或特征空间通道的依赖,特别是在两个空间的信息量在不同样本间变化时特别有效。此外,门控机制有助于抑制来自不太相关特征的噪声,并促进互补信息的选择性整合。因此,模型可以形成更具表现力和区分性的图级表示,从而提高分类结果。
表8. 不同特征融合策略对FFMPeg+Qemu数据集上TFDG性能的影响。
融合方法 准确性 精确度 召回率 F1分数
静态 61.32 57.02 66.42 61.36
门控 63.58 58.84 70.70 64.22
5.5. Q5:超参数如何影响TFDG的检测性能?
本节研究了TFDG对关键超参数的敏感性,重点关注拓扑空间通道。具体来说,我们分析了简化层数量和TOPK算法操作中使用的节点保留比例的影响。TFDG的超参数包括拓扑空间通道的简化层和节点保留比例。为了验证超参数的影响,我们进行了相应的实验,如表9所示。为了验证简化层数量的影响,我们进行了三项实验:单层、双层和三层。可以看出,单层和双层之间的F1分数没有显著差异。然而,在单层简化的结果中,召回率远高于准确性,表明模型存在过度检测的现象,即将非漏洞检测为漏洞。这是因为单层简化只能提取相对粗糙的特征,无法区分噪声和漏洞特征。我们发现,当简化层数量为3时,模型性能会下降。这是因为随着简化层数量的增加,许多图中只能保留一个节点,导致丢失太多特征,从而影响模型性能。为了研究节点保留比例的影响,我们将简化层数量固定为2,并进行了两组对比实验:(0.8, 0.5)和(0.7, 0.7)。从表9可以看出,(0.8, 0.5)的组合效果最差。这是因为与漏洞相关的大部分特征存在于代码的控制流和数据流中,而在自然语言序列中只有一小部分特征,保留太多节点的同时也保留了太多与漏洞无关的特征。当节点保留比例为(0.5, 0.8)时,模型实现了渐进式的简化过程,从而提高了检测性能。这些结果支持了早期粗略简化和后期精细保留可以最大化保留重要拓扑特征并去除模型中噪声的直觉。
表9. 特征通道超参数(简化层和节点保留比例)对FFMPeg+Qemu数据集上TFDG性能的影响。
结构 比例 准确性 精确度 召回率 F1分数
单层 0.5 57.53 53.25 66.91 59.30
双层 0.8 0.5 61.49 58.86 55.55 57.16
三层 0.7 0.7 62.00 57.46 60.37 58.88
尽管TFDG的性能很有前景,但仍存在几个潜在的有效性威胁。首先,特征空间通道中使用的节点保留比例是手动指定的。虽然我们的实验表明非对称保留策略可以提高性能,但最佳比例可能因输入图的结构复杂性或标签分布而异。缺乏自适应学习机制可能会限制TFDG在其他领域或未见过的图类型中的泛化能力。其次,我们的模型主要将结构简化过程(在拓扑空间通道中)与语义表示过程(在特征空间通道中)分开。然而,如图3所示,两个空间的特征并不完全相关。忽略这种跨空间相关性可能会导致重要信息的丢失。未来的工作应该探索统一或共享注意力的机制,以便动态捕获和整合两个通道中的共享特征。
7. 结论
在本文中,我们通过设计和实现一种名为TFDG的新双空间模型,解决了现有GNN方法在软件漏洞检测中的局限性。该模型首先将源代码转换为图表示,并从拓扑空间通道和特征空间通道对其进行处理。在拓扑空间通道中,我们使用GCN结合TOPK算法来简化图结构,从而缓解了建模长距离依赖性的挑战。在特征空间通道中,我们执行特征压缩和语义建模,并使用掩码算法来避免在处理变长图数据时引入噪声。最后,门控机制动态融合两个通道的输出以生成最终预测。为了评估我们提出方法的有效性,我们在三个公开可用的数据集上进行了广泛的实验。结果表明,我们的模型优于大多数基线方法。此外,我们观察到从拓扑空间通道和特征空间通道提取的特征是互补的,提高了模型的整体性能。当前实现的一个局限性在于节点特征初始化的选择,它依赖于相对轻量的嵌入。尽管这些嵌入足以证明TFDG的有效性,但它们可能无法充分利用大规模预训练模型捕获的丰富语义知识。结合这样的预训练模型是未来工作的一个有前景的方向,可能会进一步提高TFDG的检测准确性和泛化能力。
CRediT作者贡献声明
黄安博:撰写——原始草稿、可视化、方法论
曲海成:资源、资金获取
生物通微信公众号
生物通新浪微博
今日动态 |
人才市场 |
新技术专栏 |
中国科学人 |
云展台 |
BioHot |
云讲堂直播 |
会展中心 |
特价专栏 |
技术快讯 |
免费试用
版权所有 生物通
Copyright© eBiotrade.com, All Rights Reserved
联系信箱:
粤ICP备09063491号