翻译:VulScribeR: Exploring RAG-based Vulnerability Augmentation with LLMs

VulScribeR:探索基于 RAG 的漏洞增强与 LLMs

VulScribeR: Exploring RAG-based Vulnerability Augmentation with LLMs

![[Pasted image 20250325145540.png]]

摘要

检测漏洞对于软件安全至关重要,然而基于深度学习的漏洞检测器(DLVD)面临数据短缺的问题,这限制了其有效性。数据增强可能缓解数据短缺问题,但对漏洞代码进行增强具有挑战性,需要一种能够保持漏洞特性的生成解决方案。以往的工作仅专注于生成包含单条语句或特定类型漏洞的样本。最近,大型语言模型(LLMs)已被用于解决各种代码生成和理解任务,并取得了令人鼓舞的结果,特别是在与检索增强生成(RAG)结合时。
因此,我们提出了 VulScribeR,这是一种基于 LLM 的新颖解决方案,利用精心设计的提示模板来增强漏洞数据集。更具体地说,我们探索了三种策略,通过 LLMs 对单条语句和多条语句漏洞进行增强,即 Mutation(变异)Injection(注入)Extension(扩展)。我们在三个漏洞数据集和 DLVD 模型上使用两种 LLMs 进行了广泛的评估,结果表明,我们的方法在生成平均 5K 漏洞样本的情况下,相较于两种最先进的方法 Vulgen 和 VGX 以及随机过采样(ROS),在 F1 分数上分别提升了 27.48%、27.93% 和 15.41%;而在生成 15K 漏洞样本的情况下,分别提升了 53.84%、54.10%、69.90% 和 40.93%。我们的方法展示了其在大规模数据增强中的可行性,生成 1K 样本的成本低至仅 1.88 美元。

1 引言

漏洞检测是软件工程中的一项关键任务。近年来,基于深度学习的漏洞检测(DLVD)模型 [7, 16, 27, 28, 31, 49, 58] 因其出色的性能引起了研究界和工业界的广泛关注。然而,这些模型深受以下两个问题的困扰:1)缺乏大规模的数据集2)训练数据集与测试数据集分布之间的差异。数据采样技术 [17, 53] 可以通过增加漏洞样本的数量来平衡数据集,并选择性地添加更多易得的干净样本,在一定程度上缓解第一个问题。然而,数据增强提供了一个更具前景的解决方案,因为它可以通过生成大量多样化数据集同时解决这两个问题。例如,Zhang [57] 和 Liu 等人 [30] 专注于生成特定类型错误或漏洞的样本。类似地,Nong 等人提出了 VulGen [37] 和 VGX [35],用于挖掘单条语句的漏洞模式,并将其注入干净样本以生成漏洞样本。尽管这些方法做出了贡献,但它们仍存在局限性——要么只能生成特定类型的漏洞 [35, 37, 57],要么需要大规模数据集来学习生成漏洞代码的模型或确定注入漏洞的位置 [35-37],这使得它们在现实场景中的应用不够实际。

最近,大型语言模型(LLMs)在代码相关任务中展示了令人鼓舞的结果,例如代码理解 [25, 52]、代码生成 [20] 和漏洞理解 [13, 24]。直观地说,漏洞数据增强是一项需要代码理解能力的任务,特别是对漏洞的理解,以及代码生成能力。LLMs 在代码理解和代码生成方面的强大能力使其非常适合这项任务。

因此,为了解决上述问题并克服以往工作的局限性,我们提出了 VulScribeR,这是一种基于 LLM 的新颖解决方案,利用精心设计的提示模板来增强漏洞数据集。更具体地说,我们设计了三种策略来生成新的漏洞代码样本,分别是 Mutation(变异)Injection(注入)Extension(扩展)

  • Mutation 指通过代码转换改变漏洞代码,同时保持代码的语义和语法正确性,通过提示 LLMs 实现。
  • Injection 提示 LLMs 将漏洞样本中的漏洞片段注入一个干净样本,以创建一个新的漏洞样本。
  • Extension 旨在通过将干净样本的部分逻辑添加到漏洞样本中,扩展漏洞样本,从而提高漏洞可能出现的上下文多样性。

为了生成逼真的代码,我们在 Injection 和 Extension 中采用了检索增强生成(RAG)。具体来说,在 Injection 中,我们检索与即将注入漏洞片段的干净样本相似的漏洞样本。此外,为了增强生成的漏洞样本的多样性,我们采用聚类过程,确保所有聚类中的数据都有机会在 RAG 中被检索到。我们在 Extension 中也采用了类似的 RAG 和聚类方法。

最终,我们使用模糊解析器作为过滤机制,丢弃无效响应(例如空代码样本或存在严重语法错误的样本),以减少生成样本中的噪声。

