← AI Engineering
NLP & LLM Systems · Text Classification · CVE Data

CVE Classification: When the Label Is the Bug

Fine-tuning DistilBERT to sort CVE descriptions into CWE vulnerability categories. The real story isn't the fine-tuning itself — it's an error-analysis process that found systematic label noise in the source taxonomy, fixed what was genuinely fixable, and correctly left a second lookalike confusion alone once it turned out to be real ambiguity, not a bug.

View on GitHub Dataset source (NVD feed)
Python PyTorch Hugging Face DistilBERT scikit-learn Google Colab

The problem

Given a CVE's free-text description, predict its CWE vulnerability category — the same triage step a security team does by hand when new vulnerabilities land. The data comes from three years (2023-2025) of the NVD feed, 115,910 CVEs combined — reconstructed from a community GitHub mirror, since NIST deprecated the official bulk JSON feed in December 2023 in favor of a paginated API. Filtered to CVEs with an English description and a real (non-placeholder) CWE category, then restricted to the 15 most frequent categories: 60,858 examples. The categories are far from balanced — the most common, Cross-Site Scripting, outnumbers the rarest kept category almost 14 to 1, which shaped the loss function and the evaluation metric used throughout.

Horizontal bar chart showing the count of CVEs per CWE category, ranked from CWE-79 with over 20,000 examples down to CWE-94 with under 1,500, illustrating a steep, long-tailed class imbalance
Class counts across the 15 most frequent CWE categories, 2023-2025 combined. The ~14x gap between the largest and smallest class drove the choice of a class-weighted loss and F1 macro as the primary metric, not accuracy.

Key results

94.0%
Final test accuracy
0.89
Final F1 macro
8.9 min
Training time (fp16, T4 GPU)
14
Final vulnerability categories

Finding label noise in the source taxonomy

CWE-74 ("Improper Neutralization of Special Elements") is the generic parent category of CWE-89 (SQL Injection) in the CWE tree. A first baseline model struggled badly on CWE-74 — F1 of 0.49, far below every other class — and the confusion matrix pointed at one specific culprit: CWE-89 examples being misclassified as CWE-74 far more than chance would explain. Rather than tune around it, the actual text was checked.

Two of the misclassified examples, verbatim from the dataset, both labeled CWE-74: "...The manipulation leads to sql injection. The attack may be initiated remotely..." and "...leads to sql injection. It is possible to initiate the attack remotely..." — 73.5% of all CWE-74 examples explicitly say "sql injection" in the description, most likely CVEs filed by CNAs whose automated tooling only ever emits the generic label.

These 1,127 records were relabeled to CWE-89 before re-training. CWE-89's F1 rose from 0.83 to 0.99 — but interestingly, the F1 macro average barely moved from this step alone. Macro-averaging weights every class equally regardless of size, so a large gain on an 851-example class and almost no change on a 41-example one cancel out in the average; accuracy and weighted-F1, which are size-sensitive, are what actually showed the improvement. What remained of CWE-74 after the fix — 407 examples of CRLF injection, LDAP injection, and other unrelated mechanisms with no shared textual pattern — was excluded rather than kept as an artificially difficult 15th class.

When a confusion turns out not to be a bug

The same hypothesis was tested on the model's next-weakest classes: CWE-77 (generic Command Injection) and CWE-78 (OS Command Injection), another parent/child pair in the CWE taxonomy. This time the check came back negative — only 9.2% of CWE-77 examples explicitly mention "OS command," nowhere near the 73.5% seen with CWE-74. There is no systematic mislabeling to fix here. The confusion instead reflects genuine semantic overlap: most CWE-77 and CWE-78 descriptions read almost identically ("command injection vulnerability... allows arbitrary command execution"), and the choice between the generic and specific label is a judgment call made inconsistently by whoever files the CVE — not a bug in the model or the data.

Row-normalized confusion matrix heatmap for the final 14-class model, with most classes showing a clean diagonal and a highlighted block showing CWE-77, CWE-78, and CWE-94 confused with each other
Final model's confusion matrix. Every class outside the highlighted block is close to a clean diagonal; CWE-77/CWE-78/CWE-94 (outlined) are the one genuine trouble spot — and, per the text-evidence check above, one worth documenting rather than chasing further.

Results across all three runs

RunChangeClassesAccuracyF1 macro
1Baseline150.900.85
2Relabeled CWE-74 → CWE-89 (text-verified)150.930.85
3Excluded remaining CWE-74 (heterogeneous residual)140.940.89
Two different-looking data problems, two different verdicts: CWE-74/CWE-89 was genuine label noise, fixable and fixed with a clear before/after gain. CWE-77/CWE-78 looked identical on the confusion matrix but turned out to be real semantic ambiguity in the taxonomy itself — checked with the same method, and correctly left alone. The training loop was also sped up roughly 5.6x along the way (50 min → 8.9 min) via fp16 mixed precision and a right-sized max sequence length, verified against the actual token-length distribution rather than assumed.

Tools used

Python PyTorch Hugging Face Transformers Hugging Face Datasets scikit-learn Google Colab