| | --- |
| | license: cc-by-4.0 |
| | language: |
| | - en |
| | library_name: transformers |
| | pipeline_tag: text-classification |
| | tags: |
| | - cybersecurity |
| | - vulnerability |
| | - cwe |
| | - cve |
| | - nvd |
| | - roberta |
| | base_model: FacebookAI/roberta-base |
| | datasets: |
| | - xamxte/cve-to-cwe |
| | metrics: |
| | - accuracy |
| | - f1 |
| | model-index: |
| | - name: cwe-classifier-roberta-base |
| | results: |
| | - task: |
| | type: text-classification |
| | name: CWE Classification |
| | dataset: |
| | name: cve-to-cwe (test split) |
| | type: xamxte/cve-to-cwe |
| | split: test |
| | metrics: |
| | - name: Top-1 Accuracy |
| | type: accuracy |
| | value: 0.8744 |
| | - name: Top-3 Accuracy |
| | type: accuracy |
| | value: 0.9467 |
| | - name: Macro F1 |
| | type: f1 |
| | value: 0.6071 |
| | - task: |
| | type: text-classification |
| | name: CWE Classification (CTI-Bench) |
| | dataset: |
| | name: CTI-Bench cti-rcm |
| | type: xashru/cti-bench |
| | metrics: |
| | - name: Strict Top-1 |
| | type: accuracy |
| | value: 0.756 |
| | - name: Hierarchy-aware Top-1 |
| | type: accuracy |
| | value: 0.865 |
| | --- |
| | |
| | # 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. |
| |
|
| | | Model | Params | Type | Top-1 Accuracy | Source | |
| | |-------|--------|------|---------------|--------| |
| | | [Sec-Gemini v1](https://security.googleblog.com/2025/04/google-launches-sec-gemini-v1-new.html) (Google)* | — | closed | ~86% | Google Security Blog | |
| | | [SecLM](https://security.googlecloudcommunity.com/community-blog-42/fueling-ai-innovation-in-secops-products-the-seclm-platform-and-sec-gemini-research-pipeline-3997) (Google)* | — | closed | ~85% | Google Cloud Blog | |
| | | **This model** | **125M** | **open** | **75.6%** | — | |
| | | [Foundation-Sec-8B-Reasoning](https://arxiv.org/abs/2601.21051) (Cisco) | 8B | open | 75.3% | arXiv 2601.21051 | |
| | | [GPT-4](https://arxiv.org/abs/2406.07599) | ~1.7T | closed | 72.0% | CTI-Bench paper | |
| | | [Foundation-Sec-8B](https://arxiv.org/abs/2504.21039) (Cisco) | 8B | open | 72.0% (±1.7%) | arXiv 2504.21039 | |
| | | [WhiteRabbitNeo-V2-70B](https://arxiv.org/abs/2504.21039) | 70B | open | 71.1% | arXiv 2504.21039 | |
| | | [Foundation-Sec-8B-Instruct](https://arxiv.org/abs/2601.21051) (Cisco) | 8B | open | 70.4% | arXiv 2601.21051 | |
| | | [Llama-Primus](https://huggingface.co/trend-cybertron/Llama-Primus-Base) (Trend Micro) | 8B | open | 67.8% | HuggingFace | |
| | | [GPT-3.5](https://arxiv.org/abs/2406.07599) | ~175B | closed | 67.2% | CTI-Bench paper | |
| | | [Gemini 1.5](https://arxiv.org/abs/2406.07599) | — | closed | 66.6% | CTI-Bench paper | |
| | | [LLaMA3-70B](https://arxiv.org/abs/2406.07599) | 70B | open | 65.9% | CTI-Bench paper | |
| | | [LLaMA3-8B](https://arxiv.org/abs/2406.07599) | 8B | open | 44.7% | CTI-Bench paper | |
| |
|
| | *\*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 |
| |
|
| | ```python |
| | from transformers import pipeline |
| | |
| | classifier = pipeline("text-classification", model="xamxte/cwe-classifier-roberta-base", top_k=3) |
| | |
| | result = classifier("A SQL injection vulnerability in the login page allows remote attackers to execute arbitrary SQL commands via the username parameter.") |
| | print(result) |
| | # [[{'label': 'CWE-89', 'score': 0.95}, {'label': 'CWE-564', 'score': 0.02}, ...]] |
| | ``` |
| |
|
| | ### Manual inference |
| |
|
| | ```python |
| | from transformers import AutoTokenizer, AutoModelForSequenceClassification |
| | import torch |
| | import json |
| | |
| | model_name = "xamxte/cwe-classifier-roberta-base" |
| | tokenizer = AutoTokenizer.from_pretrained(model_name) |
| | model = AutoModelForSequenceClassification.from_pretrained(model_name) |
| | |
| | # Load label map |
| | from huggingface_hub import hf_hub_download |
| | label_map_path = hf_hub_download(repo_id=model_name, filename="cwe_label_map.json") |
| | with open(label_map_path) as f: |
| | label_map = json.load(f) |
| | id_to_label = {v: k for k, v in label_map.items()} |
| | |
| | # Predict |
| | text = "CVE Description: A buffer overflow in the PNG parser allows remote code execution via crafted image files." |
| | inputs = tokenizer(text, return_tensors="pt", max_length=384, truncation=True, padding=True) |
| | |
| | with torch.no_grad(): |
| | logits = model(**inputs).logits |
| | |
| | top3 = torch.topk(logits, 3) |
| | for score, idx in zip(top3.values[0], top3.indices[0]): |
| | print(f"{id_to_label[idx.item()]}: {score.item():.3f}") |
| | ``` |
| |
|
| | ## Training |
| |
|
| | - **Base model:** FacebookAI/roberta-base (125M params) |
| | - **Dataset:** [xamxte/cve-to-cwe](https://huggingface.co/datasets/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](https://huggingface.co/datasets/xamxte/cve-to-cwe) 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 |
| | |
| | 📄 **[Fine-tuning RoBERTa for CVE-to-CWE Classification: A 125M Parameter Model Competitive with LLMs](https://arxiv.org/abs/2603.14911)** |
| | |
| | ## Citation |
| | |
| | ```bibtex |
| | @article{mosievskiy2026cwe, |
| | title={Fine-tuning RoBERTa for CVE-to-CWE Classification: A 125M Parameter Model Competitive with LLMs}, |
| | author={Mosievskiy, Nikita}, |
| | journal={arXiv preprint arXiv:2603.14911}, |
| | year={2026} |
| | url={https://huggingface.co/xamxte/cwe-classifier-roberta-base} |
| | } |
| | ``` |
| | |