为了评估 VulScribeR,我们使用两种不同的 LLMs 为每种策略生成了漏洞样本,并评估了每个数据增强方法在三个最先进的 DLVD 模型(包括基于标记和基于图的模型)以及三个常用漏洞数据集上的有效性。我们通过以下研究问题(RQs)进行了各种实验:

  • RQ1:VulScribeR 与最先进的方法相比效果如何?
    • 结果:Injection 和 Extension 的表现远远优于基线和 Mutation,而 Injection 的性能略高于 Extension。例如,Injection 在平均 F1 分数上分别比 NoAug、Vulgen、VGX 和 ROS 高出 30.80%、27.48%、27.93% 和 15.41%。
  • RQ2:RAG 对 VulScribeR 的贡献是什么?
    • 结果:RAG 组件对 Injection 和 Extension 策略有显著贡献。
  • RQ3:生成样本的数量如何影响漏洞检测模型的有效性?
    • 结果:通过 Injection 增加更多漏洞数据有助于提高 DLVD 模型的有效性,而 VulGen、VGX 和随机过采样在生成超过 5K 漏洞样本时未能提升 DLVD 模型的性能。我们的基于 LLM 的方法更适合大规模漏洞数据增强。

总结

我们的主要贡献如下:

  • 据我们所知,我们是第一个探索使用 LLMs 进行漏洞增强的研究团队。我们精心设计了三种具有不同策略的新型提示模板,并提出了一个全面的漏洞增强流水线,可用于大规模漏洞增强(每 1K 样本的成本低至 1.88 美元)。
  • 我们使用两种不同的 LLMs、三种 DLVD 模型和三个数据集进行了广泛的评估,证明了 VulScribeR 相较于最先进的基线(包括最新的最佳技术)的优越性。
  • 我们公开了源代码、实验结果和增强后的数据集,供未来研究使用 [1]。

2 背景与相关工作

2.1 漏洞增强

在实践中,漏洞代码样本稀缺,因此已开发了多种方法 [17, 30, 35–37] 来通过生成此类数据并扩展数据集以提高漏洞检测模型的能力。例如,Nong 等人 [36] 探索了通过神经代码编辑模型 [54] 进行漏洞注入的可行性,这是一种基于深度学习(DL)的模型,经过训练可以通过对程序抽象语法树(AST)引入一组预测性更改,将干净代码转换为漏洞代码。这种方法需要高质量的数据集,因此面临“鸡与蛋”的困境,并且其应用受到限制。Ganz 等人 [17] 提出了 CodeGraphSMOTE,一种通过将 SMOTE [8] 移植到图域中生成新漏洞样本的方法。他们将代码转换为图,并使用图自动编码器将这些图编码到潜在空间中。然后,在潜在空间中应用 SMOTE(即基于插值的采样方法)以生成新的漏洞项。
Nong 等人后来提出了 VulGen [37],它通过挖掘漏洞来收集单条语句的漏洞模式,并使用修改后的 Transformer 模型定位注入位置。VGX [35] 是 VulGen 的改进版本,其中使用了一个更大的漏洞数据集,以覆盖挖掘阶段更广泛的单条语句漏洞,并用语义感知上下文化的 Transformer 替换了定位模型,以更好地预测注入上下文。然而,VulGen 和 VGX 仅针对单条语句的漏洞,并且需要全面的模式挖掘过程。根据 Bigvul [12] 数据集中的统计,单条语句漏洞仅占所有漏洞的不到 40%。

2.2 基于 RAG 和 LLMs 的源代码增强

利用 LLMs 进行源代码增强正变得流行,并已被用于语义代码搜索 [50] 和代码生成 [9] 的数据增强,这两者都高度依赖于 RAG。Wang 等人 [50] 利用 RAG 检索相似的查询-代码对,然后针对每一对,通过预定义规则提示 ChatGPT [38] 分别更改查询和代码。随后,他们使用编码器模型计算增强对的相似性分数,过滤掉低分对。Chen 等人 [9] 利用代码搜索为代码生成增强数据。具体来说,他们检索上下文-函数对(上下文可以是函数头或注释),以填充提示模板,指示 LLM 根据给定的上下文和检索到的函数生成代码。他们的架构包含三个组件:Retriever(检索器)Formulator(公式化器)Generator(生成器)。与先前的研究类似,我们也采用 RAG 来检索相关上下文以增强代码生成;不同之处在于,我们专注于漏洞增强,并采用聚类过程以确保生成数据的多样性。


3 方法论

在本节中,我们将详细阐述我们的方法论 VulScribeR。我们提出了一种基于 RAG 的解决方案,利用精心设计的提示模板和宽松的过滤机制,通过 LLMs 生成漏洞代码片段,从而生成多样化且逼真的漏洞代码片段。我们特别设计了三种提示策略来生成用于数据增强的漏洞代码样本。为此,我们为每种增强策略精心设计了三个提示模板,分别是 Mutation(变异)Injection(注入)Extension(扩展) 模板,具体内容将在第 3.1 节中详细说明。
![[Pasted image 20250325145852.png]]

