arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2406.03718cs.CRcs.AIcs.CL

通过多任务指令微调增强泛化能力的代码漏洞检测

Generalization-Enhanced Code Vulnerability Detection via Multi-Task Instruction Fine-Tuning

School of Cyber Science and Engineering, Huazhong University of Science and Technology(华中科技大学网络空间科学与技术学院);College of Computer, National University of Defense Technology(国防科技大学计算机学院);School of Computer Science and Technology, Huazhong University of Science and Technology(华中科技大学计算机科学与技术学院)

Xiaohu Du, Ming Wen, Jiahao Zhu, Zifan Xie, Bin Ji, Huijun Liu, Xuanhua Shi, Hai Jin

AI总结:

本文提出VulLLM框架,通过多任务指令微调结合漏洞定位与漏洞解释两个辅助任务,使大语言模型捕捉漏洞根本原因,从而在六个数据集上超越七个最先进模型,提升漏洞检测的有效性、泛化性与鲁棒性。

摘要:

基于代码预训练模型(CodePTMs)的漏洞检测近年来取得了可喜的成果。然而,这些模型难以泛化,因为它们通常学习的是从源代码到标签的浅层映射,而不是理解代码漏洞的根本原因,导致其在训练实例之外的真实场景中表现不佳。为应对这一挑战,我们提出了VulLLM,这是一个将多任务学习与大语言模型(LLMs)相结合以有效挖掘深层漏洞特征的新型框架。具体而言,我们在漏洞检测任务之外构建了两个辅助任务。首先,我们利用漏洞补丁构建漏洞定位任务。其次,基于从补丁中提取的漏洞特征,我们利用GPT-4构建漏洞解释任务。VulLLM创新性地通过利用生成式大语言模型来理解复杂漏洞模式,从而增强漏洞分类,促使模型捕捉漏洞的根本原因,而不是过拟合单一任务的虚假特征。在六个大型数据集上进行的实验表明,VulLLM在有效性、泛化性和鲁棒性方面均超越了七个最先进的模型。

英文摘要:

Code Pre-trained Models (CodePTMs) based vulnerability detection have achieved promising results over recent years. However, these models struggle to generalize as they typically learn superficial mapping from source code to labels instead of understanding the root causes of code vulnerabilities, resulting in poor performance in real-world scenarios beyond the training instances. To tackle this challenge, we introduce VulLLM, a novel framework that integrates multi-task learning with Large Language Models (LLMs) to effectively mine deep-seated vulnerability features. Specifically, we construct two auxiliary tasks beyond the vulnerability detection task. First, we utilize the vulnerability patches to construct a vulnerability localization task. Second, based on the vulnerability features extracted from patches, we leverage GPT-4 to construct a vulnerability interpretation task. VulLLM innovatively augments vulnerability classification by leveraging generative LLMs to understand complex vulnerability patterns, thus compelling the model to capture the root causes of vulnerabilities rather than overfitting to spurious features of a single task. The experiments conducted on six large datasets demonstrate that VulLLM surpasses seven state-of-the-art models in terms of effectiveness, generalization, and robustness.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