通过多任务指令微调实现漏洞检测的泛化性提升
密码学与安全
2024-06-07 v1 人工智能
计算与语言
摘要
基于代码预训练模型(CodePTM)的漏洞检测方法近年来取得了显著成果,但这些模型通常仅学习从源代码到标签之间的浅层映射,难以理解漏洞的根本原因,导致在训练实例之外的实际场景中表现不佳。为此,我们引入 VulLLM 框架,将多任务学习与大语言模型(LLM)集成,以有效挖掘深层次漏洞特征。具体而言,我们在漏洞检测任务之外构建了两个辅助任务。首先,我们利用漏洞补丁构建漏洞局部化任务。其次,基于补丁提取的漏洞特征,我们利用 GPT-4 构建漏洞解释任务。VulLLM 通过生成式 LLM 理解复杂漏洞模式,创新性地增强了漏洞分类,使模型能够捕获漏洞的根本原因而非对单一任务的虚假特征进行过拟合。在六个大型数据集上的实验表明,VulLLM 在有效性、泛化性和鲁棒性方面均优于七个最先进模型。
引用
@article{arxiv.2406.03718,
title = {Generalization-Enhanced Code Vulnerability Detection via Multi-Task Instruction Fine-Tuning},
author = {Xiaohu Du and Ming Wen and Jiahao Zhu and Zifan Xie and Bin Ji and Huijun Liu and Xuanhua Shi and Hai Jin},
journal= {arXiv preprint arXiv:2406.03718},
year = {2024}
}
备注
Accepted to ACL 2024 Findings