图 1 展示了这三种策略的工作流程,所有三种增强策略最多可抽象为以下四个组件(如果适用):

  • Retriever(检索器)
    给定输入代码样本(即漏洞样本和干净样本),Retriever 负责从数据库中寻找合适的漏洞和非漏洞样本对,并将其附加到提示中,以为漏洞样本生成提供上下文(如果适用)。Injection 和 Extension 策略使用 Retriever 组件,而 Mutation 不需要 Retriever。

  • Formulator(公式化器)
    对于不同的策略,我们使用相应的提示模板。Formulator 通过填充输入代码样本或将 Retriever 检索到的样本对实例化相应的提示模板(如果适用)。对于 Mutation 提示模板,输入直接来自漏洞样本数据集;而对于其他模板,输入则以 Retriever 提供的干净样本和漏洞样本对的形式出现。

  • Generator(生成器)
    在 Generator 中,我们使用实例化的模板提示 LLM 生成漏洞代码样本。如果响应中不包含代码片段或遇到 API 错误,则激活重试机制,最多三次将提示重新提交给 LLM,直到收到包含代码片段的响应,否则继续处理下一个提示。

  • Verifier(验证器)
    Verifier 的作用是对生成代码的质量进行控制,因为无法保证 LLM 生成的代码一定符合要求。我们使用 Joern [42](一种基于 Antlr [40] C 解析器的模糊 C 语言解析器)过滤掉包含严重语法错误的生成代码。选择模糊解析器而非严格解析器的原因是,在各种机器学习任务的数据增强中,生成的数据不需要严格符合原始数据的标准 [10, 11, 22, 26, 56]。需要注意的是,本研究的目标是增强现有漏洞数据集,以帮助模型捕获漏洞模式并更好地泛化,而不是生成高质量的漏洞样本。完全正确的代码对于训练模型并非必要 [51],带有细微噪声的数据有助于模型的泛化能力 [4, 15, 23]。尽管如此,我们承认可以使用更复杂的方法进行验证,但为了评估使用 LLMs 进行漏洞检测的可行性,从一个简单模块(即解析器)开始是必要的。最终,模糊 C 解析器过滤掉了 2%-13% 的生成数据。

需要注意的是,我们的三种策略共享相同的 Generator 和 Verifier 组件,而不同策略在 Retriever 和 Formulator 组件的设计上略有不同。我们在第 3.1 节中讨论每种策略的 Retriever 和 Formulator 组件的细节。

3.1 增强策略

在本节中,我们将详细分解 Mutation(变异)Injection(注入)Extension(扩展) 策略。具体来说,我们为每种策略介绍其提示模板和检索器(Retriever)。


3.1.1 Mutation(变异)

大多数先前关于程序增强的研究 [5, 6, 11, 33, 39, 55] 依赖于不改变程序流程、语义和语法正确性的代码转换技术,并通过程序分析实现。例如,更改变量名、将 for 循环替换为 while 循环(反之亦然),以及添加死代码等。然而,程序分析非常耗时,且选择转换位置具有挑战性。以往研究通常随机选择转换类型和位置,而无需任何程序理解 [5, 6, 11, 33, 39, 55]。

Mutation 策略中,我们旨在利用 LLM 的程序理解能力来变异现有的漏洞代码样本,并让 LLM 自行选择转换类型和潜在的可转换语句,从而增强漏洞代码样本。通过这种方式,我们绕过了程序分析,并依赖 LLM 的创造力生成更合适且多样化的样本。
![[Pasted image 20250325145915.png]]

我们设计了如图 2 所示的提示模板。我们指示 LLM 根据最近的一项研究 [55] 使用 18 种程序转换规则之一,并在变异漏洞代码样本时至少应用其中一条规则。在设计模板时,我们优先考虑以下几点:

  1. 生成多样化代码:为此,我们在单个提示中列出所有规则,而不是为每个规则单独设计提示,以赋予 LLM 自由选择更适合的规则并将其应用于漏洞代码片段。
  2. 保持现有漏洞不变:由于某些规则可能会改变代码的执行轨迹(例如,将 switch-case 转换为 if-else 语句),从而可能改变漏洞状态,我们指示 LLM 仅对非重要行应用这些规则。所谓重要行,实际上是代码片段中的漏洞行。我们在提示中将漏洞行称为重要行,以防止结果代码中出现不必要的更改。然而,这些行仍然可以通过不会改变执行轨迹的规则(即模板中的前五条规则)进行转换,同时保持代码片段的漏洞状态不变。
    ![[Pasted image 20250325145954.png]]

如图 1 所示,Mutation 的整体工作流程非常直接。为了增强 N N N个漏洞样本,从输入数据集中随机采样 N N N个漏洞样本,并直接送入 Formulator 实例化模板,从而生成 N N N个实例化的提示。然后,这些提示被送入 Generator 生成 N N N个漏洞样本。需要注意的是,在 Verifier 组件中,平均约 3% 的生成样本会被过滤掉。因此,如果希望最终获得至少 N N N个样本,则应在生成阶段设置更高的目标值,或者在验证后重新生成更多样本以达到目标。


3.1.2 Injection(注入)

与近期的研究 [35–37] 类似,我们也专注于将漏洞代码片段注入干净样本作为主要策略,但我们的目标是涵盖所有类型的漏洞,而不仅仅是单条语句漏洞。具体来说,我们指示 LLM 将漏洞样本的逻辑注入干净样本,优先注入漏洞片段。使用 LLM 注入漏洞片段到干净样本中,赋予 LLM 自由选择最佳位置的能力。
![[Pasted image 20250325150022.png]]

