【论文阅读】Generalization-Enhanced Code Vulnerability Detection via Multi-Task Instruction Fine-Tuning

最新推荐文章于 2025-04-29 16:09:39 发布

君的名字

最新推荐文章于 2025-04-29 16:09:39 发布

阅读量723

点赞数 33

文章标签：论文阅读

本文链接：https://blog.csdn.net/chichoxian/article/details/147502457

版权

今天介绍的文章来自华科的团队：基于多任务指令微调的泛化增强代码漏洞检测，这个文章发表在ACL2024

在这里插入图片描述
这篇文章可以和我上一篇介绍的CMU的漏洞检测（Security Vulnerability Detection with Multitask Self-Instructed
Fine-Tuning of Large Language Models）这个文件联合起来一起阅读，两个文件都是介绍利用LLM做漏洞检测，可以发现其实标题都很像，就是技术稍微有点不同，上面一篇MSVID是大模型加GNN结合起来做的，这个只用了LLM。

基于代码预训练模型的漏洞检测方法近年来取得了显著进展。然而，现有模型通常仅学习从源代码到标签的浅层映射，而未能深入理解代码漏洞的根本成因，导致其难以泛化到训练集之外的现实场景。为此，我们提出VulLLM框架，通过将多任务学习与大型语言模型相结合，实现对深层漏洞特征的有效挖掘。具体而言，我们在漏洞检测主任务基础上构建两个辅助任务：首先利用漏洞补丁构建漏洞定位任务，其次基于补丁提取的漏洞特征，借助GPT-4构建漏洞解释任务。VulLLM创新性地通过生成式大语言模型理解复杂漏洞模式，迫使模型捕捉漏洞的根源性特征而非过拟合单一任务的伪特征，从而实现漏洞分类的增强。在六个大规模数据集上的实验表明，VulLLM在检测效能、泛化能力和鲁棒性三个维度均优于七种前沿模型。

近年来，基于代码预训练模型（CodePTMs）的自动化漏洞检测方法（如CodeBERT（Feng等人，2020）、GraphCodeBERT（Guo等人，2021）和UniXcoder（Guo等人，2022））取得了最先进（SOTA）的成果。这些模型以代码片段为输入，预测代码中是否存在潜在漏洞。然而，近期研究（Du等人，2023a）揭示了此类模型泛化能力的重大缺陷，尤其是在处理分布外（OOD）数据时。这一局限源于现有方法在学习从源代码到标签的映射时，往往只能捕获浅层而非深层的漏洞特征。典型表现为：此类方法无法准确区分仅通过标识符替换生成的对抗样本（Zhang等人，2023b, 2020），表明其预测易受与漏洞无关的干扰因素影响。此外，由于不同项目的代码在编程风格和应用场景上存在差异，导致漏洞特征的分布存在显著分歧（Du等人，2023a），基于“代码到标签”映射的学习范式在处理跨项目漏洞时泛化能力受限。

值得注意的是，近期涌现的大语言模型（LLMs）在多个领域展现出卓越的推理与泛化能力（Ge等人，2023；Cao等人，2023），这启发我们将其应用于构建更鲁棒的漏洞检测模型。然而，直接应用LLMs面临特定挑战：缺乏针对漏洞检测任务的专项训练（Zhang等人，2023a；Gao等人，2023），且若沿用类似CodePTMs的微调方法仍会引入前述泛化问题。为此，本研究提出VulLLM框架，通过指令调优将代码漏洞知识注入LLMs，实现检测能力的系统性提升。

为防止模型习得伪特征，我们采用多任务学习范式，促使大语言模型（LLMs）学习深层特征而非表面关联。其核心思路在于：通过增加两个旨在揭示漏洞根本成因的辅助任务（漏洞定位与漏洞解释），增强从源代码到标签的映射深度。具体而言：

漏洞定位任务：基于漏洞补丁识别具体漏洞代码元素（如语句级定位）；
漏洞解释任务：解析漏洞根源并生成文本化解释。由于缺乏现成的解释数据，我们通过GPT-4生成此类标注。然而，LLMs在漏洞检测任务中已面临挑战（Zhang等人，2023a），识别漏洞根源更为困难——手动评估表明，ChatGPT基于简单提示所“检测”的漏洞中，仅能理解半数成因。

