DEV Community

turingrtss
turingrtss

Posted on

I Trained a Vulnerability Detection Model. It Was Reading Comments, Not Code.

I trained a simple ML model to detect code vulnerabilities and got 84.7% accuracy. Then I looked at what it actually learned.

It was reading comments, not code.

The Experiment

Hardware: ARM64, 12GB RAM, 2 CPU cores, no GPU. Everything runs on a $0/month Oracle Cloud free tier.

Dataset: 8,480 code samples from HuggingFace, labeled as safe or with specific CWE vulnerability types (SQL injection, command injection, XSS, etc.). Roughly 50/50 split between safe and vulnerable.

Models: TF-IDF character n-grams (3-6 chars) fed into three classifiers:

Model Accuracy F1 False Positive Rate Train Time
Logistic Regression 84.7% 0.842 16.7% 0.5s
Gradient Boosting 83.7% 0.835 19.5% 175s
Random Forest 82.7% 0.821 17.8% 18s

Logistic regression won. The simplest model, trained in half a second, beat both ensemble methods. That was the first hint something was off.

What the Model Actually Learned

I extracted the top features the logistic regression model uses to classify code.

Top features indicating VULNERABLE code:

+1.41  " this "
+1.33  " eval("
+1.30  " + "
+1.23  " vuln"
+1.23  " vul"
+1.20  "vul"
+1.20  "uln"
Enter fullscreen mode Exit fullscreen mode

Top features indicating SAFE code:

-2.06  " sa"
-1.81  " saf"
-1.81  " safe"
-1.78  "secu"
-1.78  "secur"
-1.77  " if"
Enter fullscreen mode Exit fullscreen mode

See the problem?

The model strongest vulnerability signal is substrings of the word "vulnerable." Its strongest safety signal is substrings of "safe" and "secure." These appear in comments, docstrings, and variable names within the dataset.

Only one genuinely code-related feature made the top 10: eval(, which is a known dangerous Python pattern.

Label Leakage

This is called label leakage: the classification labels are partially encoded in the input data through natural language descriptions.

The model learns that "vulnerable" in a comment means the label is "vulnerable." It does not need to understand the eval() call at all.

The 84.7% accuracy is partly measuring the model ability to read English comments, not detect code vulnerabilities.

Multi-class Results Confirm It

When predicting specific vulnerability types (14 classes), accuracy drops to 72.5%. Classes with fewer than 50 samples completely fail. And the model probably succeeds on deserialization partly because the word "pickle" correlates with both the label and the actual pattern.

What This Means

  1. Reported accuracy is inflated. Any model trained on code with descriptive comments risks learning descriptions instead of patterns.
  2. The model will fail on real code. Production code does not label itself as vulnerable.
  3. Dataset construction matters more than model architecture. A transformer would score higher on this dataset and be equally fooled.
  4. You can catch this. Always inspect top features. If your model strongest signals are English words rather than code patterns, your benchmark is broken.

Next Steps

  1. Strip all comments and docstrings before training
  2. Normalize variable names to generic tokens
  3. Re-evaluate accuracy on clean data
  4. Use CVE fix commits (before/after pairs from real patches)

The true model capability is whatever accuracy remains after removing the label leakage.

Reproducibility

  • Code: github.com/turingrtss/vulndetect
  • Dataset: lemon42-ai/Code_Vulnerability_Labeled_Dataset on HuggingFace
  • Full research paper (LaTeX/PDF): in the repo
  • Training time: Under 4 minutes on 2-core ARM CPU, no GPU

This is part of an ongoing research project on ML-based vulnerability detection. Next: stripping comments and re-evaluating with clean data.

Top comments (0)