Multi-Label Vulnerability Detection with GraphCodeBERT

A graph-aware transformer framework for simultaneous detection of CWE-89 (SQL Injection), CWE-362 (Race Conditions), and CWE-639 (IDOR) vulnerabilities in source code — treating vulnerability detection as a multi-label classification problem rather than per-class binary detection.

---

The Problem

Most vulnerability detectors are trained per-CWE and run sequentially — three models, three scans, three false-positive pools. Real-world code often contains co-occurring vulnerabilities (e.g., an IDOR that also enables a race condition). A single-label approach misses these patterns entirely.

This work reframes detection as multi-label classification: one pass over the code, simultaneous predictions across all three CWE classes, with inter-label dependency captured by the model.

---

Architecture

Source Code (C/Python/Java)
        │
        ▼
┌──────────────────────┐
│   GraphCodeBERT      │  ← Pre-trained on code + data-flow graphs
│   (microsoft/gcb)    │
└──────────┬───────────┘
           │  [CLS] embedding (768-dim)
           ▼
┌──────────────────────┐
│  Multi-Label Head    │  ← Linear(768 → 3) + Sigmoid
│  CWE-89 | 362 | 639  │
└──────────────────────┘
           │
           ▼
  [0.91, 0.07, 0.84]   ← per-class probability vector

Line-level scanning is handled by scanner_v2.py, which slides a context window over function bodies and aggregates predictions at the line level.

---

Results

Multi-Label Classification Performance

CWE	Precision	Recall	F1
CWE-89 (SQL Injection)	0.921	0.934	0.927
CWE-362 (Race Condition)	0.883	0.871	0.877
CWE-639 (IDOR)	0.908	0.896	0.902
Macro Average	0.904	0.900	0.902

Macro F1: 0.9049 across all three classes on held-out test set.

Comparison Against Baselines

Method	Macro F1
Semgrep (rule-based)	0.571
CodeBERT (single-label, per-class)	0.841
GraphCodeBERT Multi-Label (ours)	0.905

---

Dataset

~7,500 labeled samples across three CWE classes, sourced and curated from:

• NVD-linked open-source CVE patches
• NIST SARD synthetic benchmarks
• Manual augmentation for low-frequency patterns

Labels are non-exclusive — samples can carry multiple CWE tags, enabling the model to learn co-occurrence patterns.

Split	Samples
Train	~6,000
Val	~750
Test	~750

---

Key Design Choices

• GraphCodeBERT over CodeBERT — encodes data-flow edges in addition to token sequences, which is critical for race conditions (CWE-362) where the vulnerability is in the ordering of operations, not the operations themselves.
• Sigmoid over Softmax — each label gets an independent probability; no competition between classes.
• Macro F1 as the target metric — class imbalance is real; micro-F1 would mask poor performance on rare labels.
• Line-level scanner — function-level labels are useful for training, but developers need line-level output for triage. scanner_v2.py bridges this.

---

Internship Context

Built during a 6-month research internship at DRDO DIT-CS (Defence Research & Development Organisation, Centre for Artificial Intelligence) under Dr. Deepak Vishwakarma.

The project was part of a broader effort to integrate ML-based vulnerability detection into DRDO's static analysis pipeline for defence software auditing.

---

Publication

Paper currently under review. Code will be released upon acceptance.

If you're interested in the dataset, methodology, or collaboration — reach out via the contact below.

Contact: [j.17atishay@gmail.com]

---