为此，我们提出补丁增强的思维链自验证机制（CoT-SV），该方法通过三重验证：
• 漏洞标签一致性验证

• 通用漏洞披露（CVE）描述匹配

• 基于程序依赖图（PDG）（Li等人，2022）的上下文漏洞行校验

有效规避传统思维链（Wei等人，2022）中的错误传播与幻觉问题，显著提升LLMs的推理可靠性（Ni等人，2023；Gou等人，2023）。

辅助任务通过扩展特征维度（漏洞位置与根源）提升模型对代码漏洞领域知识的理解深度，同时多任务特征多样性迫使模型寻求跨任务最优解，从而抑制对单一任务伪特征的过拟合。

除前述用于多任务学习的生成数据外，本研究的训练数据还包含来自两项人工标注验证研究（Chen等人，2023；Croft等人，2023）的两个高标签准确率真实漏洞数据集——Devign（Zhou等人，2019）与DiverseVul（Chen等人，2023），以进一步强化大语言模型（LLMs）对多样化代码漏洞的学习能力。为验证框架的广泛适用性，我们选取三类主流基础模型构建VulLLM：通用大模型Llama-2（Touvron等人，2023）及两款代码大模型CodeLlama（Rozière等人，2023）、StarCoder（Li等人，2023）。实验表明，VulLLM在检测效能上全面超越七种现有SOTA漏洞检测模型。具体而言，相较最佳基线模型UniXcoder，VulLLM在六个数据集上的F1值平均提升8%。值得关注的是，在四个分布外（OOD）数据集上，其F1值提升达8.58%，印证了更优的泛化能力。此外，通过引入三种对抗攻击验证模型鲁棒性，VulLLM的总体F1值较UniXcoder提升68.08%，展现出显著的抗干扰能力。

本研究的主要贡献可归纳为以下三方面：

创新视角：首创性地利用GPT-4的可解释性生成漏洞成因描述，系统增强LLMs对漏洞语义的理解深度。
方法框架：提出基于多任务指令微调的VulLLM框架，据我们所知，这是首次将指令微调LLMs应用于漏洞检测领域的研究实践。
实验验证：在六个数据集上的大规模实验表明，该方法在检测效能、泛化性与鲁棒性三个维度均实现显著提升。相关代码与数据已开源：
https://github.com/CGCL-codes/VulLLM

代码漏洞检测

代码漏洞检测对保障软件系统安全稳定运行具有重要意义。基于深度学习的方法能够从海量代码样本中自动学习并泛化漏洞特征，实现漏洞模式的自动化推理，该范式近年来备受关注。早期检测方法主要利用图神经网络（GNNs）学习漏洞特征；随着Transformer架构（Vaswani等人，2017）的发展，CodeBERT（Feng等人，2020）、GraphCodeBERT（Guo等人，2021）及UniXcoder（Guo等人，2022）等代码预训练模型（CodePTMs）在漏洞检测任务中展现出更优性能。这些模型通过代码-文本对的大规模预训练，在多类代码相关下游任务中表现卓越。此外，部分方法基于CodePTMs进行扩展：ReGVD（Nguyen等人，2022）将源代码编码为图结构（节点表示代码符号，特征由CodePTMs初始化），EPVD（Zhang等人，2023c）基于控制流图（CFG）划分代码执行路径并学习不同路径的CodePTMs表征。现有方法多采用单任务学习范式，而本研究通过多任务学习增强大语言模型（LLMs）对代码漏洞的理解能力。

大语言模型自验证机制

思维链（CoT）（Wei等人，2022）是一种基于LLMs的推理提示技术，通过模拟人类解决问题的分步推理过程处理复杂任务。自验证机制（Pan等人，2023）旨在缓解LLMs的幻觉问题（Lin等人，2022）与不可靠推理（Golovneva等人，2023；Lyu等人，2023），同时减少CoT中的错误累积（Weng等人，2023）。其核心是通过反馈修正LLMs的不良行为，类似于人类“尝试-纠错-修正”的学习循环。验证来源可分为两类：
• 人工验证：如InstructGPT（Ouyang等人，2022）通过人类反馈优化GPT-3（Brown等人，2020），更符合人类偏好；

