Rethinking ML-based Vulnerability Detection Part-1

Rethinking ML-based Vulnerability Detection

Introduction

Hey everyone! It’s been a while—I’ve been down a really interesting rabbit hole in the meantime. The recent growth of AI has piqued my curiosity: how can AI be applied to security? Is it even possible to apply AI to such a mission-critical field when hallucinations can potentially cause losses in the millions? And how can we balance privacy and ethical concerns while still delivering meaningful results?

I went into this knowing full well that somebody else would’ve already explored this topic, and I could glean some insights from their efforts. Sure enough, there have been plenty of researchers tackling this very problem. But their work has also sparked new questions that I’d like to explore—and hopefully answer.

Scripts, datasets and other things have been uploaded to my github here

---

The Problem: Fundamental Questions in Vulnerability Detection

After reviewing the current state of ML-based vulnerability detection, I’ve identified several questions that remain murky(atleast to me):

1. Monolithic vs Decomposed Architectures

The field has largely focused on end-to-end monolithic models that approach vulnerability detection as a single classification task. Notable works like DiverseVul (Chen et al., 2023) and ReVeal (Chakraborty et al., 2021) leverage models such as CodeBERT, GraphCodeBERT, and NatGen - all building on a similar foundational approach:

Input Code → Single Model → Binary Classification

But here’s the thing - these papers report high precision but frustratingly low recall. Recent research has highlighted an even more concerning issue: when these models are tested on realistic datasets rather than their original benchmarks, the performance drops dramatically. For instance, DeepWukong (Cheng et al., 2021) - a sophisticated approach using graph neural networks - initially reported impressive metrics (98% recall, 93% F1-score). However, when evaluated on real-world code, researchers found that precision could decline by up to 95 percentage points. This means a model claiming 95% precision might actually achieve close to 0% on realistic data.

Source : https://surrealyz.github.io/files/pubs/raid23-diversevul.pdf

I’ve been pondering a different approach - a decomposed/cascade architecture. My reasoning is that different phases of vulnerability analysis require different optimizations. You can’t optimize for both high recall AND high precision in a single model, but what if you didn’t have to?

Key Question: Can decomposing vulnerability detection into specialized subtasks outperform monolithic approaches?

Supporting Evidence from Recent Research

A 2024 study on “Revisiting the Performance of Deep Learning-Based Vulnerability Detection” provides compelling evidence for this question. When evaluating state-of-the-art models like DeepWukong, LineVul, and ReVeal on realistic datasets, researchers observed:

• Models achieving 90%+ scores on synthetic benchmarks often fail catastrophically on real code
• The performance gap suggests these models overfit to their training data
• Even sophisticated graph neural network approaches struggle with real-world complexity

This validates the need for a fundamentally different architecture - one that doesn’t try to solve everything in a single pass.

2. Coder Models vs Reasoning Models

We now have two distinct families of LLMs:

• Code-specialized models: DeepSeek-Coder, CodeLlama, StarCoder
• Reasoning-optimized models: GPT-o1, Claude-3, QwQ

I haven’t read any papers regarding this, because as of now I have not found any :P

But I believe both offer their own pros and cons for vulnerability analysis. Some vulnerabilities are more related to syntax and some are more related to the actual logic. Logic dictates that reasoning models might perform well on logic-based flaws and code models might excel at code-specific vulnerabilities.

Key Question: Do code-pretrained models understand vulnerabilities differently than reasoning models? Which is more suitable for security analysis?

3. Program-Aware vs Semantic RAG

Recent work like Vul-RAG (Du et al., 2024) takes an interesting approach with retrieval-augmented generation, leveraging semantic similarity - essentially treating code like natural language. But this approach might not fully capture crucial program structure.

Here’s a concrete example. Consider this function:

void process_data(char* input) {
    char buffer[256];
    strcpy(buffer, input);  // No bounds checking!
    execute_command(buffer);
}

Semantic RAG might retrieve other functions using strcpy or buffer operations. But what we really need is the caller:

// Caller function
void handle_request(char* raw_input) {
    if (strlen(raw_input) > 255) {
        return;  // Input already validated!
    }
    process_data(raw_input);
}

Without this context, we’d flag a false positive. The vulnerability isn’t in process_data - it’s already protected by its caller!

Key Question: Should RAG for vulnerability detection build on program structure (caller/callee relationships, data flow) or continue exploring semantic similarity?

Quick Summary

After reading a bunch of papers regarding automated vulnerability detection, I have several questions that I would like to fully explore (to my limits):

• Systematic comparison of decomposition vs monolithic approaches
• Evaluation of modern coder vs reasoning models
• Exploring program-aware vs semantic RAG

My Approach: Mimicking Security Engineers

Based on preliminary analysis, here’s my take on these questions:

For Architecture: Task Decomposition

Current ML approaches treat vulnerability detection as a single-shot classification problem. But that’s not how security engineers work. We don’t stare at every function with equal intensity - we quickly scan for suspicious patterns, then deep-dive into anything that looks problematic.

My architecture mimics this workflow:

[High-Recall Filter]     [Context Retrieval]     [Deep Analysis]
     SLM (7B)        →    RAG (structural)   →    LLM (70B+)
   "Is this sus?"         "What calls this?"     "Is it actually vulnerable?"

A more robust explanation:

Stage 1: SLM as Security Intern

• Scans all functions quickly
• Optimized for high recall (catches ~95% of vulnerabilities)
• Lower precision is fine (false positives are cheap here)
• Think: “This function doesn’t validate input… flag it!”

Stage 2: RAG as Context Gatherer

• Retrieves callers, callees, data structures
• Builds the full vulnerability context
• Program-structure aware, not just semantic

Stage 3: LLM as Senior Engineer

• Performs expensive deep analysis
• Has full context from RAG
• Makes final determination

This decomposition lets us optimize each stage for what it does best. The SLM can be trigger-happy (high recall), while the LLM can be conservative (high precision).

For Model Selection: Both Are Needed

My hypothesis is that we need both coder and reasoning models, but at different stages. I’m experimenting with four combinations:

Model combination	SLM Model → LLM model	Comments
C-C	Coder → Coder	fast, syntax-focused
C-R	Coder → Reasoning	balanced
R-C	Reasoning → Coder	unusual, but worth testing
R-R	Reasoning → Reasoning	slow, but potentially most accurate

With a quick glance, I believe C-R might be optimal - coder models excel at spotting syntactic patterns (Stage 1), while reasoning models better understand the logic and context (Stage 3). But more rigorous evaluation is needed.

For RAG: Program Structure Matters

I believe program-aware/local RAG based on actual code structure (call graphs, data dependencies) will outperform semantic similarity. Vulnerabilities are often interprocedural - you can’t understand them without context. This requires further analysis and experimentation to correctly implement, currently I’m more focused on implementing the proper SLM → LLM pipeline

Model Selection and Dataset Details

For initial experiments, I’ve selected:

• SLM: Qwen2.5-Coder-7B (fine-tuned on 130k [~65k vulnerable/safe pairs] synthetic + 40k real vulnerable samples + clean code). I chose this because this is the best model that I can comfortablly fine tune using a consumer GPU
• LLM: [Experimenting with multiple options]

Why Qwen2.5-Coder?

1. Recent model with strong code understanding
2. 7B size offers good balance of capability and speed
3. Supports fine-tuning effectively
4. Open weights for reproducibility

An Additional Question: Classification Granularity

During dataset preparation, I encountered another fundamental question:

Should we use binary (vulnerable/safe) or ternary (vulnerable/safe/clean) classification?

Current approaches use binary classification, but I hypothesize that distinguishing between:

• Vulnerable: Contains active security flaws
• Safe/Fixed: Previously vulnerable, now patched
• Clean: Never had vulnerabilities/fixes

…might provide better signal for learning. This is something I’m actively investigating.

The SOTA Model Question

This raises another puzzling observation: Why does most research still use older models like CodeBERT (2020) or NatGen (2022)?

Recent models like:

• DeepSeek-Coder-V2
• Qwen2.5-Coder
• CodeLlama-70B
• StarCoder2

…have shown massive improvements in code understanding, yet vulnerability detection research hasn’t adopted them. Is this due to computational constraints?Or Something fundamental about vulnerability detection?

Interestingly, even newer approaches using advanced architectures face challenges. DeepWukong, despite using sophisticated graph neural networks to capture both control-flow and data-flow information, still struggles when moving from controlled benchmarks to real-world scenarios. This suggests the problem might be deeper than just model architecture - it could be about how we approach the task itself.

Dataset Preprocessing

My approach uses two data sources:

Synthetic Samples: NIST SARD

This dataset gave me the most trouble—not just because of its sheer size, but because of all the modifications needed to make it suitable for model training. It’s primarily designed as test cases for a massive 118 different CWEs.

The testcases are available here

Step 1: Function Extraction

The test cases are organized into multiple folders, one for each CWE. My first challenge was extracting individual functions to create a perfectly balanced dataset—ensuring every vulnerable function has a corresponding safe implementation.

The biggest hurdle? The wildly different formats. Two C programs with the same CWE114 vulnerability can look completely different because of various patterns—Sources, Sinks, Wrappers, multi-file vulnerabilities, constructors/destructors, you name it.

After analyzing the patterns, I identified 3 main formats in the dataset:

1. 2 function format (simplest) - 1 bad main function, 1 good main function - both self-contained with no external dependencies.

2. 4 function format - 2 bad functions (main + helper that serves as either sink or source) and 2 good functions (same structure but safe)

3. 6 function format - 3 bad functions (main, constructor, destructor) and 3 good functions (safe versions of each)

I wrote a Python script using regex parsing to extract all functions from each file. (Sure, AST parsing, tree-sitter, or ctags would’ve been more robust, but regex got the job done). BTW, some vulnerabilities span multiple files through function calls—too complex for diminishing returns, so I excluded anything spanning 2+ files.

Before extraction, I also randomized all variable names using ctags to prevent the model from learning based on revealing names. The final format is as follows,

  {
    "code": "... code here ...",
    "type": "bad",
    "file_hash": "a84f62d8198a961f98d07d7db30845b31aacd5519a8906c066e1167b8c6a5574",
    "function_name": "CWE114_Process_Control__w32_char_connect_socket_01_bad",
    "source_file": "CWE114_Process_Control__w32_char_connect_socket_01.c",
    "namespace": null
  }
  {
    "code": "... code here ...",
    "type": "good",
    "file_hash": "a84f62d8198a961f98d07d7db30845b31aacd5519a8906c066e1167b8c6a5574",
    "function_name": "goodG2B",
    "source_file": "CWE114_Process_Control__w32_char_connect_socket_01.c",
    "namespace": null
  }

Step 2: Preprocessing the Extracted Functions

Next came the cleanup phase. I wrote a script that:

• Removes wrappers and duplicate functions
• Keeps only one good implementation per bad function
• Ensures equal numbers of good and bad functions

I then paired these functions according to their formats (2, 4, or 6 functions). The final dataset has 6 fields: bad input, good input, bad output, good output, hash, and metadata. The input and output fields are identical except the output includes revealing comments. (I originally wanted to try contrastive learning with this structure, but asking a 7B model to both classify AND annotate vulnerabilities proved way too ambitious).

Finally, I randomized the remaining identifiers—structs, unions, global variables, and function names—to prevent the model from creating biased associations. the final format is as follows:

  {
    "hash": "0000fba7ab0607d14f9fb9c359b9b6e9bde01b2bf42e2178725b6e0ee5478d4c",
    "bad_input": "... randomized code here ...",
    "good_input": "... randomized code here ...",
    "bad_output": "... randomized code here ...",
    "good_output": "... randomized code here ...",
    "metadata": {
      "num_bad_functions": 1,
      "num_good_functions": 1,
      "bad_function_names": [
        "ahft_pac"
      ],
      "good_function_names": [
        "fqofpvau_hftcj"
      ],
      "source_file": "CWE121_Stack_Based_Buffer_Overflow__CWE805_wchar_t_alloca_loop_31.c",
      "name_mapping": {
        "CWE121_Stack_Based_Buffer_Overflow__CWE805_wchar_t_alloca_loop_31_bad": "ahft_pac",
        "goodG2B": "fqofpvau_hftcj"
      }
    }
  }

