[论文解读] Vul-LMGNNs: Fusing language models and online-distilled graph neural networks for code vulnerability detection

Ruitong Liu, Yanbin Wang+5 more|arXiv (Cornell University)|Apr 23, 2024

Software Engineering Research被引用 4

一句话总结

Vul-LMGNN 提出了一种统一模型，将预训练代码语言模型（CodeBERT）与融合了语法、控制流和数据依赖信息的代码属性图（CPG）相结合，利用门控图神经网络（GGNN）联合学习序列特征与结构特征；该模型实现了最先进性能，在小规模数据集上的 F1 分数相比先前方法高出约 10%。

ABSTRACT

Code Language Models (codeLMs) and Graph Neural Networks (GNNs) are widely used in code vulnerability detection. However, GNNs often rely on aggregating information from adjacent nodes, limiting structural information propagation across layers. While codeLMs can supplement GNNs with semantic information, existing integration methods underexplore their collaborative potential. To address these challenges, we propose Vul-LMGNNs, integrating pre-trained codeLMs with GNNs to enable cross-layer propagation of semantic and structural information. Vul-LMGNNs leverage Code Property Graphs (CPGs) to incorporate syntax, control flow, and data dependencies, using gated GNNs for structural extraction. An online knowledge distillation (KD) mechanism allows a student GNN to capture structural information from a trained counterpart via alternating training. Additionally, an "implicit-explicit" joint training framework leverages codeLMs to initialize embeddings and propagate code semantics. In the explicit phase, it performs late fusion via linear interpolation. Evaluations on real-world vulnerability datasets show Vul-LMGNNs outperform 17 state-of-the-art approaches. Source code is available at: https://github.com/Vul-LMGNN/vul-LMGNN.

研究动机与目标

• 为解决基于序列和基于图的代码漏洞检测方法的局限性，这些方法常会遗漏结构依赖关系或长距离上下文信息。
• 统一预训练代码语言模型中的局部语义特征与异构代码属性图中的全局结构信息。
• 通过端到端联合训练代码语言模型与门控 GNN，捕捉复杂代码依赖关系，提升漏洞检测准确率。
• 通过线性插值分类器融合 Vul-LMGNN 与 CodeBERT 的预测结果，增强模型鲁棒性与泛化能力。
• 在多个真实世界漏洞数据集上验证所提出架构的有效性。

提出的方法

• 构建一个代码属性图（CPG），整合抽象语法树（AST）、控制流图（CFG）和程序依赖图（PDG），以表示语法、控制与数据流信息。
• 使用微调后的 CodeBERT 模型初始化节点嵌入，将局部语义上下文注入 CPG。
• 采用具有 GRU 单元的门控图神经网络（GGNN），通过迭代聚合邻居消息，捕捉长距离依赖关系与异构代码特征。
• 端到端联合训练 CodeBERT 与 GGNN 模块，实现序列与结构代码表示的隐式融合。
• 使用线性插值分类器，结合 Vul-LMGNN 与预训练 CodeBERT 模型的预测结果，提升最终检测性能。
• 在目标数据集上微调 CodeBERT 模型，使其嵌入更适应漏洞特定的模式。

实验结果

研究问题

• RQ1 将预训练代码语言模型与统一的代码属性图相结合，能否提升漏洞检测性能？
• RQ2 与标准 GNN 相比，门控 GNN 在捕捉代码图中异构依赖关系（语法、控制流、数据流）方面效果如何？
• RQ3 代码语言模型与 GNN 的联合训练是否能带来比独立训练更优的特征学习效果？
• RQ4 Vul-LMGNN 与 CodeBERT 预测结果的线性插值对整体检测准确率有何影响？
• RQ5 当与代码语言模型结合时，不同 GNN 架构（GGNN、GCN、GAT）对漏洞检测性能的影响如何？

主要发现

• Vul-LMGNN 在四个真实世界漏洞数据集上均达到最先进性能，超越六种 SOTA 方法。
• 在小规模数据集上，Vul-LMGNN 的 F1 分数相比先前方法高出约 10%，表明其在数据有限情况下具有强大的泛化能力。
• 在目标数据集上微调 CodeBERT 以初始化节点嵌入可获得最佳性能（F1: 83.87%），优于直接使用预训练权重。
• GGNN 模型在所有指标上均优于 GCN 与 GAT，F1 分数达 83.87%，高于 GCN 的 82.15% 与 GAT 的 78.39%。
• GGNN 的基于 GRU 的消息传递机制能有效捕捉长距离依赖关系与异构代码特征，解释了其优越性能。
• 辅助的线性插值分类器通过显式融合 Vul-LMGNN 与 CodeBERT 的预测结果，提升了模型性能，证明了集成学习在漏洞检测中的价值。