• 自动验证：可基于LLM自身或外部知识。例如SelfCheck（Miao等人，2023）利用LLM独立修正CoT中的错误（无需外部资源），通过多阶段结果的置信度加权投票提升响应准确性；而基于维基百科（Varshney等人，2023）或搜索引擎（Gou等人，2023）等外部知识的验证则依赖领域信息。本研究从漏洞补丁提取特征实现验证（详见第3.2节）。

方法论

3.1 总体架构
图1展示了VulLLM的三阶段框架：

漏洞特征提取（3.2节）：从漏洞补丁中挖掘代码语义线索
漏洞解释生成（3.3节）：基于GPT-4生成漏洞成因描述
多任务指令微调（3.4节）：通过联合训练增强模型泛化能力

3.2 漏洞特征提取
漏洞特征是支撑漏洞解释的核心语义线索。本研究聚焦三类特征以增强LLMs的泛化能力：
• 漏洞行：直接标识漏洞代码元素

• 漏洞上下文：包含漏洞触发条件的周边代码

• CVE描述：官方漏洞披露文档的技术分析

漏洞行提取：

补丁通过增删特定代码元素修复漏洞，其中被删除的代码行直接反映漏洞语义（Nguyen等人，2016）。例如，若被删代码涉及不安全编码实践（如内存管理缺陷或输入验证漏洞），则这些代码行即为漏洞根源。具体提取流程如下：

解析漏洞补丁的diff文件
定位标记为"-"的删除代码行
通过代码抽象（如变量重命名）提取语法无关的语义特征

漏洞上下文提取

漏洞上下文指围绕漏洞代码的关联语句（如条件判断、安全检查等），为理解漏洞提供语义支撑。本研究通过程序依赖图（Program Dependency Graph, PDG）（Li等人，2022）提取与漏洞行存在数据/控制依赖的代码，具体流程如下：

技术流程：

PDG构建：使用JOERN工具（JOERN，2023）生成漏洞函数的程序依赖图。PDG为有向无环图，节点表示代码元素，边表示数据依赖（如变量传递）或控制依赖（如条件分支）。
依赖路径挖掘：以漏洞行为起点，沿数据依赖边与控制依赖边双向遍历k跳（k-hop）范围内的邻接节点。
上下文合成：将遍历所得节点对应的代码行纳入漏洞上下文。

实施细节：

• 跳数控制：设置k=1以平衡上下文广度与LLMs输入长度限制。如图2所示CVE-2018-7751漏洞（CWE-835类型），当k=1时，第9行漏洞行的控制依赖边（←第8行）与数据依赖边（←第3行、→第14/16行）共同构成上下文范围（第3、8、14、16行）。

• 复杂依赖处理：真实代码依赖关系可能高度复杂（如循环嵌套），k值控制可避免上下文爆炸式增长。

CVE描述整合

CVE描述通过官方漏洞披露平台NVD（NVD，2023）爬取，其从软件/硬件层面系统阐述漏洞的根本成因、触发条件及潜在攻击场景。例如，针对缓冲区溢出类漏洞，CVE描述会详细说明未进行输入长度校验的具体代码位置与内存操作异常机制。这些描述为LLMs理解漏洞语义提供关键背景知识。

3.3 漏洞解释生成

基于前文提取的漏洞特征，本节通过四阶段思维链自验证（CoT-SV）生成可靠漏洞解释，具体流程如图3所示。

Step 1：角色化提示初始化
受提示工程中角色扮演策略的启发（Kong等人，2023；Shanahan等人，2023），我们设计角色约束型提示以聚焦漏洞检测任务。基础提示模板如下（改编自Zhang等人，2023a）：

Prompt1：请扮演漏洞检测模型，判断以下代码是否存在漏洞：[代码片段]