我们展示了如图 3 所示的 Injection 模板提示。模板包含两到三个占位符,分别是 input_vulnerable_sampleinput_clean_sampleinput_vulnerable_lines

  • input_clean_sample 是输入的干净样本,我们希望在其中注入漏洞,这在 Mutation 策略中不存在。
  • input_vulnerable_sample 是一个与干净输入样本相似的漏洞代码示例。我们检索相似的漏洞代码,因为我们认为类似漏洞代码的逻辑可以更容易、更自然地集成到干净代码中。
  • input_vulnerable_lines 提供元信息,指示检索到的漏洞示例中的漏洞行。我们指示 LLM 优先包含这些漏洞行,因为我们的目标是生成漏洞代码。

为了填充模板,我们需要构建一个包含干净样本及其对应的相似漏洞代码样本(即 clean-vul 对)的数据集。注意,我们选择提供干净代码样本作为输入并搜索相似的漏洞样本,而不是提供漏洞代码样本并搜索相似的非漏洞样本。由于漏洞样本数量较少,这种策略更快、更高效,减少了检索过程中的数据集规模。
![[Pasted image 20250325150055.png]]

算法 1 概述了我们通过检索过程构建 clean-vul 对数据集的方法。给定一个包含漏洞样本( V V V)和干净样本( C C C)的数据集,我们的目标是检索 N N N个 clean-vul 对。一种直接的方法是为每个干净样本检索最相似的漏洞样本,按相似度降序排序,并选择前 N N N对。然而,仅关注相似性会降低检索样本的多样性,这对数据增强是不利的。先前研究表明,训练数据集中更高的多样性可以提高深度学习模型的泛化能力 [19, 47]。因此,为了增强多样性,我们引入了一个聚类阶段,将漏洞样本分组为 G G G个簇,确保所有簇的样本都被考虑用于选择。

首先,我们将漏洞样本聚类为 G G G组(第 2 行)。为此,我们使用 CodeBERT [14] 将每个代码样本嵌入为向量,特别是利用每个样本的 [CLS] token 的嵌入。然后,我们使用 KMeans 进行聚类,并通过余弦相似度衡量每对样本的相似性。完成聚类后,我们使用 Lucene 为每个簇创建索引,以促进高效的搜索(第 7-10 行)。对于每个干净样本,我们使用簇索引在每个簇中搜索最相似的漏洞样本,并存储生成的 clean-vul 对以供进一步处理(第 12-17 行)。为了为每个干净样本找到相似的漏洞样本,我们使用 BM-25 算法 [45],因为先前研究表明,在利用 RAG 时,稀疏方法(如 BM-25)与密集方法(如 CodeBERT)之间没有显著差异 [9, 18, 34]。我们使用 Lucene 的 BM-25 实现进行本研究。

一旦计算出每对 clean-vul 的相似度分数,我们根据相似度分数对每个簇内的结果进行排序(第 19-21 行)。然后,我们迭代选择每个簇中的最佳对,确保选择既多样化又相关(第 25-32 行)。我们还根据簇的大小对簇进行排序,并从最大的簇开始,以确保在更高 G G G值下覆盖范围更大(第 23 行),遵循先前的研究 [32]。在本研究中,我们设置 G = 5 G = 5 G=5。在检索到 clean-vul 对后,这些对将被送入 Formulator 实例化 Injection 模板,并通过工作流的其余组件生成 N N N个新的漏洞样本,类似于 Mutation

3.1.3 Extension(扩展)

尽管将漏洞代码片段注入干净代码是增强漏洞的默认且有效方法,但通过添加额外逻辑扩展已有的漏洞代码片段同样可以生成新的漏洞样本,从而丰富漏洞代码出现的上下文环境。与将漏洞片段注入干净代码相比,从干净样本中添加某些部分到已有的漏洞样本中有更高的概率保持原始漏洞代码的上下文完整。通过扩展已有的漏洞样本,只有小部分代码与漏洞部分无关,而这些漏洞部分对于模型来说是关键内容;相比之下,将漏洞片段注入干净样本会导致一个代码片段,其中漏洞部分与原始上下文几乎没有关联,容易暴露其人工注入的痕迹。
![[Pasted image 20250325150114.png]]

因此,我们旨在探索这种策略作为漏洞注入的替代方案,并将其称为 Extension(扩展)。我们在图 4 中展示了 Extension 的提示模板。与 Injection 策略类似,我们指示 LLM 将干净样本的部分逻辑添加到漏洞样本中,同时对漏洞行实施“不更改”策略。唯一的区别在于,Extension 策略并非将漏洞样本注入干净样本,而是反向操作。因此,我们可以通过修改 Retriever 完全复用 Injection 策略的工作流程,目标是从给定的漏洞样本中检索相似的干净样本,同时保持其他所有组件不变。

需要注意的是,我们也尝试过指示 LLM 在没有提供干净样本的情况下自由扩展漏洞样本,这为 LLM 提供了更大的自由度来扩展漏洞样本。然而,这种方法效果不佳,通常会导致幻觉现象(hallucination)。具体来说,LLM 通常会保持漏洞样本不变,或者表现得像一个变异器(例如,更改变量名或添加死代码)。当 LLM 没有收到明确指令时,它在生成代码方面的表现较差。因此,我们设计策略的方式是简单要求 LLM 遵循清晰的指令(即混合两段代码),而不提及“漏洞”一词,也不提出任何模糊或过于宽泛的要求。这避免了 LLM 过于关注代码的漏洞特性,而是鼓励它遵循给定的指令,而不是产生幻觉。

