Training a Small Model to Classify 200+ Vulnerability Types — and Beating LLMs

Every CVE gets assigned a CWE — a label that tells you what kind of vulnerability it is. Buffer overflow, SQL injection, path traversal. There are hundreds of these categories, and they matter: CWEs drive prioritization, remediation workflows, and trend analysis across entire vulnerability databases.

The problem is that CWE labels in the NVD are noisy. They're inconsistently assigned, often too broad (hello, CWE-119), and sometimes just wrong. MITRE and NVD analysts do their best, but manually classifying 30,000+ CVEs per year across 200+ categories doesn't scale.

I wanted to see if a small, fine-tuned model could do this reliably. Turns out it can — and it competes with models 64 times its size.

The Dataset #

This was the hardest part. You can't just take NVD labels at face value — they're the thing you're trying to improve.

I started with the full NVD dataset and built a pipeline to refine the labels:

1. Pull all CVE descriptions with their assigned CWEs from NVD
2. Use Claude Sonnet 4.6 to re-classify each CVE independently based on the description text
3. Keep only samples where the NVD label and Claude's label agree — this filters out ambiguous or mislabeled entries

This agreement-based filtering is key. If a human analyst and a frontier LLM both arrive at the same CWE independently from the description alone, that label is probably correct. If they disagree, toss it — you don't want noisy labels in your training data.

The result: 234,770 training samples and 27,780 test samples across 205 CWE classes. The dataset is public on HuggingFace .

The Model #

Nothing exotic here. RoBERTa-base with a classification head. 125 million parameters.

I chose RoBERTa over newer architectures for a few reasons:

• CVE descriptions are short text (usually under 512 tokens). You don't need a 7B model for this.
• RoBERTa's tokenizer and pretraining are well-suited for technical English.
• Inference is fast and cheap — you can run this on a single CPU in production.

Fine-tuning was straightforward. Standard cross-entropy loss, AdamW optimizer, linear warmup. Nothing that would surprise you if you've fine-tuned a transformer before.

Results #

The headline numbers on the test set:

The Macro F1 looks lower than the accuracy because it equally weights all 205 classes — including rare CWEs with very few training samples. For the top 50 most common CWEs, per-class F1 scores are significantly higher.

Compared to a TF-IDF + logistic regression baseline (84.9% accuracy), the model gains +2.5 points on accuracy and +15.5 points on Macro F1. That F1 jump is where the real improvement lives — the model is much better at rare categories that TF-IDF can't reliably distinguish.

CTI-Bench Benchmark #

This is where it gets interesting. CTI-Bench is a NeurIPS 2024 benchmark for evaluating LLMs on cyber threat intelligence tasks, including CVE-to-CWE classification.

My 125M model scores 75.6% strict accuracy (95% CI: 72.8–78.2%) on CTI-Bench Task 2.

For comparison, Cisco Foundation-Sec-8B-Reasoning — an 8 billion parameter model specifically designed for security tasks — scores 75.3% on the same benchmark.

A model with 64x fewer parameters, trained on a single task, matches a purpose-built security LLM. The specialized model doesn't need to understand language in general — it just needs to map technical descriptions to vulnerability categories.

What Surprised Me #

Label quality matters more than model size. Early experiments with raw NVD labels topped out around 82% accuracy. The agreement-filtered dataset pushed that to 87.4% without changing anything about the model or training.

Rare classes are genuinely hard. Some CWEs have fewer than 20 training examples. The model struggles there — but so do humans. Some of these categories are inherently ambiguous even for security experts.

The gap between accuracy and F1 tells a story. 87.4% accuracy vs 60.7% Macro F1 means the model is great at common vulnerabilities but has room to improve on the long tail. Future work: data augmentation for rare classes, or a hierarchical classification approach using the CWE tree structure.

Using the Model #

The model and dataset are both publicly available:

• Model: xamxte/cwe-classifier-roberta-base
• Dataset: xamxte/cve-to-cwe
• Paper: arXiv:2603.14911
• Live demo: HuggingFace Space

Loading and running inference takes a few lines:

from transformers import pipeline

classifier = pipeline(
    "text-classification",
    model="xamxte/cwe-classifier-roberta-base",
    top_k=5
)

cve_description = """
A buffer overflow vulnerability in the HTTP parser of libcurl
allows remote attackers to execute arbitrary code via a crafted
URL that exceeds the expected buffer length.
"""

results = classifier(cve_description)
for r in results:
    print(f"{r['label']}: {r['score']:.3f}")

The model returns top-k predictions with confidence scores. In practice, looking at the top-3 predictions covers most edge cases where the primary CWE might be debatable.

What This Means #

There's a practical gap in vulnerability management: organizations need to classify and triage thousands of CVEs, and the existing labels aren't reliable enough for automation. A small, fast model that runs locally changes what's possible:

• Triage automation: Route vulnerabilities to the right team based on predicted CWE
• Label validation: Flag CVEs where the NVD label and model prediction disagree — those are worth a human review
• Trend analysis: Track vulnerability categories over time with consistent classification

The model is 125M parameters. It runs in milliseconds on CPU. You can deploy it anywhere — no GPU, no API calls, no data leaving your network.

The full paper has all the details on training, evaluation methodology, and benchmark comparisons.