CWE Classifier (RoBERTa-base)
A fine-tuned RoBERTa-base model that maps CVE (Common Vulnerabilities and Exposures) descriptions to CWE (Common Weakness Enumeration) categories. 125M parameters, 205 CWE classes.
Performance
Internal Test Set (27,780 agreement-filtered samples)
| Metric | Score |
|---|
| Top-1 Accuracy | 87.4% |
| Top-3 Accuracy | 94.7% |
| Macro F1 | 0.607 |
| Weighted F1 | 0.872 |
CTI-Bench External Benchmark (NeurIPS 2024, zero training overlap)
| Benchmark | Strict Top-1 | Hierarchy-aware Top-1 |
|---|
| cti-rcm (2023-2024 CVEs) | 75.6% | 86.5% |
| cti-rcm-2021 (2011-2021 CVEs) | 71.8% | 82.8% |
Comparison on CTI-Bench cti-rcm (strict exact match)
All scores below use the official CTI-Bench evaluation protocol: strict exact CWE ID match.
*Sec-Gemini and SecLM scores are approximate, estimated from published comparison charts. Exact values were not reported.
Competitive with the best open-weight models at 64x fewer parameters (125M vs 8B). Note: the 0.3pp difference vs Cisco Foundation-Sec-8B-Reasoning is not statistically significant (95% CIs overlap on n=1000). The Cisco models are general-purpose LLMs; ours is a task-specific encoder.
TF-IDF baseline comparison
A TF-IDF + Logistic Regression baseline reaches 84.9% top-1 on the same test set, but only 45.2% Macro F1 vs our 60.7% — a +15.5pp Macro F1 gap showing the model's advantage on rare CWE classes that keyword matching cannot handle.
Hierarchy-aware evaluation (supplementary)
This model predicts specific child CWEs (e.g., CWE-121 Stack Buffer Overflow) while CTI-Bench ground truth often uses generic parent CWEs (e.g., CWE-119 Buffer Overflow). When parent↔child equivalences are counted as correct:
| Benchmark | Strict Top-1 | Hierarchy-aware Top-1 |
|---|
| cti-rcm (2023-2024 CVEs) | 75.6% | 86.5% (+10.9pp) |
| cti-rcm-2021 (2011-2021 CVEs) | 71.8% | 82.8% (+11.0pp) |
Note: Other models in the table above were evaluated with strict matching only. Hierarchy-aware scores are not directly comparable and are shown separately for transparency.
Usage
1from transformers import pipeline
2
3classifier = pipeline("text-classification", model="xamxte/cwe-classifier-roberta-base", top_k=3)
4
5result = classifier("A SQL injection vulnerability in the login page allows remote attackers to execute arbitrary SQL commands via the username parameter.")
6print(result)
7# [[{'label': 'CWE-89', 'score': 0.95}, {'label': 'CWE-564', 'score': 0.02}, ...]]
Manual inference
1from transformers import AutoTokenizer, AutoModelForSequenceClassification
2import torch
3import json
4
5model_name = "xamxte/cwe-classifier-roberta-base"
6tokenizer = AutoTokenizer.from_pretrained(model_name)
7model = AutoModelForSequenceClassification.from_pretrained(model_name)
8
9# Load label map
10from huggingface_hub import hf_hub_download
11label_map_path = hf_hub_download(repo_id=model_name, filename="cwe_label_map.json")
12with open(label_map_path) as f:
13 label_map = json.load(f)
14id_to_label = {v: k for k, v in label_map.items()}
15
16# Predict
17text = "CVE Description: A buffer overflow in the PNG parser allows remote code execution via crafted image files."
18inputs = tokenizer(text, return_tensors="pt", max_length=384, truncation=True, padding=True)
19
20with torch.no_grad():
21 logits = model(**inputs).logits
22
23top3 = torch.topk(logits, 3)
24for score, idx in zip(top3.values[0], top3.indices[0]):
25 print(f"{id_to_label[idx.item()]}: {score.item():.3f}")
Training
- Base model: FacebookAI/roberta-base (125M params)
- Dataset: xamxte/cve-to-cwe — 234,770 training samples with Claude Sonnet 4.6 refined labels
- Training method: Two-phase fine-tuning
- Phase 1: Freeze first 8/12 transformer layers, train classifier head (4 epochs, lr=1e-4)
- Phase 2: Unfreeze all layers, full fine-tuning (9 epochs, lr=2e-5)
- Key hyperparameters: max_length=384, batch_size=32, label_smoothing=0.1, cosine scheduler, bf16
- Hardware: NVIDIA RTX 5080 (16GB), ~4 hours total
- Framework: HuggingFace Transformers + PyTorch
Label Quality
Training labels were refined using Claude Sonnet 4.6 via the Anthropic Batch API (~$395 total cost). The test/validation sets contain only agreement-filtered samples where NVD and Sonnet labels agree (73.1% exact match; 84.5% with hierarchy-aware matching). This biases evaluation toward unambiguous cases — real-world accuracy on arbitrary NVD entries will be lower. See the
dataset card for details.
CWE Hierarchy
This model predicts specific (child) CWE categories where possible. For example, buffer overflows are classified as CWE-121 (Stack) or CWE-122 (Heap) rather than the generic CWE-119. This provides more actionable information for vulnerability triage, but means strict accuracy on benchmarks using parent CWEs appears lower than actual performance.
Limitations
- 205 CWE classes only: Covers the most common CWEs in NVD. Rare CWEs not in the training set will be mapped to the closest known class.
- English only: Trained on English CVE descriptions from NVD.
- Description-based: Uses only the text description, not CVSS scores, CPE, or other metadata.
- Single-label: Predicts one CWE per CVE, though some vulnerabilities may involve multiple weakness types.
Paper
Citation
1@article{mosievskiy2026cwe,
2 title={Fine-tuning RoBERTa for CVE-to-CWE Classification: A 125M Parameter Model Competitive with LLMs},
3 author={Mosievskiy, Nikita},
4 journal={arXiv preprint arXiv:2603.14911},
5 year={2026}
6 url={https://huggingface.co/xamxte/cwe-classifier-roberta-base}
7}