4 实验设置

在本节中,我们将介绍研究问题(RQs)、数据集、DLVD 模型、LLMs、评估指标、针对每个 RQ 的分析方法以及实现细节。


4.1 研究问题

我们从不同方面评估 VulScribeR,以回答以下研究问题:

  • RQ1:VulScribeR 相较于最先进的方法(SOTA)有多有效?
  • RQ2:RAG 对 VulScribeR 的贡献是什么?
  • RQ3:生成样本的数量如何影响漏洞检测模型的有效性?

RQ1 中,我们通过将 VulScribeR 与当前的 SOTA 方法进行比较,评估其通过增强漏洞数据来提升 DLVD 模型性能的有效性。在 RQ2 中,我们旨在调查 Retriever 组件设计的有效性。在 RQ3 中,我们探讨我们的方法和 SOTA 方法生成的样本数量如何影响 DLVD 模型的性能。


4.2 数据集

在本研究中,我们在三个广泛使用的漏洞检测数据集 [7, 16, 35, 37, 53, 58] 上评估我们的方法,分别是 Devign [58]Reveal [7]BigVul [12]。表 1 显示了这些数据集的详细信息。所有这些数据集都包含从实际项目中收集的 C/C++ 函数。
![[Pasted image 20250325150144.png]]

4.2.1 Devign

根据先前的研究 [35, 37],我们使用 Devign 作为主要训练数据集,因为它在漏洞样本和干净样本(即非漏洞样本)之间具有更好的平衡,这一点遵循了先前的研究 VulGen [37]。我们还使用该数据集中的干净样本作为 Retriever 组件的干净输入。

4.2.2 BigVul

我们使用 BigVul [12] 的全部三组数据(即训练集、验证集和测试集)用于不同目的。我们使用 BigVul 的训练集作为收集漏洞样本的来源,并将其提供给 Retriever 作为漏洞样本。然而,并非所有样本都可以使用,因为它们缺少 Injection 和 Extension 提示模板所需的漏洞行元数据。因此,我们仅使用 BigVul 训练集中包含漏洞行元数据的 6,610 个漏洞样本。我们使用 BigVul 的测试集作为本研究的一个测试集。

4.2.3 Reveal

VulGen [37] 类似,我们使用 Reveal 作为另一个测试集。Reveal 是一个比 BigVul 测试集稍大的数据集,并且漏洞样本的数量是其两倍。


4.3 基于深度学习的漏洞检测(DLVD)模型

DLVD 模型分为两类:基于标记的模型和基于图的模型。为了代表这两类模型,我们选择了一个 SOTA 基于标记的模型和两个 SOTA 基于图的模型,分别是 LineVul [16]Devign [58]Reveal [7] 模型。

对于 DevignReveal 模型,我们采用与 VulGen [37] 相同的设置,使用随机种子值训练这两个模型 5 次,测试并报告最高 F1 分数的结果。对于 LineVul,我们遵循 VGX [35] 的设置,训练 10 个 epoch,并选择在 BigVul 验证集上获得最高 F1 分数的检查点。

我们特别使用不同的数据集训练和测试模型,遵循 VulGen [37]VGX [35] 的设置,以确保没有信息泄露,并且结果可以应用于真实世界未见过的数据。


4.4 评估指标

DLVD 本质上是一个分类任务,因为代码片段会被分配到某一类别。因此,为了评估所研究方法的有效性,我们采用了常见的分类指标 [41],与先前相关研究 [28, 35, 37, 41, 43, 44, 49, 53] 一致,包括 Precision(精确率)Recall(召回率)F1-Score(F1 分数)


4.5 基础 LLMs

在本研究中,我们使用两种不同的 LLMs 来生成漏洞代码,分别是 ChatGPT3.5 [38](一种商业通用 LLM)和 CodeQwen1.5 [48](一种专注于代码任务的开源 LLM)。

  • ChatGPT:我们使用 OpenAI 的 ChatGPT [38] 的“gpt-3.5-turbo”变体,它在先前的研究中展现了巨大潜力 [9, 18, 46, 50]。我们将温度超参数设置为 0.5,因为我们希望在一定程度上减少响应的随机性,使 ChatGPT 能够遵循我们的指令,但又不会限制其创造力。其他超参数均保持默认值。
    • CodeQwen:我们使用 CodeQwen1.5 [48] 的 7B 变体“CodeQwen1.57B-Chat”,这是通用 Qwen [3] 的一个变体。与 ChatGPT 类似,CodeQwen1.5 是一个仅解码器的 Transformer 模型,不同之处在于它在代码数据上进行了预训练,并支持更大的上下文长度。我们未更改此模型的任何超参数,除了新生成 token 的数量,我们将其设置为 ChatGPT 的默认值(即 4096 个 token)。

4.6 基线方法