Step 3: Creating the Classification Dataset

With everything randomized, I extracted just the input fields (without comments) from the contrastive dataset and used the metadata to create two classification categories—one major, one minor. I’m still debating whether to include this in the classification task, since the reasoning/explanation is handled by the LLM, not the SLM. BTW, “SLM” and “LLM” are relative terms here—my 7B model is definitely small, just relatively smaller compared to SOTA LLMs. Final format is as follows,

  {
    "code": "... randomized input code ...",
    "is_vulnerable": true,
    "cwe": "CWE-121",
    "major_category": "Memory_Related",
    "specialized_category": "Memory_Management",
    "vulnerability_status": "Vulnerable"
  }

  {
    "code": "... randomized input code ...",
    "is_vulnerable": false,
    "cwe": "CWE-121",
    "major_category": "Memory_Related",
    "specialized_category": "Memory_Management",
    "vulnerability_status": "Fixed"
  }

Real-World Samples: DiverseVul

Extracting real-world samples was refreshingly straightforward, thanks to the excellent work behind DiverseVul. It combines well-known existing datasets with new additions, creating a substantial collection of ~41K vulnerable functions. I just needed a simple script to filter out wrapper functions and near-empty functions that only appeared in commits due to whitespace changes. the only thing to do now is to convert this format to the systhetic dataset format.

  {
    "project": "qemu",
    "commit_id": "bbc5842211cdd90103cfe52f2ca24afac880694f",
    "target": 1,
    "func": "... code here ...",
    "idx": 335528
  }
  {
    "project": "gpac",
    "commit_id": "bceb03fd2be95097a7b409ea59914f332fb6bc86",
    "target": 1,
    "func": "... code here ...",
    "idx": 3003
  }

Future Work: Questions for Part 2

In the next part, I’ll address these questions empirically:

1. Architecture Comparison:
   • Implement both monolithic and decomposed approaches
   • Rigorous comparison on same datasets
   • Analysis of failure modes
2. Model Comparison Matrix:
   • Systematic evaluation of C-C, C-R, R-C, R-R combinations
   • Task-specific benchmarks (memory bugs vs logic bugs)
   • Speed/accuracy tradeoffs
3. RAG Approaches:
   • Implement semantic similarity baseline
   • Build program-structure-aware retrieval
   • Compare on interprocedural vulnerabilities
4. Classification Granularity:
   • Train models with 2-class vs 3-class
   • Analyze impact on different vulnerability types
   • Developer study on usefulness
5. Modern Model Evaluation:
   • Benchmark latest models against papers’ baselines
   • Understand why adoption is slow
   • Cost-benefit analysis

Conclusion

The field of ML-based vulnerability detection is at a crossroads. We have powerful new models and techniques, but fundamental questions remain unanswered:

• Is end-to-end always better, or should we decompose?
• Which model family understands security better?
• How should we retrieve relevant context?
• Why aren’t we using SOTA models?

Part 1 has laid out these questions and my proposed approach. The key insight? Maybe we’ve been thinking about this wrong. Instead of trying to build a perfect monolithic model, we should mimic how security engineers actually work - quick scans followed by deep dives with full context.

Stay tuned for Part 2 where I’ll share experimental results!

---

References

1. Chen, Y., et al. (2023). “DiverseVul: A New Vulnerable Source Code Dataset for Deep Learning Based Vulnerability Detection”
2. Chakraborty, S., et al. (2021). “Deep Learning based Vulnerability Detection: Are We There Yet?”
3. Du, X., et al. (2024). “Vul-RAG: Enhancing LLM-based Vulnerability Detection via Knowledge-level RAG”
4. NIST. “Software Assurance Reference Dataset (SARD)”

---