其中[代码片段]为待检测代码。研究表明，单一提示下LLMs的检测精度有限（Cheshkov等人，2023）。为此，我们采用正负样本均衡策略：
• 漏洞代码：若LLM输出错误，后续步骤将逐步修正

• 非漏洞代码：要求GPT-4生成"无漏洞"解释，用于构建负样本数据集

该阶段通过强制LLM确认漏洞存在性，为后续推理奠定基础。

Step 2-4：多特征渐进式验证
基于自验证模板（Pan等人，2023；Ling等人，2023），我们构建统一的三段式验证框架。每个步骤包含：

待验证信息：当前步骤需确认的特征（如CVE描述）
验证指令：如"请重新检查答案并分析其正确性"（Ling等人，2023）
输出约束：规定下一阶段LLM需生成的内容

具体提示词设计如下：

Prompt2：该代码存在漏洞。请检查答案并分析正确性，随后描述漏洞成因。  
Prompt3：漏洞描述为[CVE描述]。请检查正确性，随后定位直接相关的漏洞代码行。  
Prompt4：漏洞行位于[漏洞行]。请检查正确性，随后列出与之存在数据/控制依赖的代码行。

其中[CVE描述]和[漏洞行]来自3.2节提取的漏洞特征。通过三阶段交叉验证，逐步修正LLM的幻觉与推理偏差。

Step 5：漏洞上下文整合与解释生成
漏洞上下文作为CoT-SV的最终验证特征，驱动大语言模型（LLM）融合多源特征生成完整漏洞解释。具体流程如下：

Prompt5设计：

Prompt5：依赖代码行为[依赖行]。请重新检查答案并分析正确性，随后综合漏洞描述与依赖关系，生成完整的漏洞解释。

其中[依赖行]为步骤4提取的数据/控制依赖代码行。

知识整合机制：
• CVE描述：约束漏洞的语义边界（如漏洞类型、攻击向量）

• 漏洞行：定位具体缺陷代码位置

• 依赖行：揭示漏洞触发路径的上下文逻辑

抗幻觉设计：

语义抽象：将代码依赖流（数据流/控制流）转化为自然语言描述
人工验证：所有生成数据均经过人工验证，剔除GPT-4最终判断仍错误的样本

以图2漏洞代码为例，附录A对比展示了初始提示（Step1）与完整CoT-SV流程生成的解释差异，验证多步验证机制的有效性。

在这里插入图片描述

3.4 多任务指令微调

数据准备
多任务数据主体源自PatchDB（Wang等人，2021），其独特的补丁差异（diff）信息为提取漏洞解释提供关键支持。鉴于该数据集规模有限，我们进一步整合两个高精度人工标注数据集：DiverseVul（Chen等人，2023）与Devign（Zhou等人，2019）（在人工评估抽样中准确率最高，参见Croft等人，2023）。三任务数据构建策略如下：
• 主任务（漏洞检测）：融合797个开源项目代码，通过随机标识符替换（10%变量/函数名替换为跨项目随机选择标识符）缓解跨项目特征分布偏移问题

• 辅助任务1（漏洞定位）：输入源代码，输出图2所示漏洞行定位结果

• 辅助任务2（漏洞解释）：输入源代码，输出CoT-SV生成的文本化解释（示例见附录A）

指令微调机制
基于Alpaca框架（Taori等人，2023）构建指令模板，引导LLMs生成任务特定输出：

### Instruction:  
[任务指令]  
### Input:  
[源代码]  
### Response:  
[目标输出]

其中：
• [任务指令]：差异化指令（如“检测漏洞”、“定位漏洞行”、“生成解释”）

• [目标输出]：对应任务标签（0/1）、漏洞行编号或自然语言解释

通过计算模型输出与目标输出的差异，反向传播更新LLMs权重，实现领域知识注入。

4 实验设置

4.1 数据集

选取六大规模C/C++漏洞检测基准数据集：
• 训练集（Dataset1）：DiverseVul、Devign

• 测试集（Dataset2）：BigVul（Fan等人，2020）、CVEfixes（Bhandari等人，2021）、ReVeal（Chakraborty等人，2022）、Juliet（Boland & Black，2012）