为了将我们的研究与先前的工作进行比较并展示 VulScribeR 的有效性,我们考虑以下基线方法:

  • NoAug:这是最基础的基线方法,我们直接使用原始数据集进行训练,不引入任何更改。
    • ROS:正如 Yang 等人 [53] 所示,随机过采样(ROS)对提升 DLVD 性能有显著效果,因此我们也将 ROS 作为另一个基线。
    • VulGen:VulGen [37] 是一种 SOTA 漏洞生成方法,它挖掘单条语句漏洞并使用 Transformer 模型定位漏洞模式的注入位置。我们将其作为基线,因为它可以生成大量数据以提升 DLVD 模型的性能。
    • VGX:VGX [35] 是 VulGen 的改进版本,它使用更大的数据集挖掘单条语句模式,并采用更复杂的定位模型。因此,我们使用 VGX,因为它可能是一种更强大的方法,能够生成比 VulGen 更多样化且高质量的数据。

4.7 针对研究问题的方法

4.7.1 RQ1 的方法

为了展示 VulScribeR 的有效性,并检查哪种提议的策略优于其他策略,我们首先通过使用不同的方法(即 VulScribeR 和基线方法)生成 5K 漏洞样本。更具体地说,对于任何策略生成 5K 样本,我们首先通过 Generator 生成多达 6K 样本,然后由 Verifier 过滤以生成超过 5K 样本,从中随机采样得到 5K 样本。

  • 对于 VGXVulGen,我们使用 VGX [35] 中生成的结果,并从其结果中随机采样 5K 样本。
  • 对于 ROS,我们从 Devign 数据集中的漏洞样本中随机过采样 5K 样本。

然后,我们使用增强后的 Devign 数据集训练三个 DLVD 模型。为了增强 Devign 数据集,我们将 5K 生成的漏洞样本添加到其中,并根据 [29] 提出的漏洞基准数据集按比例添加干净样本,以保持数据集的原始比例,这一过程遵循先前的研究 [35, 37]。我们使用 BigVul 的测试集和 Reveal 数据集评估所有这些模型的性能。我们通过评估特定 DLVD 模型在指定策略增强数据集上的表现来衡量漏洞数据增强方法的性能。因此,每种策略将在 12 种情况下进行测试(即 3 个 DLVD 模型 * 2 个测试数据集 * 2 个 LLMs)。


4.7.2 RQ2 的方法

在本研究问题中,我们调查了所提出的 RAG 对 Injection 和 Extension 策略的有效性。具体来说,我们移除了 Retriever 组件,而是随机匹配干净样本和漏洞样本,以检查 RAG 的效果(即 Injection w/o RetrieverExtension w/o Retriever)。

此外,我们还对 Retriever 组件的聚类阶段进行了消融研究,以检查多样性对生成样本质量的影响。为此,我们从 Retriever 组件中移除聚类阶段,并在所有漏洞样本中(而非每个簇内)为每个干净输入项搜索相似样本。然后,我们取前 5 个检索到的样本,以获得相似数量的 clean-vul 对。在采样阶段,我们根据相关性分数对所有样本对进行排序,并从顶部开始选取。我们将没有聚类阶段的 Injection 策略记为 Injection w/o Clustering,将没有聚类阶段的 Extension 策略记为 Extension w/o Clustering

与 RQ1 类似,我们使用所研究的方法生成 5K 漏洞样本,并通过评估在增强数据集上训练的 DLVD 模型来比较它们的有效性。


4.7.3 RQ3 的方法

最后,我们调查生成样本的数量如何影响 DLVD 模型的性能,以及通过提供更多数据可以获得多少性能提升。为了观察影响并确定 VulScribeR 是否可用于大规模漏洞增强,我们使用第 5 节中描述的最佳策略 Injection,生成多达 16.5K 样本,以便在过滤后最终得到至少 15K 漏洞样本,这比原始数据集中漏洞样本的数量高出约 40%。然后,我们使用两个 LLMs 重复 Injection 实验。我们分别采样 10K 和 15K 漏洞样本,同时按比例添加干净样本以保持数据集的比例。我们对 VGX [35]VulGen [37]ROS 也执行相同的操作。


4.8 实现细节

我们使用四块 24GB Nvidia RTX 3090 显卡加载和提示 CodeQwen1.5,以及训练和测试 LineVul、Devign 和 Reveal 模型。我们为所有研究问题总共进行了 616 次实验,仅训练和测试 DLVD 模型就花费了超过 1500 个 GPU 小时(不包括生成样本的时间)。

5 结果

5.1 RQ1:VulScribeR 相较于最先进的方法(SOTA)有多有效?

Extension 和 Injection 的表现优于 Mutation,而 Injection 略微优于 Extension。
![[Pasted image 20250325150402.png]]

表 2 展示了 VulScribeR 和基线方法在不同实验实例中的结果。可以观察到,在大多数(12 次实验中的 9 次)情况下,ExtensionInjection 的表现优于 MutationExtensionInjection 的表现非常接近,但在大多数(12 次实验中的 9 次)设置中,Injection 略占优势。在 F1 分数方面,Injection 实现了 17.60% 的 F1 分数,分别比 Extension(F1 分数为 17.38%)和 Mutation(F1 分数为 15.75%)高出 0.96% 和 12.44%。

一个可能的原因是,InjectionExtension 丰富了漏洞可能出现的上下文,而 Mutation 并未改变漏洞代码的语义。Mutation 并未丰富原始漏洞代码语义的上下文。我们提出的基于 LLM 的所有增强策略均优于基线方法。通常,InjectionExtension 在所有实例中始终优于基线方法。具体而言,Injection 在平均 F1 分数上分别比 NoAugVulgenVGXROS 高出 30.80%、27.48%、27.93% 和 15.41%。

