[Paper Review] 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 Research4 citations

TL;DR

Vul-LMGNN proposes a unified model that fuses pre-trained code language models (CodeBERT) with a code property graph (CPG) enhanced by syntax, control flow, and data dependency information, using a gated Graph Neural Network (GGNN) to jointly learn sequence and structural features; it achieves state-of-the-art performance, with a 10% higher F1 score on small-scale datasets compared to prior methods.

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.

Motivation & Objective

• To address the limitations of sequence-based and graph-based code vulnerability detection methods, which often miss structural dependencies or long-range context.
• To unify local semantic features from pre-trained code language models with global structural information from heterogeneous code property graphs.
• To improve vulnerability detection accuracy by jointly training a code language model and a gated GNN to capture complex code dependencies.
• To enhance model robustness and generalization through a linear interpolation classifier that fuses predictions from Vul-LMGNN and CodeBERT.
• To demonstrate the effectiveness of the proposed architecture across multiple real-world vulnerability datasets.

Proposed method

• Constructs a code property graph (CPG) integrating Abstract Syntax Trees (AST), Control Flow Graphs (CFG), and Program Dependence Graphs (PDG) to represent syntax, control, and data flow.
• Initializes node embeddings using a fine-tuned CodeBERT model to inject local semantic context into the CPG.
• Employs a Gated Gated Neural Network (GGNN) with GRU units to iteratively aggregate messages from neighbors, capturing long-range dependencies and heterogeneous code features.
• Jointly trains the CodeBERT and GGNN modules end-to-end to enable implicit fusion of sequential and structural code representations.
• Uses a linear interpolation classifier that combines predictions from Vul-LMGNN and a pre-trained CodeBERT model to improve final detection performance.
• Fine-tunes the CodeBERT model on target datasets to adapt its embeddings to vulnerability-specific patterns.

Experimental results

Research questions

• RQ1 Can combining pre-trained code language models with a unified code property graph improve vulnerability detection performance?
• RQ2 How effective is a gated GNN in capturing heterogeneous dependencies (syntax, control, data flow) in code graphs compared to standard GNNs?
• RQ3 Does joint training of code language models and GNNs lead to better feature learning than isolated training?
• RQ4 How does the linear interpolation of Vul-LMGNN and CodeBERT predictions affect overall detection accuracy?
• RQ5 What is the impact of different GNN architectures (GGNN, GCN, GAT) on vulnerability detection performance when combined with code language models?

Key findings

• Vul-LMGNN achieves state-of-the-art performance across four real-world vulnerability datasets, outperforming six SOTA methods.
• On small-scale datasets, Vul-LMGNN achieves approximately 10% higher F1 score than prior approaches, indicating strong generalization on limited data.
• Fine-tuning CodeBERT on the target dataset for node embedding initialization yields the best performance (F1: 83.87%), surpassing direct use of pre-trained weights.
• The GGNN-based model outperforms GCN and GAT across all metrics, with an F1 score of 83.87% compared to 82.15% (GCN) and 78.39% (GAT).
• The GGNN's GRU-based message passing mechanism effectively captures long-range dependencies and heterogeneous code features, explaining its superior performance.
• The auxiliary linear interpolation classifier enhances model performance by explicitly fusing predictions from Vul-LMGNN and CodeBERT, demonstrating the value of ensemble learning in vulnerability detection.