Vul-CGBT：通过对比语义学习和图嵌入技术提升代码漏洞检测能力

《Knowledge-Based Systems》：Vul-CGBT: Enhancing code vulnerability detection via contrastive semantic learning and graph embedding

　　

杨张 | 王一凡

摘要

代码漏洞检测对于保护系统免受攻击至关重要。然而，现有的方法依赖于预训练模型的微调机制，无法理解漏洞代码的语义，因此难以捕捉漏洞代码与非漏洞代码之间的细微语义差异。此外，这些方法将代码视为序列输入，忽略了源代码的结构信息。为了解决这些问题，本文提出了一种新的代码漏洞检测方法 Vul-CGBT，该方法结合了对比语义学习和图嵌入技术，通过提取漏洞的深层语义和代码结构特征来提高检测性能。首先，使用掩码语言模型和对比语义学习对 Transformer 模型进行预训练，动量对比学习增强了其对漏洞语义模式的理解，从而生成代码序列嵌入。其次，通过 Tree-sitter 解析源代码的抽象语法树，构建控制流图和数据流图，然后利用图卷积网络学习图嵌入。最后，将序列嵌入和图嵌入融合后输入到全连接网络中进行检测。在 ReVeal、Big-Vul 和 CodeXGLUE-Defect 数据集上的实验结果表明，Vul-CGBT 的 F1 分别达到了 51.17%、61.15% 和 66.19%，比 PDBERT 和 LineVul 等基线方法高出 3.07%、2.25% 和 50.49%，同时其泛化能力也得到了验证。