Dataset2用于验证模型在分布外场景（OOD）的泛化能力。数据集统计信息详见附录B。

译文：
4.2 基线模型
为全面评估VulLLM性能，我们选取七种SOTA模型作为基线，涵盖以下技术范式：

基于图神经网络的模型
• Devign（Zhou等人，2019）

• ReVeal（Chakraborty等人，2022）
代码预训练模型（CodePTMs）
• CodeBERT（Feng等人，2020）

• GraphCodeBERT（Guo等人，2021）

• UniXcoder（Guo等人，2022）
CodePTMs扩展模型
• ReGVD（Nguyen等人，2022）

• EPVD（Zhang等人，2023c）

模型细节见附录C，实现参数配置见附录E。

5 实验结果

5.1 检测效能与泛化能力

我们分别选择Llama-2、CodeLlama和StarCoder的两种不同参数量版本作为基模型，以验证本框架的广泛适用性。具体而言，我们在六个测试数据集上评估了各种方法的有效性。值得注意的是，未参与模型训练的Dataset2可进一步反映模型的泛化能力。

表1所示结果表明，基于CodeLlama-13B的VulLLM展现出最高性能，其整体F1值达到66.54%。在所有模型和数据集上，基于CodeLlama-13B的VulLLM始终排名第一或第二，显著优于七个选定基线模型。因此，我们后续对VulLLM的分析均基于CodeLlama-13B版本。相较于最佳基线模型UniXcoder，VulLLM在全部六个数据集上的整体效能提升8%（即(66.54-61.61)/61.61）。特别地，在Dataset2上的F1值提升达8.58%（即(64.16-59.09)/59.09），这印证了其卓越的泛化能力。

此外，现有模型普遍表现出较差的泛化能力。在CodePTMs的20组（5模型×4数据集）泛化实验中，基线模型的平均F1值较Dataset1下降11.36%至38.36%（平均降幅24.33%）。相比之下，VulLLM的性能降幅显著更小（仅下降10.03%），且关键的是在所有数据集上均保持F1值高于60%。尽管基于GNN的模型在Dataset2上似乎比CodePTMs表现出更小的性能降幅（表面上显示更好的泛化能力），但其较低的检测效能和复杂的数据预处理流程使其实际适用性远低于其他模型。

我们选取Llama-2、CodeLlama与StarCoder的不同参数量版本（7B/13B）作为基模型，验证框架的广泛适用性。表1结果显示：
• 最优性能：基于CodeLlama-13B的VulLLM以66.54%的F1值位居榜首

• 跨数据集稳定性：在全部6个数据集上均保持前两名，显著优于7种基线模型

• 泛化优势：在4个分布外数据集（Dataset2）上，F1值较最佳基线UniXcoder提升8.58%（64.16% vs 59.09%）

关键发现：

CodePTMs的泛化局限：基线模型在分布外数据上的F1值平均下降24.33%（从38.36%降至11.36%）
VulLLM鲁棒性突破：性能降幅仅10.03%，且所有数据集F1值均超60%
GNN模型的实用性缺陷：尽管在Dataset2上性能降幅较小，但其检测效能低下且数据预处理复杂，严重制约实际应用

在这里插入图片描述

鲁棒性评估

在这里插入图片描述

本实验分别从Dataset1和Dataset2选取DiverseVul与ReVeal数据集。选择依据如下：
• DiverseVul：涵盖多类型项目代码，可全面评估模型抗干扰能力

• ReVeal：数据源为Chromium与Debian软件包（区别于NVD平台或GitHub仓库的其他数据集）

对抗攻击方法选取基于随机标识符替换的两种最新技术：
• MHM（张等人，2020）

• WIR-Random（曾等人，2022）

其在近期评估中展现出最高攻击成功率（杜等人，2023b）。此外，我们基于随机死代码插入构建第三种攻击，其形式参考DIP方法（Na等人，2023）。由于DIP代码未开源，且本研究目标为对比模型抗攻击能力（非追求最大攻击成功率），故未直接复现DIP。攻击细节详见附录F。