如表 2 所示,InjectionExtension 策略在所有实验实例中均优于基线方法(即 NoAugVulGenVGXROS)。更具体地说,Extension 在平均 F1 分数上分别比基线方法 NoAugVulgenVGXROS 高出 29.68%、26.27%、26.90% 和 14.35%,而 Injection 策略则分别高出 30.80%、27.48%、27.93% 和 15.41%。

可以观察到,Mutation 在大多数情况下显著优于基线方法。然而,MutationROS 并没有绝对优势,并且在某些情况(12 次实验中的 4 次)下,ROS 的表现优于 Mutation。值得注意的是,ROS 在平均 F1 分数上分别比基线方法 VulgenVGX 高出 13.96%、10.72% 和 11.21%。

正如第 2 节所述,像 VGXVulGen 这样的 SOTA 方法仅专注于单条语句漏洞,这限制了生成漏洞的多样性。然而,我们的方法没有这一限制,能够生成更多样化的漏洞样本。为了验证这一点,我们测量了通过我们的方法以及 VGXVulGen 增强后的漏洞样本的多样性。按照先前的研究方法,我们首先对生成的漏洞样本的 CodeBERT 嵌入应用主成分分析(PCA),将维度降至三维。接下来,我们使用 10 个区间计算直方图,并从中计算向量的熵值。较高的熵值表明漏洞样本的多样性更高,而较低的熵值则表明分布更加集中,样本相似度更高。
![[Pasted image 20250325150425.png]]

如图 5 所示,对于 ChatGPTMutationInjectionExtension 的熵值分别为 4.79、4.5 和 4.62;对于 CodeQwen,熵值分别为 4.66、4.45 和 4.42。而对于 VulgenVGX,熵值分别为 4.42 和 4.31。结果表明,我们的方法在生成多样化样本方面具有更好的潜力,因为我们仅使用了其中一个用于漏洞挖掘的数据集就超越了它们。VulScribeR 生成的漏洞样本比 SOTA 方法 VGXVulGen 更加多样化。

最后,值得注意的是,两种 LLM 的结果非常接近。CodeQwen1.5-7B-Chat 在所有三种策略中平均略微优于 ChatGPT3.5 Turbo(差距为 1.49%)。然而,我们不能因此得出结论认为 CodeQwen 是更适合漏洞生成的 LLM,因为我们并未探索每种 LLM 的最佳超参数设置。不过,我们可以得出结论:我们的策略适用于与 ChatGPTCodeQwen 类似的 LLM(例如 GPT-4 [2]DeepSeek-Coder [21]),并且不需要专门针对代码训练的 LLM。

总体而言,Injection 和 Extension 的表现远远优于基线方法和 Mutation,而 Injection 的性能略高于 Extension。例如,Injection 在平均 F1 分数上分别比 NoAug、Vulgen、VGX 和 ROS 高出 30.80%、27.48%、27.93% 和 15.41%。


5.2 RQ2:RAG 对 VulScribeR 的贡献是什么?

![[Pasted image 20250325150513.png]]

对于 Injection 策略,Retriever 的聚类阶段在所有设置中平均为 F1 分数增加了 0.49%,而无聚类版本相比无 Retriever 版本则将平均 F1 分数提高了 6.34%。如表 3 所示,相似性和多样性均对模型的有效性产生积极影响,其中相似性的影响更为显著。

当比较 InjectionInjection w/o Retriever 时,我们观察到完整的 Retriever 组件将 Injection 的有效性提升了 4.99% 的平均 F1 分数,并在 12 次实验中的 8 次中表现更优。具体来说,当我们将 Injection 与无聚类版本进行比较时,发现聚类阶段将 RAG 的有效性提升了 0.49% 的平均 F1 分数,并在 12 次实验中的 8 次中表现更优;而通过比较无聚类版本与无 Retriever 版本,我们发现 RAG 将其有效性提升了 6.34% 的平均 F1 分数,并在 12 次实验中的 8 次中表现更优。
![[Pasted image 20250325150545.png]]

对于 Extension 策略,Retriever 的聚类阶段在所有设置中平均为 F1 分数带来了 2.54% 的提升,而无聚类版本相比无 Retriever 版本则将平均 F1 分数提高了 5.61%。如表 4 所示,相似性和多样性均对模型的有效性产生积极影响,其中相似性的影响更为显著。总体而言,完整的 Retriever 模块将 Extension 策略的有效性提升了 10.77% 的平均 F1 分数,并在 12 次实验中的 6 次中超越了其他设置。具体来说,通过比较 Extension 与无聚类版本,我们发现聚类阶段将 RAG 的有效性提升了 2.54% 的平均 F1 分数,并在 12 次实验中的 7 次中表现更优;而通过比较无聚类版本与无 Retriever 版本,我们发现 RAG 将其有效性提升了 5.61% 的平均 F1 分数,并在 12 次实验中的 10 次中表现更优。

