Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Fine-tuning RoBERTa for CVE-to-CWE Classification: A 125M Parameter Model Competitive with LLMs

Nikita Mosievskiy

Abstract

We present a fine-tuned RoBERTa-base classifier (125M parameters) for mapping Common Vulnerabilities and Exposures (CVE) descriptions to Common Weakness Enumeration (CWE) categories. We construct a large-scale training dataset of 234,770 CVE descriptions with AI-refined CWE labels using Claude Sonnet 4.6, and agreement-filtered evaluation sets where NVD and AI labels agree. On our held-out test set (27,780 samples, 205 CWE classes), the model achieves 87.4% top-1 accuracy and 60.7% Macro F1 -- a +15.5 percentage-point Macro F1 gain over a TF-IDF baseline that already reaches 84.9% top-1, demonstrating the model's advantage on rare weakness categories. On the external CTI-Bench benchmark (NeurIPS 2024), the model achieves 75.6% strict accuracy (95% CI: 72.8-78.2%) -- statistically indistinguishable from Cisco Foundation-Sec-8B-Reasoning (75.3%, 8B parameters) at 64x fewer parameters. We release the dataset, model, and training code.

Fine-tuning RoBERTa for CVE-to-CWE Classification: A 125M Parameter Model Competitive with LLMs

Abstract

We present a fine-tuned RoBERTa-base classifier (125M parameters) for mapping Common Vulnerabilities and Exposures (CVE) descriptions to Common Weakness Enumeration (CWE) categories. We construct a large-scale training dataset of 234,770 CVE descriptions with AI-refined CWE labels using Claude Sonnet 4.6, and agreement-filtered evaluation sets where NVD and AI labels agree. On our held-out test set (27,780 samples, 205 CWE classes), the model achieves 87.4% top-1 accuracy and 60.7% Macro F1 -- a +15.5 percentage-point Macro F1 gain over a TF-IDF baseline that already reaches 84.9% top-1, demonstrating the model's advantage on rare weakness categories. On the external CTI-Bench benchmark (NeurIPS 2024), the model achieves 75.6% strict accuracy (95% CI: 72.8-78.2%) -- statistically indistinguishable from Cisco Foundation-Sec-8B-Reasoning (75.3%, 8B parameters) at 64x fewer parameters. We release the dataset, model, and training code.
Paper Structure (24 sections, 2 figures, 6 tables)

This paper contains 24 sections, 2 figures, 6 tables.

Table of Contents

  1. Introduction
  2. Dataset Construction
  3. Source Data
  4. Label Refinement with Claude Sonnet 4.6
  5. Split Construction
  6. MITRE ATT&CK Mapping
  7. Data Decontamination
  8. Method
  9. Model
  10. Two-Phase Training
  11. Training Details
  12. Baseline
  13. Results
  14. Internal Evaluation
  15. CTI-Bench External Benchmark
  16. ...and 9 more sections

Figures (2)

  • Figure 1: CTI-Bench RCM comparison. Our task-specific 125M-parameter model (blue) is competitive with general-purpose 8B-parameter models (gray). Hatched bars indicate Google's closed systems with approximate scores estimated from published charts.
  • Figure 2: Per-CWE accuracy on CTI-Bench cti-rcm (top 10 classes by support). CWE-787 is an outlier due to the hierarchy issue---the model predicts specific children (CWE-121, CWE-122) that map to CWE-787.