选取表1中表现最佳的基线模型UniXcoder与VulLLM进行鲁棒性测试。表2显示两模型在三种对抗攻击下的F1值：
• UniXcoder：各类攻击下性能显著下降，且在ReVeal数据集（分布外数据）上的降幅更大，印证其泛化能力不足

• VulLLM：所有攻击场景下均保持优越鲁棒性，F1值较UniXcoder平均提升68.08%。更重要的是，其在分布外样本上的鲁棒性未出现进一步衰减，表明其鲁棒性与泛化能力全面领先

鲁棒性差异归因分析
为探究两模型鲁棒性差异的内在机理，我们对VulLLM与UniXcoder在两个数据集上的正确预测概率密度分布进行深度解析。在高风险安全场景中，模型对预测结果的高置信度（接近1.0）尤为重要。如图4所示，整体分析表明：

关键发现：

DiverseVul数据集
• VulLLM：高概率区（>0.9）密度显著高于UniXcoder，曲线在0.9附近达到峰值后快速衰减，表明其预测高度集中于高置信区间

• UniXcoder：密度分布相对分散，峰值位置低于0.8
ReVeal数据集
• VulLLM：在>0.8区域仍保持密度优势，0.8-1.0区间密度曲线始终高于UniXcoder，峰值稳定在0.9附近

• UniXcoder：密度曲线在0.6-0.8区间出现平台期，高置信预测占比显著降低

结论：VulLLM的预测置信度分布更集中于高概率区间（尤其是>0.8区域），这种对自身预测的强确信性是其卓越鲁棒性的重要成因。该特性使其在对抗攻击下仍能维持稳定的高精度输出。

5.3 消融实验

本节探究多任务学习（MT）与数据增强（DA）对模型性能的贡献。如表3所示：

组件移除影响：

移除多任务学习（w/o MT）
• 所有数据集性能下降，整体相对降幅达9.30%

• 证实MT对提升模型效能与泛化能力具有核心作用
移除数据增强（w/o DA）
• 平均性能下降2.98%

• 四类数据集性能降低，但ReVeal与Juliet数据集性能上升

• 表明DA虽有效增强模型效能，但可能对特定数据集泛化产生负面影响
同时移除MT&DA（w/o DA&MT）
• 五个数据集性能全面下降

• 在Devign数据集上性能低于三种CodePTMs

• 在BigVul与CVEfixes数据集上表现不及ReGVD

核心结论：
• 多任务学习与数据增强对模型整体性能提升缺一不可

• 多任务学习贡献度更高（相对影响权重9.30% vs 2.98%），是驱动性能优化的关键组件
在这里插入图片描述

5.4 超参数敏感性分析

表4：不同辅助任务样本量下VulLLM的F1值

上下文长度	辅助数据量	DiverseVul	Reveal	总平均
512	694	70.99	64.52	66.54
1,024	1,509	69.66	64.36	66.76
2,048	2,138	68.68	67.00	66.89

实验设置：

为探究辅助任务样本量对模型性能的影响，我们在保持训练效率的前提下，将上下文长度从基线值512扩展至1,024与2,048，以纳入更多辅助任务数据。实验选取5.2节的两个典型数据集（DiverseVul/Reveal）进行重点分析，并统计六数据集平均表现。

关键发现：

辅助数据增益效应
• 辅助样本量从694增至2,138时，模型在分布外样本（Reveal）上的F1值提升2.48%（64.52%→67.00%）

• 六数据集总平均F1值微增0.35%（66.54%→66.89%）
分布内性能权衡
• 在原始训练分布数据集（DiverseVul）上，F1值下降2.31%（70.99%→68.68%）

归因分析：
• 正向驱动：多任务学习迫使模型提取跨任务共性特征（如漏洞语义模式），增强对未知分布的适应能力

• 负向权衡：辅助任务数据占比增加可能稀释主任务（漏洞检测）的细粒度特征学习，导致分布内性能轻微衰减

工程启示：
上下文长度的选择需权衡计算资源与任务需求——资源受限场景建议维持512长度；追求跨域泛化时可扩展至2,048以激活多任务协同效应。