Retriever 组件对 Injection 和 Extension 做出了显著贡献。 Extension 从 Retriever 组件中获得的提升超过两倍,这表明它对检索对的相似性和多样性的敏感度更高。

5.3 RQ3:生成样本的数量如何影响漏洞检测模型的有效性?

![[Pasted image 20250325150615.png]]

![[Pasted image 20250325150648.png]]

通过 Injection 增加更多漏洞数据有助于提升 DLVD 模型的有效性。表 5 和图 6 展示了在使用不同漏洞样本(由 Injection 和其他基线方法生成,例如 VulGenVGXROSNoAug)增强的数据集上训练的 DLVD 模型的性能(以 F1 分数衡量)。我们可以看到,在几乎所有实验实例中(12 次中的 11 次),随着更多增强的漏洞样本被添加到训练数据中,DLVD 模型的性能有所提升。这表明,通过 Injection 增强更多数据为模型捕捉漏洞提供了更有用的信息。

总体而言,VulGenVGXROS 在增加更多漏洞样本时未能提升 DLVD 模型的性能。从图 6 中的基线结果可以看出,VulGenVGXROS 在通过增加更多数据来提升 DLVD 模型性能方面表现不佳。例如,当我们比较在增强 5K 和 15K 漏洞样本时 DLVD 模型的性能时,发现 VGXVulGenROS 分别在 6 次实验中的 5 次、4 次和 4 次中降低了性能(注意,这些基线方法不依赖于 LLM,但为了便于比较在此展示)。这一发现与 Yang 等人 [53] 报告的结果一致,表明过度随机过采样并不能提高 DLVD 模型的有效性,甚至可能降低其性能。一个可能的解释是,随机过采样并未引入任何新信息,而过度过采样可能导致模型对过采样数据的过拟合。对于 VGXVulGen,这可能是由于它们的方法仅专注于单条语句漏洞,这极大地限制了此类方法的实用性,并使它们无法用于大规模数据增强。

总之,Injection 在所有规模下均优于所有基线方法(即 VulGen、VGX 和 ROS),尤其是在用于大规模数据增强时。所研究的基线方法在增强超过 5K 样本时未能提升 DLVD 模型的有效性。我们的基于 LLM 的方法更适合大规模漏洞数据增强。具体来说,Injection 在 15K 样本时分别比 NoAug、VulGen(15K)、VGX(15K) 和 ROS(15K) 的性能高出 53.84%、54.10%、69.90% 和 40.93%。


6 讨论

6.1 成本分析

根据我们的实验结果,每次生成 1,000 个漏洞样本(无并发处理)需要花费 3.4 小时,使用 GPT-3.5-Turbo 的成本约为 1.88 美元,而在两块 RTX 3090 上使用 CodeQwen1.5-7B-Chat 则需要约 9 个 GPU 小时。这表明 VulScribeR 的可行性,因为使用 ChatGPT 将像 Devign 这样的数据集扩展至两倍大小的成本不到 19 美元,同时可将 DLVD 模型的性能提升高达 128.95%。

6.2 对有效性的威胁

内部有效性

我们研究的一个常见威胁是对 DLVD 模型和 LLM 的超参数设置。对于 DLVD 模型,鉴于我们研究的规模,超参数调优极其昂贵,因此我们遵循了先前研究 [35, 37] 的设置。此外,在比较不同的数据增强方法时,我们为 DLVD 模型使用了相同的设置,这确保了我们的比较是公平的。对于 LLM,我们主要使用默认设置,仅将新生成 token 的数量设置为 ChatGPT 的最大值,并将 ChatGPT 的温度降低至 0.5。通过微调这些超参数可能会获得更好的结果。

另一个威胁是,LLM 可能并不总是按照提示指令生成结果,并可能出现幻觉现象。为了缓解这一威胁,我们添加了一个验证组件以过滤掉低质量代码,并使用重试机制,在出现错误(包括空代码响应)时最多提示 LLM 三次。更重要的是,我们本研究的目标是增强现有的漏洞数据集,以帮助模型捕捉漏洞模式并更好地泛化,而不是生成高质量的漏洞样本。完全正确的代码对于训练模型并非必要 [51],而带有细微噪声的数据有助于模型的泛化 [4, 15, 23]。我们承认可以使用更复杂的方法进行验证,但从简单模块(即解析器)开始是评估使用 LLM 进行漏洞检测是否可行的必要步骤。

外部有效性

外部有效性涉及我们研究结果的普适性。尽管我们使用了两种不同的 LLM(ChatGPT 3.5 TurboCodeQwen1.5-7B-Chat),并在三个常用数据集上进行了评估,涵盖了两类 SOTA DLVD 模型,但我们的研究结果可能无法很好地推广到使用不同模型、LLM 和数据集的真实场景中。我们鼓励未来的研究探索更多 LLM(包括基于编码器-解码器的模型)、数据集和 DLVD 模型。


7 结论

在本文中,我们提出了 VulScribeR,这是一种新颖且有效的漏洞样本增强框架,利用定制化的 RAG 机制设计精心设计的提示模板,然后用于提示 LLM。我们还使用模糊解析器作为验证方法,以确保生成的代码没有严重的语法问题。我们的评估结果表明,我们的方法显著优于 SOTA 方法,并且适合低成本的大规模漏洞增强。

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值