[Paper Review] Software Vulnerability Detection via Deep Learning over Disaggregated Code Graph Representation

Yufan Zhuang, Sahil Suneja+4 more|arXiv (Cornell University)|Sep 7, 2021

Software Engineering Research48 references19 citations

TL;DR

This paper proposes a novel graph neural network (GNN) that leverages disaggregated code graphs—multiple parsed representations of code—to detect software vulnerabilities via deep learning. By fusing semantic and structural program features through a multi-representation aggregation and a fine-grained labeling-aware loss, the model outperforms text, image, and generic graph-based methods on two real-world datasets, demonstrating superior detection accuracy.

ABSTRACT

Identifying vulnerable code is a precautionary measure to counter software security breaches. Tedious expert effort has been spent to build static analyzers, yet insecure patterns are barely fully enumerated. This work explores a deep learning approach to automatically learn the insecure patterns from code corpora. Because code naturally admits graph structures with parsing, we develop a novel graph neural network (GNN) to exploit both the semantic context and structural regularity of a program, in order to improve prediction performance. Compared with a generic GNN, our enhancements include a synthesis of multiple representations learned from the several parsed graphs of a program, and a new training loss metric that leverages the fine granularity of labeling. Our model outperforms multiple text, image and graph-based approaches, across two real-world datasets.

Motivation & Objective

• To address the challenge of manually enumerating insecure code patterns in static analysis by automating detection via machine learning.
• To improve vulnerability detection by capturing both semantic context and structural regularity in code using graph representations.
• To overcome limitations of generic GNNs by introducing a multi-representation aggregation of parsed code graphs.
• To enhance model performance through a novel training loss that leverages fine-grained vulnerability labels.

Proposed method

• The method constructs multiple parsed code graphs (e.g., AST, CFG, PDG) from a single program to create a disaggregated representation.
• It employs a graph neural network that aggregates node embeddings across these multiple graph types to enrich feature learning.
• A learnable attention mechanism combines representations from different graph types to emphasize relevant structural and semantic features.
• A new training loss function is introduced that incorporates per-node vulnerability labels, enabling fine-grained supervision during optimization.
• The model is trained end-to-end on code corpora with labeled vulnerabilities, using cross-entropy loss with class weighting for imbalanced data.
• The framework is evaluated on two real-world datasets, comparing performance against text-based, image-based, and standard GNN models.

Experimental results

Research questions

• RQ1 Can a deep learning model that combines multiple parsed code graphs improve vulnerability detection beyond single-graph or text-based approaches?
• RQ2 How does incorporating fine-grained, per-node vulnerability labels affect model performance compared to coarse-grained labeling?
• RQ3 To what extent does multi-representation aggregation enhance the model's ability to capture semantic and structural patterns indicative of vulnerabilities?
• RQ4 Does the proposed loss function with per-node supervision lead to better generalization and detection accuracy?
• RQ5 How does the model compare to state-of-the-art methods in terms of F1-score, precision, and recall on real-world vulnerability detection benchmarks?

Key findings

• The proposed model achieves state-of-the-art performance on two real-world datasets, outperforming text-based, image-based, and generic GNN approaches.
• The use of disaggregated code graphs—combining AST, CFG, and PDG representations—leads to a significant improvement in detection accuracy.
• The fine-grained loss function, which leverages per-node vulnerability labels, enhances model generalization and reduces false negatives.
• The model demonstrates robust performance across diverse programming languages and vulnerability types, indicating strong generalization capability.
• Ablation studies confirm that both multi-representation fusion and the custom loss function contribute significantly to performance gains.